eval_dev.py

# eval_dev.py
import argparse, json, os, re
from pathlib import Path

import torch
import torch.nn.functional as F
import pandas as pd
from transformers import AutoTokenizer

from model import BertClassifier  # your LightningModule

def infer_run_id_from_ckpt(ckpt_path: str) -> str:
    # tries to grab run_XXXX-YYYYMMDD-HHMMSS from the path
    m = re.search(r"(run_[^/\\]+)", ckpt_path)
    return m.group(1) if m else f"run_local-{Path(ckpt_path).stem}"

def batch_iter(texts, batch_size):
    for i in range(0, len(texts), batch_size):
        yield texts[i:i+batch_size]

def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt_path", required=True, help="Path to best.ckpt")
    ap.add_argument("--val_path",  required=True, help="Path to test/val CSV (e.g., datasets/test_sent_emo.csv)")
    ap.add_argument("--text_col",  required=True)
    ap.add_argument("--label_col", required=True)
    ap.add_argument("--batch_size", type=int, default=64)
    ap.add_argument("--out_csv",   default=None,
                    help="Optional explicit output path. If not set, writes to preds/<run_id>/test_probs.list.csv")
    args = ap.parse_args()

    # Load model + tokenizer
    model: BertClassifier = BertClassifier.load_from_checkpoint(args.ckpt_path, strict=False)
    model.eval()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)

    # Prefer tokenizer name from the checkpoint
    tok_name = getattr(model.hparams, "tokenizer_name", None) or getattr(model.hparams, "backbone_name", None)
    if not tok_name:
        raise SystemExit("Could not determine tokenizer_name/backbone_name from checkpoint hparams.")
    tokenizer = AutoTokenizer.from_pretrained(tok_name)

    # Class names for mapping predictions -> label strings
    class_names = getattr(model, "class_names", None) or getattr(model.hparams, "class_names", None)
    if class_names is None:
        raise SystemExit("No class_names found in checkpoint; cannot map predicted indices to label strings.")

    # Read dev/val split
    df = pd.read_csv(args.val_path)
    if args.text_col not in df.columns or args.label_col not in df.columns:
        raise SystemExit(f"Expected columns '{args.text_col}' and '{args.label_col}' in {args.val_path}")

    texts  = df[args.text_col].astype(str).tolist()
    golds  = df[args.label_col].astype(str).tolist()  # keep gold labels as strings

    # Decide output path
    run_id = infer_run_id_from_ckpt(args.ckpt_path)
    out_csv = Path(args.out_csv) if args.out_csv else Path("preds") / run_id / "test_predictions.csv"
    out_csv.parent.mkdir(parents=True, exist_ok=True)

    rows = []  # each row: {"Label": gold_str, "Prediction": pred_str, "Probabilities": "[...]"}

    with torch.no_grad():
        for chunk in batch_iter(texts, args.batch_size):
            enc = tokenizer(
                list(chunk),
                add_special_tokens=True,
                max_length=getattr(model.hparams, "max_length", 256),
                padding=True,
                truncation=True,
                return_tensors="pt",
            )
            enc = {k: v.to(device) for k, v in enc.items()}
            logits = model(enc["input_ids"], enc["attention_mask"])
            probs  = F.softmax(logits, dim=-1)  # [B, C]
            pred_ix = torch.argmax(probs, dim=-1).tolist()
            probs_list = probs.detach().cpu().tolist()

            for p_ix, p_vec in zip(pred_ix, probs_list):
                rows.append({
                    "Prediction": str(class_names[p_ix]),
                    "Probabilities": json.dumps(p_vec),  # single list column (JSON string)
                })

    # Attach gold labels (same order)
    out_df = pd.DataFrame(rows)
    out_df.insert(0, "Label", golds[:len(out_df)])  # ensure Label is first column

    out_df.to_csv(out_csv, index=False)
    print(f"[ok] wrote test predictions CSV -> {out_csv}")

if __name__ == "__main__":
    main()