""" 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()