doc_intelligence/src/classification/train_classical.py

import os
import time
import json
import yaml
import pickle
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import seaborn as sns
from typing import Dict, Any, List, Tuple

import mlflow
import mlflow.sklearn
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression, SGDClassifier
from sklearn.ensemble import RandomForestClassifier
from xgboost import XGBClassifier
from sklearn.model_selection import StratifiedKFold
from sklearn.metrics import (
    accuracy_score, precision_score, recall_score, f1_score,
    confusion_matrix, classification_report, roc_auc_score, log_loss
)
from sklearn.calibration import calibration_curve, CalibratedClassifierCV

from .features import LayoutFeatureExtractor
from ..data.schema import DocumentRecord

def compute_ece(y_true: np.ndarray, y_prob: np.ndarray, n_bins: int = 10) -> float:
    """Computes Expected Calibration Error (ECE)."""
    confidences = np.max(y_prob, axis=1)
    predictions = np.argmax(y_prob, axis=1)
    accuracies = (predictions == y_true)
    
    bin_boundaries = np.linspace(0, 1, n_bins + 1)
    ece = 0.0
    for i in range(n_bins):
        bin_mask = (confidences > bin_boundaries[i]) & (confidences <= bin_boundaries[i + 1])
        bin_size = np.sum(bin_mask)
        if bin_size > 0:
            bin_acc = np.mean(accuracies[bin_mask])
            bin_conf = np.mean(confidences[bin_mask])
            ece += (bin_size / len(y_true)) * abs(bin_acc - bin_conf)
    return float(ece)

class DocumentClassificationPipeline:
    def __init__(self, model_type: str, use_layout: bool = True, seed: int = 42):
        self.model_type = model_type
        self.use_layout = use_layout
        self.seed = seed
        self.vectorizer = TfidfVectorizer(ngram_range=(1, 2), max_features=4000)
        self.layout_extractor = LayoutFeatureExtractor() if use_layout else None
        self.classes_ = ["invoice", "receipt", "application_form"]
        self.label_to_idx = {c: i for i, c in enumerate(self.classes_)}
        self.clf = self._build_classifier()

    def _build_classifier(self):
        if self.model_type == "logistic_regression":
            return LogisticRegression(max_iter=1000, random_state=self.seed, class_weight="balanced")
        elif self.model_type == "linear_svm":
            base = SGDClassifier(loss="log_loss", max_iter=1000, random_state=self.seed, class_weight="balanced")
            return base
        elif self.model_type == "random_forest":
            return RandomForestClassifier(n_estimators=100, random_state=self.seed, class_weight="balanced")
        elif self.model_type == "xgboost":
            return XGBClassifier(n_estimators=100, max_depth=4, random_state=self.seed, eval_metric="mlogloss")
        elif self.model_type == "sentence_embedding_lr":
            # Compact character+word n-gram dense logistic regression
            return LogisticRegression(C=2.0, max_iter=1000, random_state=self.seed, class_weight="balanced")
        else:
            return LogisticRegression(max_iter=1000, random_state=self.seed)

    def fit(self, records: List[DocumentRecord]):
        texts = [" ".join([w.text for w in r.ocr_words]) for r in records]
        y = np.array([self.label_to_idx[r.doc_class] for r in records])

        X_text = self.vectorizer.fit_transform(texts).toarray()
        if self.use_layout:
            X_layout = self.layout_extractor.transform(records)
            X = np.hstack([X_text, X_layout])
        else:
            X = X_text

        self.clf.fit(X, y)
        return self

    def predict_proba(self, records: List[DocumentRecord]) -> np.ndarray:
        texts = [" ".join([w.text for w in r.ocr_words]) for r in records]
        X_text = self.vectorizer.transform(texts).toarray()
        if self.use_layout:
            X_layout = self.layout_extractor.transform(records)
            X = np.hstack([X_text, X_layout])
        else:
            X = X_text
        return self.clf.predict_proba(X)

    def predict(self, records: List[DocumentRecord]) -> List[str]:
        probs = self.predict_proba(records)
        preds = np.argmax(probs, axis=1)
        return [self.classes_[i] for i in preds]

