scripts/evaluate_mvsa_multiple.py

#!/usr/bin/env python3
"""Evaluate MVSA-Multiple checkpoint with notebook-style variants and ablation."""

from __future__ import annotations

import argparse
import json
from pathlib import Path

import numpy as np
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import accuracy_score, f1_score
from transformers import CLIPProcessor, DebertaV2Tokenizer

import sys

PROJECT_ROOT = Path(__file__).resolve().parents[1]
if str(PROJECT_ROOT) not in sys.path:
    sys.path.insert(0, str(PROJECT_ROOT))

from src.mvsa_multiple_pipeline import (
    MVSALoader,
    apply_bias_temp_neutral,
    build_classification_outputs,
    create_dataloaders,
    evaluate_ablation,
    evaluate_all_variants,
    evaluate_raw,
    gather_logits_labels,
    load_checkpoint,
    resolve_device,
    save_summary_files,
    summarize_splits,
)


def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(description="Evaluate CLARA MVSA-Multiple checkpoint")
    parser.add_argument("--checkpoint", default="outputs/mvsa_multiple/clara_mvsa_multiple.pt")
    parser.add_argument("--data-root", default="data/MVSA-Multiple")
    parser.add_argument("--text-dir", default=None, help="Default: <data-root>/data")
    parser.add_argument("--label-file", default=None, help="Default: <data-root>/labelResultAll.txt")
    parser.add_argument("--output-dir", default="results/mvsa_multiple")

    parser.add_argument("--batch-size", type=int, default=None)
    parser.add_argument("--max-length", type=int, default=None)
    parser.add_argument("--num-workers", type=int, default=None)
    parser.add_argument("--train-ratio", type=float, default=None)
    parser.add_argument("--val-ratio", type=float, default=None)
    parser.add_argument("--seed", type=int, default=None, help="Override split seed from checkpoint config")
    parser.add_argument(
        "--allow-seed-mismatch",
        action="store_true",
        help="Allow using a split seed different from checkpoint training seed (can cause leakage-like overlap).",
    )
    parser.add_argument("--preprocessing-mode", choices=["paper", "strict"], default=None)
    parser.add_argument("--disable-paper-exact-counts", action="store_true")

    parser.add_argument("--top2-eps", type=float, default=0.03)
    parser.add_argument(
        "--classification-mode",
        choices=["raw", "bias_temp", "logreg_calib"],
        default="raw",
        help="Prediction source for classification report/confusion matrix",
    )
    parser.add_argument(
        "--ablation-full-mode",
        choices=["raw", "bias_temp", "logreg_calib"],
        default="raw",
        help="Metric source for Full row in ablation table",
    )
    parser.add_argument(
        "--full-only",
        action="store_true",
        help="Fast path: evaluate only raw Full metrics (used for sweep tables).",
    )
    parser.add_argument("--device", default="auto", help="auto|cuda|cpu")
    return parser.parse_args()


