src/evaluate.py

"""
src/evaluate.py
---------------
Run full test-set evaluation on the fine-tuned DistilBERT model.
Saves per-example predictions + probabilities to results/test_predictions.csv
for use in the error analysis notebook.

Usage:
    python -m src.evaluate
"""

import os
import json
import numpy as np
import pandas as pd
import torch
from torch.utils.data import Dataset, DataLoader
from transformers import DistilBertTokenizerFast, DistilBertForSequenceClassification
from sklearn.metrics import f1_score, hamming_loss, precision_score, recall_score

from src.config import (
    DISTORTION_LABELS, NUM_LABELS,
    DATA_PROC_DIR, MODELS_DIR, RESULTS_DIR,
    BATCH_SIZE, SEED
)

MODEL_PATH = os.path.join(MODELS_DIR, "distilbert_cognitive_distortion")
MAX_LEN    = 128
THRESHOLD  = 0.2   # matches Colab training


class CognitiveDistortionDataset(Dataset):
    def __init__(self, df, tokenizer, max_len):
        self.texts     = df["text"].astype(str).tolist()
        self.labels    = df[DISTORTION_LABELS].values.astype(np.float32)
        self.tokenizer = tokenizer
        self.max_len   = max_len

    def __len__(self):
        return len(self.texts)

    def __getitem__(self, idx):
        enc = self.tokenizer(
            self.texts[idx],
            max_length=self.max_len,
            padding="max_length",
            truncation=True,
            return_tensors="pt"
        )
        return {
            "input_ids"     : enc["input_ids"].squeeze(0),
            "attention_mask": enc["attention_mask"].squeeze(0),
            "labels"        : torch.tensor(self.labels[idx], dtype=torch.float32)
        }


def load_baseline():
    path = os.path.join(RESULTS_DIR, "baseline_results.json")
    if not os.path.exists(path):
        return {}
    with open(path) as f:
        return json.load(f)


def run():
    torch.manual_seed(SEED)
    np.random.seed(SEED)

    if not os.path.isdir(MODEL_PATH):
        raise FileNotFoundError(
            f"Model not found at: {MODEL_PATH}\n"
            "Extract the Colab zip and place the folder there."
        )

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Device   : {device}")

    # ── Load data ─────────────────────────────────────────────────────────────
    test_df = pd.read_csv(os.path.join(DATA_PROC_DIR, "test.csv"))
    print(f"Test set : {len(test_df)} rows")

    # ── Load model ────────────────────────────────────────────────────────────
    print(f"Loading  : {MODEL_PATH}")
    tokenizer = DistilBertTokenizerFast.from_pretrained(MODEL_PATH)
    model     = DistilBertForSequenceClassification.from_pretrained(
        MODEL_PATH, num_labels=NUM_LABELS
    )
    model = model.to(device)
    model.eval()

    dataset = CognitiveDistortionDataset(test_df, tokenizer, MAX_LEN)
    loader  = DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=False)

    # ── Inference ─────────────────────────────────────────────────────────────
    all_probs  = []
    all_preds  = []
    all_labels = []

    with torch.no_grad():
        for batch in loader:
            input_ids      = batch["input_ids"].to(device)
            attention_mask = batch["attention_mask"].to(device)
            labels         = batch["labels"]

            outputs = model(input_ids=input_ids, attention_mask=attention_mask)
            probs   = torch.sigmoid(outputs.logits)
            preds   = (probs >= THRESHOLD).int()

            all_probs.append(probs.cpu().numpy())
            all_preds.append(preds.cpu().numpy())
            all_labels.append(labels.numpy())

    all_probs  = np.vstack(all_probs)
    all_preds  = np.vstack(all_preds)
    all_labels = np.vstack(all_labels).astype(int)

    # ── Metrics ───────────────────────────────────────────────────────────────
    per_label_f1  = f1_score(all_labels, all_preds, average=None,    zero_division=0)
    per_label_p   = precision_score(all_labels, all_preds, average=None, zero_division=0)
    per_label_r   = recall_score(all_labels, all_preds, average=None,    zero_division=0)
    micro_f1      = f1_score(all_labels, all_preds, average="micro",  zero_division=0)
    macro_f1      = f1_score(all_labels, all_preds, average="macro",  zero_division=0)
    h_loss        = hamming_loss(all_labels, all_preds)

    # ── Print comparison table ────────────────────────────────────────────────
    baseline    = load_baseline()
    base_labels = baseline.get("per_label_f1", {})

    print("\n" + "=" * 72)
    print("DISTILBERT vs BASELINE — TEST SET EVALUATION")
    print("=" * 72)
    print(f"  {'Label':<35} {'Base F1':>7} {'P':>6} {'R':>6} {'F1':>6} {'Delta':>7}")
    print(f"  {'-'*35} {'-'*7} {'-'*6} {'-'*6} {'-'*6} {'-'*7}")
    for label, p, r, f1 in zip(DISTORTION_LABELS, per_label_p, per_label_r, per_label_f1):
        base  = base_labels.get(label, float("nan"))
        delta = f1 - base if not np.isnan(base) else float("nan")
        sign  = "+" if delta >= 0 else ""
        d_str = f"{sign}{delta:.4f}" if not np.isnan(delta) else "   N/A"
        print(f"  {label:<35} {base:>7.4f} {p:>6.4f} {r:>6.4f} {f1:>6.4f} {d_str:>7}")

    print("=" * 72)
    print(f"  {'Micro F1':<35} {baseline.get('micro_f1', float('nan')):>7.4f} {'':>6} {'':>6} {micro_f1:>6.4f}")
    print(f"  {'Macro F1':<35} {baseline.get('macro_f1', float('nan')):>7.4f} {'':>6} {'':>6} {macro_f1:>6.4f}")
    print(f"  {'Hamming Loss':<35} {baseline.get('hamming_loss', float('nan')):>7.4f} {'':>6} {'':>6} {h_loss:>6.4f}")
    print("=" * 72)

    # ── Save results JSON ─────────────────────────────────────────────────────
    os.makedirs(RESULTS_DIR, exist_ok=True)
    results = {
        "model"        : "DistilBERT fine-tuned",
        "threshold"    : THRESHOLD,
        "micro_f1"     : round(float(micro_f1), 4),
        "macro_f1"     : round(float(macro_f1), 4),
        "hamming_loss" : round(float(h_loss),   4),
        "per_label_f1" : {l: round(float(s), 4) for l, s in zip(DISTORTION_LABELS, per_label_f1)},
        "per_label_precision": {l: round(float(s), 4) for l, s in zip(DISTORTION_LABELS, per_label_p)},
        "per_label_recall"   : {l: round(float(s), 4) for l, s in zip(DISTORTION_LABELS, per_label_r)},
    }
    results_path = os.path.join(RESULTS_DIR, "distilbert_results.json")
    with open(results_path, "w") as f:
        json.dump(results, f, indent=2)
    print(f"\nSaved: {results_path}")

    # ── Save per-example predictions CSV (for error analysis notebook) ─────────
    pred_df = test_df[["text"]].copy()
    for i, label in enumerate(DISTORTION_LABELS):
        short = label.replace("/", "_").replace(" ", "_").replace("-", "_")
        pred_df[f"true_{short}"]  = all_labels[:, i]
        pred_df[f"pred_{short}"]  = all_preds[:, i]
        pred_df[f"prob_{short}"]  = all_probs[:, i].round(4)

    preds_path = os.path.join(RESULTS_DIR, "test_predictions.csv")
    pred_df.to_csv(preds_path, index=False)
    print(f"Saved: {preds_path}")
    print("Evaluation complete.")


if __name__ == "__main__":
    run()