Spaces:
Runtime error
Runtime error
File size: 4,327 Bytes
d4ec2fa | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | """
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()
|