training/eval_cefr.py

"""Evaluate any CEFR checkpoint on the canonical split — including the published
baseline, for the ADR 0003 comparison.

Usage (from the repo root):
    uv run --group train python training/eval_cefr.py \
        --model UniversalCEFR/xlm-roberta-base-cefr-all-classifier \
        --config training/configs/en_only.toml

    uv run --group train python training/eval_cefr.py \
        --model models/cefr/en_chunked_weighted/model \
        --config training/configs/en_only.toml

The --config defines the data and the split; use the same config as the run
you compare against so both models see the *same* test documents.
"""

import argparse
import json
from pathlib import Path

from evaluation import evaluate_views, lang_filtered, predict_probs
from train_cefr import build_parts, load_config


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True, help="local checkpoint dir or Hub id")
    parser.add_argument("--config", type=Path, required=True)
    parser.add_argument("--batch-size", type=int, default=64)
    args = parser.parse_args()

    config = load_config(args.config)
    parts = build_parts(config)
    test_passages = parts["test"]
    print(f"Evaluating {args.model} on {len(test_passages)} test passages")

    import mlflow
    from transformers import AutoModelForSequenceClassification, AutoTokenizer

    tokenizer = AutoTokenizer.from_pretrained(args.model)
    model = AutoModelForSequenceClassification.from_pretrained(args.model)
    if hasattr(model, "cuda"):
        import torch

        if torch.cuda.is_available():
            model = model.cuda()

    max_length = config["model"].get("max_length", 512)
    probs = predict_probs(
        model, tokenizer, test_passages, max_length=max_length, batch_size=args.batch_size
    )
    en_passages, en_probs = lang_filtered(test_passages, probs, "en")
    results = {
        "en": evaluate_views(en_passages, en_probs),
        "all": evaluate_views(test_passages, probs),
    }

    tracking = config.get("tracking", {})
    mlflow.set_tracking_uri(tracking.get("uri", "sqlite:///mlflow.db"))
    mlflow.set_experiment(tracking.get("experiment", "cefr-classifier"))
    run_name = f"eval__{args.model.replace('/', '_')}"
    with mlflow.start_run(run_name=run_name):
        mlflow.log_param("model", args.model)
        mlflow.log_param("config", str(args.config))
        for scope, views in results.items():
            for view, report in views.items():
                mlflow.log_metrics({f"test_{scope}_{view}_{k}": v for k, v in report.items()})
        mlflow.log_dict(results, "test_metrics.json")

    print(json.dumps(results, indent=2))


if __name__ == "__main__":
    main()