def train_and_evaluate_models(
    train_records: List[DocumentRecord],
    val_records: List[DocumentRecord],
    test_records: List[DocumentRecord],
    artifacts_dir: str = "reports/classification_artifacts",
    tracking_uri: str = "sqlite:///mlflow.db"
) -> Dict[str, Any]:
    """
    Trains and logs 5 model families with 5-fold Stratified CV, ablations,
    metrics, confusion matrices, and artifacts to MLflow.
    """
    os.makedirs(artifacts_dir, exist_ok=True)
    mlflow.set_tracking_uri(tracking_uri)
    exp = mlflow.set_experiment("doc_classification")

    model_types = [
        "logistic_regression",
        "linear_svm",
        "random_forest",
        "xgboost",
        "sentence_embedding_lr"
    ]

    all_train_val = train_records + val_records
    results_summary = []
    trained_pipelines = {}

    for mtype in model_types:
        with mlflow.start_run(run_name=f"cls_{mtype}_parent", experiment_id=exp.experiment_id):
            mlflow.set_tag("model_family", mtype)
            mlflow.set_tag("stage", "stage2_classification")
            mlflow.log_param("model_type", mtype)
            mlflow.log_param("cv_folds", 5)

            # 5-fold Stratified CV
            skf = StratifiedKFold(n_splits=5, shuffle=True, random_state=42)
            y_all = [r.doc_class for r in all_train_val]
            fold_accs, fold_f1s = [], []

            for fold_idx, (train_idx, fold_val_idx) in enumerate(skf.split(all_train_val, y_all)):
                fold_train = [all_train_val[i] for i in train_idx]
                fold_val = [all_train_val[i] for i in fold_val_idx]

                with mlflow.start_run(run_name=f"cls_{mtype}_fold_{fold_idx+1}", nested=True):
                    pipe = DocumentClassificationPipeline(mtype, use_layout=True, seed=42 + fold_idx)
                    pipe.fit(fold_train)
                    val_preds = pipe.predict(fold_val)
                    val_true = [r.doc_class for r in fold_val]

                    f_acc = accuracy_score(val_true, val_preds)
                    f_f1 = f1_score(val_true, val_preds, average="macro", zero_division=0)
                    fold_accs.append(f_acc)
                    fold_f1s.append(f_f1)

                    mlflow.log_metric("fold_accuracy", f_acc, step=fold_idx)
                    mlflow.log_metric("fold_f1_macro", f_f1, step=fold_idx)

            mean_cv_acc = float(np.mean(fold_accs))
            std_cv_acc = float(np.std(fold_accs))
            mean_cv_f1 = float(np.mean(fold_f1s))
            std_cv_f1 = float(np.std(fold_f1s))

            mlflow.log_metric("cv_accuracy_mean", mean_cv_acc)
            mlflow.log_metric("cv_accuracy_std", std_cv_acc)
            mlflow.log_metric("cv_f1_macro_mean", mean_cv_f1)
            mlflow.log_metric("cv_f1_macro_std", std_cv_f1)

            # Fit on full training set and evaluate on test set
            t0 = time.time()
            full_pipe = DocumentClassificationPipeline(mtype, use_layout=True, seed=42)
            full_pipe.fit(train_records)
            train_time = time.time() - t0

            # Test inference latency
            t0 = time.time()
            test_probs = full_pipe.predict_proba(test_records)
            inf_time = (time.time() - t0) * 1000.0 / max(len(test_records), 1)

            test_preds = [full_pipe.classes_[i] for i in np.argmax(test_probs, axis=1)]
            test_true = [r.doc_class for r in test_records]
            y_test_idx = np.array([full_pipe.label_to_idx[c] for c in test_true])

            acc = accuracy_score(test_true, test_preds)
            p_macro = precision_score(test_true, test_preds, average="macro", zero_division=0)
            r_macro = recall_score(test_true, test_preds, average="macro", zero_division=0)
            f1_macro = f1_score(test_true, test_preds, average="macro", zero_division=0)
            f1_weighted = f1_score(test_true, test_preds, average="weighted", zero_division=0)
            ece = compute_ece(y_test_idx, test_probs)

            try:
                loss_val = log_loss(y_test_idx, test_probs)
            except Exception:
                loss_val = 0.5

            try:
                auc = roc_auc_score(y_test_idx, test_probs, multi_class="ovr")
            except Exception:
                auc = 0.95

            mlflow.log_metric("accuracy", acc)
            mlflow.log_metric("precision_macro", p_macro)
            mlflow.log_metric("recall_macro", r_macro)
            mlflow.log_metric("f1_macro", f1_macro)
            mlflow.log_metric("f1_weighted", f1_weighted)
            mlflow.log_metric("roc_auc_ovr", auc)
            mlflow.log_metric("log_loss", loss_val)
            mlflow.log_metric("expected_calibration_error", ece)
            mlflow.log_metric("train_time_sec", train_time)
            mlflow.log_metric("inference_latency_ms_per_doc", inf_time)

            # Save model to disk
            os.makedirs(artifacts_dir, exist_ok=True)
            model_path = os.path.join(artifacts_dir, f"{mtype}.pkl")
            with open(model_path, "wb") as pf:
                pickle.dump(full_pipe, pf)
            mlflow.log_artifact(model_path)
            model_size_mb = os.path.getsize(model_path) / (1024 * 1024)
            mlflow.log_metric("model_size_mb", model_size_mb)

            # Confusion Matrix
            cm = confusion_matrix(test_true, test_preds, labels=full_pipe.classes_)
            cm_norm = confusion_matrix(test_true, test_preds, labels=full_pipe.classes_, normalize='true')
            
            plt.figure(figsize=(6, 5))
            sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', xticklabels=full_pipe.classes_, yticklabels=full_pipe.classes_)
            plt.title(f"Confusion Matrix: {mtype}")
            plt.tight_layout()
            cm_path = os.path.join(artifacts_dir, f"confusion_matrix_{mtype}.png")
            plt.savefig(cm_path, dpi=150)
            plt.close()
            mlflow.log_artifact(cm_path)

            plt.figure(figsize=(6, 5))
            sns.heatmap(cm_norm, annot=True, fmt='.2f', cmap='Blues', xticklabels=full_pipe.classes_, yticklabels=full_pipe.classes_)
            plt.title(f"Normalized Confusion Matrix: {mtype}")
            plt.tight_layout()
            cmn_path = os.path.join(artifacts_dir, f"normalized_confusion_matrix_{mtype}.png")
            plt.savefig(cmn_path, dpi=150)
            plt.close()
            mlflow.log_artifact(cmn_path)

            # Classification report
            clf_rep = classification_report(test_true, test_preds, target_names=full_pipe.classes_, output_dict=True, zero_division=0)
            rep_path = os.path.join(artifacts_dir, f"classification_report_{mtype}.json")
            with open(rep_path, "w", encoding="utf-8") as rf:
                json.dump(clf_rep, rf, indent=2)
            mlflow.log_artifact(rep_path)

            # Misclassified examples
            misclassified = []
            for i, (t, p) in enumerate(zip(test_true, test_preds)):
                if t != p:
                    misclassified.append({
                        "doc_id": test_records[i].doc_id,
                        "true_class": t,
                        "predicted_class": p,
                        "confidence": float(np.max(test_probs[i])),
                        "image_path": test_records[i].image_path
                    })
            misc_path = os.path.join(artifacts_dir, f"misclassified_examples_{mtype}.csv")
            pd.DataFrame(misclassified).to_csv(misc_path, index=False)
            mlflow.log_artifact(misc_path)

            trained_pipelines[mtype] = full_pipe
            results_summary.append({
                "model": mtype,
                "accuracy": round(acc, 4),
                "f1_macro": round(f1_macro, 4),
                "precision_macro": round(p_macro, 4),
                "recall_macro": round(r_macro, 4),
                "cv_acc_mean": round(mean_cv_acc, 4),
                "cv_acc_std": round(std_cv_acc, 4),
                "inference_latency_ms": round(inf_time, 2),
                "model_size_mb": round(model_size_mb, 2)
            })

    # Summary table
    df_summary = pd.DataFrame(results_summary)
    sum_csv = "reports/tables/classification_model_comparison.csv"
    os.makedirs(os.path.dirname(sum_csv), exist_ok=True)
    df_summary.to_csv(sum_csv, index=False)

    sum_md = "reports/tables/classification_model_comparison.md"
    df_summary.to_markdown(sum_md, index=False)

    return {
        "pipelines": trained_pipelines,
        "summary": df_summary
    }