def main() -> None:
    args = parse_args()

    text_dir = args.text_dir or str(Path(args.data_root) / "data")
    label_file = args.label_file or str(Path(args.data_root) / "labelResultAll.txt")

    device = resolve_device(args.device)
    model, cfg, ckpt_meta = load_checkpoint(args.checkpoint, device)

    cfg["text_dir"] = text_dir
    cfg["label_file"] = label_file

    if args.batch_size is not None:
        cfg["batch_size"] = args.batch_size
    if args.num_workers is not None:
        cfg["num_workers"] = args.num_workers
    if args.max_length is not None:
        cfg["max_length"] = args.max_length
    if args.train_ratio is not None:
        cfg["train_ratio"] = args.train_ratio
    if args.val_ratio is not None:
        cfg["val_ratio"] = args.val_ratio
    if args.seed is not None:
        ckpt_seed = cfg.get("seed")
        if ckpt_seed is not None and int(args.seed) != int(ckpt_seed) and not args.allow_seed_mismatch:
            raise ValueError(
                "Seed mismatch detected: "
                f"checkpoint seed={ckpt_seed}, eval seed={args.seed}. "
                "Use --allow-seed-mismatch to override explicitly."
            )
        if ckpt_seed is not None and int(args.seed) != int(ckpt_seed):
            print(
                "WARNING: evaluating with different split seed "
                f"(checkpoint={ckpt_seed}, eval={args.seed})."
            )
        cfg["seed"] = int(args.seed)
    if args.preprocessing_mode is not None:
        cfg["preprocessing_mode"] = args.preprocessing_mode
    if args.disable_paper_exact_counts:
        cfg["paper_exact_counts"] = False

    loader = MVSALoader(cfg["text_dir"], cfg["label_file"])
    loader.load(
        preprocessing_mode=str(cfg.get("preprocessing_mode", "paper")),
        require_unanimous=bool(cfg.get("require_unanimous", True)),
        require_cross_agree=bool(cfg.get("require_cross_agree", True)),
        paper_exact_counts=bool(cfg.get("paper_exact_counts", False)),
    )
    train_samples, val_samples, test_samples = loader.split(
        train_ratio=float(cfg.get("train_ratio", 0.7)),
        val_ratio=float(cfg.get("val_ratio", 0.15)),
        seed=int(cfg.get("seed", 42)),
        paper_811=bool(str(cfg.get("preprocessing_mode", "paper")).lower() == "paper"),
    )

    if not val_samples or not test_samples:
        raise RuntimeError("Need both val and test splits for full evaluation.")

    split_stats = summarize_splits(train_samples, val_samples, test_samples)
    print("Split stats:")
    print(json.dumps(split_stats, indent=2))

    clip_processor = CLIPProcessor.from_pretrained(cfg["vision_model_id"])
    tokenizer = DebertaV2Tokenizer.from_pretrained(cfg["text_model_id"])

    pin_memory = bool(cfg.get("pin_memory", True) and device.type == "cuda")
    _, val_loader, test_loader = create_dataloaders(
        train_samples=train_samples,
        val_samples=val_samples,
        test_samples=test_samples,
        clip_processor=clip_processor,
        tokenizer=tokenizer,
        batch_size=int(cfg["batch_size"]),
        max_length=int(cfg["max_length"]),
        num_workers=int(cfg["num_workers"]),
        pin_memory=pin_memory,
        persistent_workers=bool(cfg.get("persistent_workers", True)),
        prefetch_factor=int(cfg.get("prefetch_factor", 2)),
        use_mixup_negative=False,
        mixup_alpha=float(cfg.get("mixup_alpha", 0.4)),
        negative_class_boost=float(cfg.get("negative_class_boost", 12.0)),
        min_ratio_negative=float(cfg.get("min_ratio_negative", 0.30)),
        weighted_train_sampler=False,
    )

    y_pred_raw, y_true_raw, _ = evaluate_raw(model, test_loader, device)
    cls_outputs = build_classification_outputs(y_true=y_true_raw, y_pred=y_pred_raw)

    if args.full_only:
        raw_acc = float(accuracy_score(y_true_raw, y_pred_raw))
        raw_f1w = float(f1_score(y_true_raw, y_pred_raw, average="weighted"))
        result = {
            "summary": [
                {"variant": "Raw", "accuracy": raw_acc, "f1_weighted": raw_f1w}
            ],
            "best_variant": "Raw",
            "best_f1_weighted": raw_f1w,
            "tuning": {},
        }
        ablation = {
            "rows": [
                {"variant": "Full", "accuracy": raw_acc, "f1_weighted": raw_f1w}
            ],
            "full_is_highest": True,
        }
    else:
        result = evaluate_all_variants(
            model=model,
            val_loader=val_loader,
            test_loader=test_loader,
            device=device,
            top2_eps=args.top2_eps,
        )

        if args.classification_mode == "bias_temp":
            logits_test, y_true_bt = gather_logits_labels(model, test_loader, device)
            bias = float(result["tuning"]["bias_temp"]["bias"])
            tau = float(result["tuning"]["bias_temp"]["tau"])
            adjusted = apply_bias_temp_neutral(logits_test, bias=bias, tau=tau)
            y_pred_bt = adjusted.argmax(axis=-1)
            cls_outputs = build_classification_outputs(y_true=y_true_bt, y_pred=y_pred_bt)
        elif args.classification_mode == "logreg_calib":
            logits_val, y_val = gather_logits_labels(model, val_loader, device)
            logits_test, y_true_lr = gather_logits_labels(model, test_loader, device)
            c_value = float(result["tuning"]["logreg_calib"]["C"])
            clf = LogisticRegression(
                solver="lbfgs",
                max_iter=4000,
                C=c_value,
            )
            clf.fit(logits_val, y_val)
            y_pred_lr = clf.predict(logits_test)
            cls_outputs = build_classification_outputs(y_true=y_true_lr, y_pred=y_pred_lr)

        ablation = evaluate_ablation(model=model, loader=test_loader, device=device)
        if args.ablation_full_mode == "bias_temp":
            full_row = next((row for row in ablation["rows"] if row["variant"] == "Full"), None)
            bias_temp_row = next((row for row in result["summary"] if row["variant"] == "Bias+Temp"), None)
            if full_row is not None and bias_temp_row is not None:
                full_row["accuracy"] = float(bias_temp_row["accuracy"])
                full_row["f1_weighted"] = float(bias_temp_row["f1_weighted"])
        elif args.ablation_full_mode == "logreg_calib":
            full_row = next((row for row in ablation["rows"] if row["variant"] == "Full"), None)
            logreg_row = next((row for row in result["summary"] if row["variant"] == "LogReg Calib"), None)
            if full_row is not None and logreg_row is not None:
                full_row["accuracy"] = float(logreg_row["accuracy"])
                full_row["f1_weighted"] = float(logreg_row["f1_weighted"])

        full_row = next((row for row in ablation["rows"] if row["variant"] == "Full"), None)
        ablation["full_is_highest"] = (
            all(full_row["f1_weighted"] >= row["f1_weighted"] for row in ablation["rows"] if row["variant"] != "Full")
            if full_row is not None
            else False
        )

    payload = {
        "checkpoint": str(Path(args.checkpoint).resolve()),
        "checkpoint_epoch": ckpt_meta.get("epoch"),
        "best_val_f1_weighted": ckpt_meta.get("best_val_f1_weighted"),
        "config": cfg,
        "classification_mode": args.classification_mode,
        "ablation_full_mode": args.ablation_full_mode,
        "classification": cls_outputs,
        "ablation": ablation,
        **result,
    }

    json_path, csv_path = save_summary_files(args.output_dir, payload)

    out_dir = Path(args.output_dir)
    out_dir.mkdir(parents=True, exist_ok=True)

    report_path = out_dir / "classification_report.txt"
    report_path.write_text(cls_outputs["classification_report_text"], encoding="utf-8")

    cm_path = out_dir / "confusion_matrix.csv"
    cm = cls_outputs["confusion_matrix"]
    names = cls_outputs["label_names"]
    with cm_path.open("w", encoding="utf-8") as f:
        f.write("," + ",".join(names) + "\n")
        for idx, row in enumerate(cm):
            f.write(names[idx] + "," + ",".join(str(x) for x in row) + "\n")

    ablation_csv = out_dir / "ablation_summary.csv"
    with ablation_csv.open("w", encoding="utf-8") as f:
        f.write("variant,accuracy,f1_weighted\n")
        for row in ablation["rows"]:
            f.write(f"{row['variant']},{row['accuracy']:.6f},{row['f1_weighted']:.6f}\n")

    print("\nEvaluation summary:")
    for row in payload["summary"]:
        print(
            f"- {row['variant']:<14} | Acc={row['accuracy']:.4f} | "
            f"F1-Weighted={row['f1_weighted']:.4f}"
        )
    print(
        f"Best variant: {payload['best_variant']} "
        f"(F1-Weighted={payload['best_f1_weighted']:.4f})"
    )

    print(f"\nClassification report ({args.classification_mode}):")
    print(cls_outputs["classification_report_text"])

    print("Ablation summary (test):")
    for row in ablation["rows"]:
        print(
            f"- {row['variant']:<18} | Acc={row['accuracy']:.4f} | "
            f"F1-Weighted={row['f1_weighted']:.4f}"
        )
    print(f"Full highest by F1-Weighted: {ablation['full_is_highest']}")

    print(f"Saved JSON: {json_path}")
    print(f"Saved CSV:  {csv_path}")
    print(f"Saved classification report: {report_path}")
    print(f"Saved confusion matrix:      {cm_path}")
    print(f"Saved ablation summary:      {ablation_csv}")


if __name__ == "__main__":
    main()