ml/model/evaluate.py

"""
ml/model/evaluate.py - Honest evaluation + reproducibility evidence
===================================================================
Produces the exact audit-ready numbers: accuracy, per-class
precision/recall/F1, macro-F1, ROC-AUC, and a confusion matrix on the
OUT-OF-SPEAKER held-out test split (voices never seen in training).

Usage:
    python -m ml.model.evaluate --ckpt ml/models/stutter/stutter_lora \
        --data data/synthetic_lattice/dataset --out reports/ev
"""
from __future__ import annotations
import argparse
import json
from pathlib import Path

import numpy as np
import torch
from sklearn.metrics import (
    accuracy_score, precision_recall_fscore_support, confusion_matrix, roc_auc_score,
)
import json as _json
from peft import PeftModel
from transformers import Wav2Vec2FeatureExtractor, Wav2Vec2ForSequenceClassification

from ml.model.stutter_trainer import (
    SR, MAX_SECONDS, ID2LABEL, BIN_ID2LABEL, prepare_dataset, MODEL_BASE,
    clean_cache,
)


def _batch_input(row, device):
    """One tokenized row -> keyword tensors for model forward."""
    x = np.asarray(row["input_values"])
    out = {"input_values": torch.tensor(x, dtype=torch.float32).unsqueeze(0).to(device)}
    if "attention_mask" in row:
        m = np.asarray(row["attention_mask"])
        out["attention_mask"] = torch.tensor(m, dtype=torch.long).unsqueeze(0).to(device)
    return out


def evaluate(data_dir, ckpt_dir, out="reports/ev", device=None, threshold: float = 0.5):
    device = device or ("cuda" if torch.cuda.is_available() else "cpu")

    cm_path = Path(ckpt_dir).parent / "class_map.json"
    binary = True
    if cm_path.exists():
        try:
            cm = _json.loads(cm_path.read_text(encoding="utf-8"))
            binary = bool(cm.get("binary", True))
        except Exception:
            binary = True
    id2l = BIN_ID2LABEL if binary else ID2LABEL
    n_classes = len(id2l)

    feat = Wav2Vec2FeatureExtractor(sampling_rate=SR)
    tr, va, te = prepare_dataset(data_dir, feat, binary=binary)

    base = Wav2Vec2ForSequenceClassification.from_pretrained(
        MODEL_BASE, num_labels=n_classes, ignore_mismatched_sizes=True)
    model = PeftModel.from_pretrained(base, str(ckpt_dir))
    model.to(device)
    model.eval()

    y_true, y_pred, y_probs = [], [], []
    for row in te:
        x = _batch_input(row, device)
        with torch.no_grad():
            logits = model(**x).logits
            probs = torch.softmax(logits, dim=1)[0].cpu().numpy()
        
        y_true.append(int(row["labels"]))
        y_probs.append(probs)
        if binary:
            pred = 1 if probs[1] >= threshold else 0
        else:
            pred = int(np.argmax(probs))
        y_pred.append(pred)

    y_true = np.array(y_true)
    y_pred = np.array(y_pred)
    y_probs = np.array(y_probs)

    cids = list(range(n_classes))
    acc = accuracy_score(y_true, y_pred)
    p, r, f, _ = precision_recall_fscore_support(
        y_true, y_pred, labels=cids, zero_division=0)
    macro_f1 = float(np.mean(f))
    cm = confusion_matrix(y_true, y_pred, labels=cids).tolist()

    auc_score = None
    if binary and len(np.unique(y_true)) > 1:
        try:
            auc_score = float(roc_auc_score(y_true, y_probs[:, 1]))
        except Exception:
            auc_score = None

    report = {
        "model": str(ckpt_dir),
        "base_model": "facebook/wav2vec2-base",
        "adapter": "LoRA (r=8, alpha=16, target q/k/v)",
        "n_train": len(tr),
        "n_val": len(va),
        "n_test": len(te),
        "split": "by-speaker (test voices never seen in training)",
        "binary": binary,
        "threshold": threshold,
        "accuracy": round(float(acc), 4),
        "macro_f1": round(float(macro_f1), 4),
        "roc_auc": round(float(auc_score), 4) if auc_score is not None else None,
        "per_class": {
            id2l[i]: {
                "precision": round(float(p[i]), 4),
                "recall": round(float(r[i]), 4),
                "f1": round(float(f[i]), 4),
            }
            for i in cids
        },
        "confusion_matrix": cm,
        "class_map": id2l,
        "metric_definitions": {
            "accuracy": "correct / total on out-of-speaker test set",
            "precision": "class TP / (TP+FP)",
            "recall": "class TP / (TP+FN)",
            "macro_f1": "mean of per-class F1",
            "roc_auc": "area under ROC curve",
        },
    }

    out = Path(out)
    out.mkdir(parents=True, exist_ok=True)
    report_file = out / ("evaluation.json" if "synthetic" not in str(data_dir) else "synthetic_eval.json")
    text_file = out / ("evaluation.txt" if "synthetic" not in str(data_dir) else "synthetic_eval.txt")
    
    report_file.write_text(json.dumps(report, indent=2), encoding="utf-8")
    text_file.write_text(render(report), encoding="utf-8")
    clean_cache()
    print(f"[eval] -> {report_file} and {text_file}")
    print(f"  Accuracy:  {report['accuracy']:.4f}")
    print(f"  Macro-F1:  {report['macro_f1']:.4f}")
    if auc_score is not None:
        print(f"  ROC-AUC:   {report['roc_auc']:.4f}")
    for k, v in report["per_class"].items():
        print(f"  {k:16} Prec: {v['precision']:.4f} | Rec: {v['recall']:.4f} | F1: {v['f1']:.4f}")
    return report


def render(r):
    L = [f"EVALUATION  base={r['base_model']}  adapter={r['adapter']}",
         f"Test set: {r['n_test']} clips, split by speaker (unseen voices)",
         f"Accuracy {r['accuracy']:.4f}   Macro-F1 {r['macro_f1']:.4f}" + (f"   ROC-AUC {r['roc_auc']:.4f}" if r.get('roc_auc') else ""),
         "Per-class (precision / recall / F1):"]
    for lab, m in r["per_class"].items():
        L.append(f"  {lab:18} {m['precision']:.3f}  {m['recall']:.3f}  {m['f1']:.3f}")
    L.append("Confusion matrix (rows=true, cols=pred):")
    hdr = "              " + "  ".join(f"{c:>8}" for c in r["class_map"].values())
    L.append(hdr)
    for i, row in enumerate(r["confusion_matrix"]):
        L.append(f"{r['class_map'][i]:>12} " + "  ".join(f"{v:>8}" for v in row))
    return "\n".join(L)


def _main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", default="ml/models/stutter/stutter_lora")
    ap.add_argument("--data", default="data/synthetic_lattice/dataset")
    ap.add_argument("--out", default="reports/ev")
    ap.add_argument("--threshold", type=float, default=0.5)
    a = ap.parse_args()
    evaluate(a.data, a.ckpt, a.out, threshold=a.threshold)


if __name__ == "__main__":
    _main()