Spaces:
Runtime error
Runtime error
| """ | |
| Ticket 6 — Evaluation & Confusion Matrix Report (ATA classifier). | |
| Runs the trained model against the TEST split (never used for training or | |
| checkpoint selection -- that was val) and produces: | |
| - overall accuracy / macro-F1 (comparable to the PRD's success metric) | |
| - per-class precision/recall/F1 (exposes which classes are actually weak, | |
| which a single headline number hides) | |
| - a confusion matrix image | |
| - a metrics JSON, for the Model Performance Dashboard (Ticket 16) to read later | |
| """ | |
| import json | |
| from pathlib import Path | |
| import matplotlib | |
| matplotlib.use("Agg") | |
| import matplotlib.pyplot as plt | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| from sklearn.metrics import ( | |
| accuracy_score, | |
| classification_report, | |
| confusion_matrix, | |
| precision_recall_fscore_support, | |
| ) | |
| from transformers import AutoModelForSequenceClassification, AutoTokenizer | |
| MODEL_DIR = "models/ata_classifier" | |
| TEST_CSV = "data/processed/test.csv" | |
| MAX_LENGTH = 256 | |
| OUT_DIR = "reports/ata_classifier" | |
| def main(): | |
| Path(OUT_DIR).mkdir(parents=True, exist_ok=True) | |
| tokenizer = AutoTokenizer.from_pretrained(MODEL_DIR, local_files_only=True) | |
| model = AutoModelForSequenceClassification.from_pretrained(MODEL_DIR, local_files_only=True) | |
| model.eval() | |
| with open(f"{MODEL_DIR}/label_map.json") as f: | |
| label_map = json.load(f) | |
| id2label = {int(k): v for k, v in label_map["id2label"].items()} | |
| label2id = label_map["label2id"] | |
| df = pd.read_csv(TEST_CSV) | |
| df = df[df["ata_chapter_label"].notna()].reset_index(drop=True) | |
| print(f"Test examples (ATA-labeled): {len(df)}") | |
| true_labels = df["ata_chapter_label"].astype(str).tolist() | |
| y_true = [label2id[t] for t in true_labels] | |
| y_pred = [] | |
| with torch.no_grad(): | |
| for text in df["narrative_text"].tolist(): | |
| inputs = tokenizer(text, truncation=True, max_length=MAX_LENGTH, return_tensors="pt") | |
| logits = model(**inputs).logits | |
| pred_id = int(torch.argmax(logits, dim=1).item()) | |
| y_pred.append(pred_id) | |
| acc = accuracy_score(y_true, y_pred) | |
| precision, recall, f1, support = precision_recall_fscore_support( | |
| y_true, y_pred, average="macro", zero_division=0 | |
| ) | |
| print(f"\n=== TEST SET RESULTS (never used in training or checkpoint selection) ===") | |
| print(f"Accuracy: {acc:.1%} (PRD target: >=80%)") | |
| print(f"Macro F1: {f1:.3f}") | |
| print(f"Macro Precision: {precision:.3f} Macro Recall: {recall:.3f}") | |
| labels_present = sorted(set(y_true) | set(y_pred)) | |
| target_names = [id2label[i] for i in labels_present] | |
| report_dict = classification_report( | |
| y_true, y_pred, labels=labels_present, target_names=target_names, | |
| zero_division=0, output_dict=True | |
| ) | |
| print("\nPer-class report:") | |
| print(classification_report(y_true, y_pred, labels=labels_present, target_names=target_names, zero_division=0)) | |
| cm = confusion_matrix(y_true, y_pred, labels=labels_present) | |
| fig, ax = plt.subplots(figsize=(10, 9)) | |
| im = ax.imshow(cm, cmap="Blues") | |
| ax.set_xticks(range(len(target_names))) | |
| ax.set_yticks(range(len(target_names))) | |
| ax.set_xticklabels(target_names, rotation=90) | |
| ax.set_yticklabels(target_names) | |
| ax.set_xlabel("Predicted") | |
| ax.set_ylabel("True") | |
| ax.set_title("ATA Chapter Classifier — Confusion Matrix (Test Set)") | |
| for i in range(len(target_names)): | |
| for j in range(len(target_names)): | |
| if cm[i, j] > 0: | |
| ax.text(j, i, str(cm[i, j]), ha="center", va="center", | |
| color="white" if cm[i, j] > cm.max() / 2 else "black", fontsize=7) | |
| fig.colorbar(im) | |
| fig.tight_layout() | |
| fig.savefig(f"{OUT_DIR}/confusion_matrix.png", dpi=150) | |
| print(f"\nSaved confusion matrix to {OUT_DIR}/confusion_matrix.png") | |
| metrics_out = { | |
| "test_accuracy": acc, | |
| "test_macro_f1": f1, | |
| "test_macro_precision": precision, | |
| "test_macro_recall": recall, | |
| "n_test_examples": len(df), | |
| "per_class": report_dict, | |
| "prd_target_accuracy": 0.80, | |
| "meets_prd_target": acc >= 0.80, | |
| } | |
| with open(f"{OUT_DIR}/metrics.json", "w") as f: | |
| json.dump(metrics_out, f, indent=2) | |
| print(f"Saved metrics JSON to {OUT_DIR}/metrics.json") | |
| if __name__ == "__main__": | |
| main() | |