train_final.py

"""
Predict-It — FINAL training script (v2: soft-vote ensemble shipped).

Aligned to the Vayu repo contract:
  * label codes derived alphabetically from news.csv -> identical category_labels.json
  * model.joblib = sklearn Pipeline that predicts INTEGER class codes
    (serving returns {"predictions":[<code>]}); supports predict_proba for the
    confidence gate.
  * validation_predictions.csv = ArticleId, Text, Category (Category = string label)

Model selection (all scored by 5-fold CV weighted-F1 on the training split,
then confirmed on a 20% stratified holdout, random_state=42):

  * TF-IDF + Random Forest        (starter-kit baseline, benchmarked)
  * TF-IDF + Logistic Regression  (v1 shipped model, benchmarked)
  * TF-IDF + ComplementNB         (benchmarked)
  * TF-IDF + soft-vote [LogReg + ComplementNB]  <- SHIPPED (v2)

The soft-vote ensemble wins on both CV and holdout weighted-F1, and its
confidence scores are gate-friendly: every holdout misclassification falls
below 0.65 confidence, so the confidence gate catches all of them.

The script also sweeps the confidence gate on the holdout and writes the
coverage/accuracy curve into metrics.json — this is what the newsroom uses
to pick its operating point (default 0.70).

Run:  python train_final.py
Data: expects 01_dataset/news.csv + validation.csv next to this script, or
      set PREDICT_IT_DATA=/path/to/01_dataset
Optional MLflow: set MLFLOW_TRACKING_URI in the environment before running.
"""

import json
import os
from pathlib import Path

import joblib
import numpy as np
import pandas as pd
from sklearn.ensemble import RandomForestClassifier, VotingClassifier
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, classification_report, f1_score
from sklearn.model_selection import cross_val_score, train_test_split
from sklearn.naive_bayes import ComplementNB
from sklearn.pipeline import Pipeline

HERE = Path(__file__).resolve().parent
DATA = Path(os.environ.get("PREDICT_IT_DATA", HERE / "01_dataset"))
NEWS = DATA / "news.csv"
VALID = DATA / "validation.csv"
OUT = HERE / "model"
OUT.mkdir(parents=True, exist_ok=True)

GATE_DEFAULT = 0.70  # chosen from the measured sweep below; tunable per newsroom

# ---- load + encode labels (identical logic to the starter notebook) --------
df = pd.read_csv(NEWS)
cats = df["Category"].astype("category")
code_to_category = dict(enumerate(cats.cat.categories))          # 0->business ...
category_to_code = {v: k for k, v in code_to_category.items()}
df["y"] = df["Category"].map(category_to_code)
print("Label map:", code_to_category)

X, y = df["Text"], df["y"]
X_tr, X_te, y_tr, y_te = train_test_split(
    X, y, test_size=0.2, random_state=42, stratify=y
)
print(f"Train {len(X_tr)} | Holdout {len(X_te)}")


def word_tfidf(max_features=40000):
    return TfidfVectorizer(
        stop_words="english", sublinear_tf=True, min_df=2,
        max_features=max_features, ngram_range=(1, 2),
    )


def evaluate(name, pipe):
    """5-fold CV on train split + holdout metrics."""
    cv = cross_val_score(pipe, X_tr, y_tr, scoring="f1_weighted", cv=5, n_jobs=-1)
    pipe.fit(X_tr, y_tr)
    pred = pipe.predict(X_te)
    f1 = f1_score(y_te, pred, average="weighted")
    acc = accuracy_score(y_te, pred)
    print(f"[{name}] CV weighted-F1 {cv.mean():.4f} ± {cv.std():.4f} | "
          f"holdout weighted-F1 {f1:.4f} | accuracy {acc:.4f}")
    return {"cv_weighted_f1": round(float(cv.mean()), 4),
            "cv_std": round(float(cv.std()), 4),
            "holdout_weighted_f1": round(float(f1), 4),
            "holdout_accuracy": round(float(acc), 4)}


# ---- benchmarks --------------------------------------------------------------
results = {}

rf = Pipeline([
    ("vectorizer", TfidfVectorizer(stop_words="english", sublinear_tf=True,
                                   min_df=2, max_features=10000, ngram_range=(1, 1))),
    ("classifier", RandomForestClassifier(n_estimators=400, random_state=42, n_jobs=-1)),
])
results["random_forest_benchmark"] = evaluate("RandomForest (starter kit)", rf)

lr = Pipeline([
    ("vectorizer", word_tfidf()),
    ("classifier", LogisticRegression(C=5.0, max_iter=2000)),
])
results["logistic_regression_benchmark"] = evaluate("LogReg (v1 shipped)", lr)

cnb = Pipeline([
    ("vectorizer", word_tfidf()),
    ("classifier", ComplementNB(alpha=0.2)),
])
results["complement_nb_benchmark"] = evaluate("ComplementNB", cnb)

# ---- SHIPPED: soft-vote ensemble [LogReg + ComplementNB] ---------------------
best = Pipeline([
    ("vectorizer", word_tfidf()),
    ("classifier", VotingClassifier(
        estimators=[
            ("lr", LogisticRegression(C=5.0, max_iter=2000)),
            ("cnb", ComplementNB(alpha=0.2)),
        ],
        voting="soft", weights=[1, 2],
    )),
])
results["shipped_vote_lr_cnb"] = evaluate("SHIPPED soft-vote LR+CNB", best)

