src/train.py

"""
Phase 3 — Stage-2 classifier (fine-tuned DeBERTa-v3-small).

Trains a binary Consistent/Mismatched classifier on the pseudo-labeled data.
- Inputs: model_text  (assigned priority + channel/category/tier tags + real issue
          sentence) -> satisfies "text + >=1 structured metadata feature".
- Imbalance: class-weighted cross-entropy (inverse frequency).
- Held-out evaluation on the FIXED `split=='test'` rows; reports the three
  verification metrics and PASS/FAIL vs the thresholds in config.

Run:  python3 src/train.py
Out:  artifacts/models/deberta-sia/   (model + tokenizer)
      artifacts/metrics/classifier_metrics.json
      artifacts/data/test_predictions.parquet
"""
from __future__ import annotations
import os, sys, json
os.environ.setdefault("PYTORCH_ENABLE_MPS_FALLBACK", "1")
os.environ.setdefault("TOKENIZERS_PARALLELISM", "false")
import numpy as np
import pandas as pd
import torch
import torch.nn.functional as F

sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
from src import config as C

from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, f1_score, recall_score, confusion_matrix
from transformers import (AutoTokenizer, AutoModelForSequenceClassification,
                          TrainingArguments, Trainer, DataCollatorWithPadding)

torch.manual_seed(C.SEED); np.random.seed(C.SEED)

# --------------------------------------------------------------------------- #
class TicketDS(torch.utils.data.Dataset):
    def __init__(self, enc, labels):
        self.enc, self.labels = enc, labels
    def __len__(self): return len(self.labels)
    def __getitem__(self, i):
        item = {k: v[i] for k, v in self.enc.items()}
        item["labels"] = int(self.labels[i])
        return item

class WeightedTrainer(Trainer):
    def __init__(self, *a, class_weights=None, **kw):
        super().__init__(*a, **kw)
        self.class_weights = class_weights
    def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
        labels = inputs.pop("labels")
        outputs = model(**inputs)
        loss = F.cross_entropy(outputs.logits, labels,
                               weight=self.class_weights.to(outputs.logits.device))
        return (loss, outputs) if return_outputs else loss

def compute_metrics(eval_pred):
    logits = eval_pred.predictions
    logits = logits[0] if isinstance(logits, tuple) else logits
    labels = eval_pred.label_ids
    preds = logits.argmax(-1)
    rec = recall_score(labels, preds, average=None, labels=[0, 1], zero_division=0)
    return {
        "accuracy": float(accuracy_score(labels, preds)),
        "macro_f1": float(f1_score(labels, preds, average="macro", zero_division=0)),
        "recall_consistent": float(rec[0]),
        "recall_mismatch": float(rec[1]),
    }

def load_tokenizer():
    try:
        return AutoTokenizer.from_pretrained(C.BASE_MODEL)
    except Exception as e:
        print("[tok] fast failed, falling back to slow:", e)
        return AutoTokenizer.from_pretrained(C.BASE_MODEL, use_fast=False)