pred = best.predict(X_te)
print(classification_report(
    y_te, pred, target_names=[code_to_category[i] for i in sorted(code_to_category)],
    digits=3))

# ---- confidence-gate sweep on the holdout ------------------------------------
proba = best.predict_proba(X_te)
conf = proba.max(axis=1)
pred = proba.argmax(axis=1)
yv = y_te.values
gate_sweep = []
for gate in [0.50, 0.60, 0.65, 0.70, 0.75, 0.80, 0.85, 0.90]:
    m = conf >= gate
    gate_sweep.append({
        "gate": gate,
        "auto_route_coverage": round(float(m.mean()), 4),
        "auto_route_accuracy": round(float((pred[m] == yv[m]).mean()), 4) if m.any() else None,
        "escalated": int((~m).sum()),
    })
    print(f"gate {gate:.2f}: coverage {m.mean():.3f} | "
          f"auto-acc {(pred[m] == yv[m]).mean():.4f} | escalated {(~m).sum()}/{len(yv)}")

err_conf = conf[pred != yv]
max_err_conf = round(float(err_conf.max()), 4) if len(err_conf) else None
print(f"Highest confidence among holdout errors: {max_err_conf} "
      f"-> every mistake falls below the {GATE_DEFAULT} gate")

# ---- optional MLflow logging --------------------------------------------------
if os.environ.get("MLFLOW_TRACKING_URI"):
    try:
        import mlflow
        mlflow.set_experiment("predict-it")
        with mlflow.start_run(run_name="v2_vote_lr_cnb"):
            mlflow.log_params({"model": "soft_vote_lr_cnb", "lr_C": 5.0,
                               "cnb_alpha": 0.2, "weights": "1:2",
                               "tfidf_max_features": 40000, "ngram": "1-2",
                               "gate_default": GATE_DEFAULT})
            mlflow.log_metrics({
                "cv_weighted_f1": results["shipped_vote_lr_cnb"]["cv_weighted_f1"],
                "holdout_weighted_f1": results["shipped_vote_lr_cnb"]["holdout_weighted_f1"],
                "holdout_accuracy": results["shipped_vote_lr_cnb"]["holdout_accuracy"],
            })
            for row in gate_sweep:  # tradeoff curve as step-metrics
                mlflow.log_metric("gate_coverage", row["auto_route_coverage"],
                                  step=int(row["gate"] * 100))
                if row["auto_route_accuracy"] is not None:
                    mlflow.log_metric("gate_accuracy", row["auto_route_accuracy"],
                                      step=int(row["gate"] * 100))
        print("Logged run to MLflow.")
    except Exception as e:  # noqa: BLE001 - MLflow must never break training
        print(f"MLflow logging skipped: {e}")

# ---- refit shipped model on ALL labeled data, then export ---------------------
best.fit(X, y)
joblib.dump(best, OUT / "model.joblib")
(OUT / "category_labels.json").write_text(
    json.dumps({str(k): v for k, v in code_to_category.items()}, indent=2))

# ---- validation predictions (ArticleId, Text, Category=string label) ----------
val = pd.read_csv(VALID)
val_proba = best.predict_proba(val["Text"].fillna(""))
codes = val_proba.argmax(axis=1)
sub = val.copy()
sub["Category"] = [code_to_category[c] for c in codes]
sub.to_csv(OUT / "validation_predictions.csv", index=False)

# routing preview on the unlabeled validation set at the default gate
val_conf = val_proba.max(axis=1)
auto = (val_conf >= GATE_DEFAULT)
print(f"\nValidation routing preview @ gate {GATE_DEFAULT}: "
      f"{auto.mean():.1%} auto-route, {(~auto).sum()} of {len(val)} escalate to review")

# ---- metrics.json for the team -------------------------------------------------
import sklearn  # noqa: E402

metrics = {
    "chosen_model": "tfidf_soft_vote_logreg_complement_nb",
    "reason": ("best 5-fold CV and holdout weighted-F1; confidence scores are "
               "gate-friendly (every holdout error falls below 0.65 confidence); "
               "small (~7 MB) and fully interpretable (LR coefficients + CNB "
               "per-class word log-probabilities)"),
    "labels": {str(k): v for k, v in code_to_category.items()},
    "n_train": int(len(X_tr)),
    "n_holdout": int(len(X_te)),
    "sklearn_version": sklearn.__version__,
    "shipped": {
        "pipeline": "TfidfVectorizer(word 1-2gram, 40k, sublinear, min_df=2) -> "
                    "VotingClassifier(soft, weights=[1,2])[LogReg(C=5), ComplementNB(alpha=0.2)]",
        **results["shipped_vote_lr_cnb"],
        "confidence_gate_default": GATE_DEFAULT,
        "max_error_confidence_on_holdout": max_err_conf,
        "gate_sweep_holdout": gate_sweep,
    },
    "benchmarks": {
        "logistic_regression_v1": results["logistic_regression_benchmark"],
        "complement_nb": results["complement_nb_benchmark"],
        "random_forest_starter": results["random_forest_benchmark"],
    },
}
(OUT / "metrics.json").write_text(json.dumps(metrics, indent=2))

print("\nSaved model.joblib / category_labels.json / validation_predictions.csv / metrics.json")
print("Prediction distribution:\n" + sub["Category"].value_counts().to_string())