# --------------------------------------------------------------------------- #
def main():
    df = pd.read_parquet(C.PROC_DIR / "pseudo_labeled.parquet")
    train_full = df[df.split == "train"].reset_index(drop=True)
    test_df = df[df.split == "test"].reset_index(drop=True)

    tr_df, val_df = train_test_split(
        train_full, test_size=C.TRAIN["val_size"], random_state=C.SEED,
        stratify=train_full["mismatch"])

    print(f"[data] train={len(tr_df):,} val={len(val_df):,} test={len(test_df):,}")
    print(f"[label balance] train mismatch rate={tr_df['mismatch'].mean():.3f}")

    tok = load_tokenizer()
    def enc(texts):
        return tok(list(texts), truncation=True, max_length=C.MAX_LEN)
    ds_tr = TicketDS(enc(tr_df["model_text"]), tr_df["mismatch"].values)
    ds_val = TicketDS(enc(val_df["model_text"]), val_df["mismatch"].values)
    ds_te = TicketDS(enc(test_df["model_text"]), test_df["mismatch"].values)

    # inverse-frequency class weights
    counts = np.bincount(tr_df["mismatch"].values, minlength=2)
    cw = torch.tensor(counts.sum() / (2.0 * counts), dtype=torch.float)
    print(f"[class weights] consistent={cw[0]:.3f} mismatch={cw[1]:.3f}")

    model = AutoModelForSequenceClassification.from_pretrained(
        C.BASE_MODEL, num_labels=C.NUM_LABELS,
        id2label={0: C.CLASS_NAMES[0], 1: C.CLASS_NAMES[1]},
        label2id={C.CLASS_NAMES[0]: 0, C.CLASS_NAMES[1]: 1})

    args = TrainingArguments(
        output_dir=str(C.CLASSIFIER_DIR),
        num_train_epochs=C.TRAIN["epochs"],
        per_device_train_batch_size=C.TRAIN["batch_size"],
        per_device_eval_batch_size=32,
        gradient_accumulation_steps=C.TRAIN["grad_accum"],
        learning_rate=C.TRAIN["lr"], weight_decay=C.TRAIN["weight_decay"],
        warmup_ratio=C.TRAIN["warmup_ratio"],
        eval_strategy="epoch", save_strategy="epoch",
        load_best_model_at_end=True, metric_for_best_model="macro_f1",
        greater_is_better=True, save_total_limit=1,
        logging_steps=50, report_to=[], seed=C.SEED,
        dataloader_num_workers=0, disable_tqdm=False)

    trainer = WeightedTrainer(
        model=model, args=args, train_dataset=ds_tr, eval_dataset=ds_val,
        processing_class=tok, data_collator=DataCollatorWithPadding(tok),
        compute_metrics=compute_metrics, class_weights=cw)

    trainer.train()

    # ---- final held-out evaluation ----
    pred = trainer.predict(ds_te)
    logits = pred.predictions[0] if isinstance(pred.predictions, tuple) else pred.predictions
    preds = logits.argmax(-1)
    m = compute_metrics(pred)
    cm = confusion_matrix(test_df["mismatch"].values, preds).tolist()

    th = C.THRESHOLDS
    passed = (m["accuracy"] >= th["accuracy"] and m["macro_f1"] >= th["macro_f1"]
              and m["recall_consistent"] >= th["per_class_recall"]
              and m["recall_mismatch"] >= th["per_class_recall"])
    result = {**m, "confusion_matrix": cm, "thresholds": th, "PASSED": bool(passed),
              "n_test": int(len(test_df))}

    trainer.save_model(str(C.CLASSIFIER_DIR))
    tok.save_pretrained(str(C.CLASSIFIER_DIR))
    with open(C.METRICS_DIR / "classifier_metrics.json", "w") as f:
        json.dump(result, f, indent=2)
    test_df = test_df.copy()
    test_df["pred_mismatch"] = preds
    test_df["pred_prob_mismatch"] = F.softmax(torch.tensor(logits), -1)[:, 1].numpy()
    test_df.to_parquet(C.PROC_DIR / "test_predictions.parquet", index=False)

    print("\n==================  HELD-OUT TEST METRICS  ==================")
    print(f"  Accuracy          : {m['accuracy']*100:.2f}%   (>= {th['accuracy']*100:.0f}%)")
    print(f"  Macro F1          : {m['macro_f1']:.4f}    (>= {th['macro_f1']})")
    print(f"  Recall Consistent : {m['recall_consistent']:.4f}    (>= {th['per_class_recall']})")
    print(f"  Recall Mismatch   : {m['recall_mismatch']:.4f}    (>= {th['per_class_recall']})")
    print(f"  Confusion matrix  : {cm}")
    print(f"  >>> VERIFICATION: {'PASSED ✅' if passed else 'FAILED ❌'}")
    print("============================================================")
    return result

if __name__ == "__main__":
    main()