scripts/03_training/evaluate.py

"""Run BabyLM evaluation pipeline on an exported HF model.

Calls the evaluation pipeline Python modules directly (not shell scripts)
to run fast zero-shot evaluation (BLiMP, supplement, EWoK, entity_tracking,
wug_past, wug_adj, reading) and parse results into a JSON summary.

Usage (standalone):
    python -m scripts.03_training.evaluate \
        --model_path models/hf_export/exp_A \
        --backend causal \
        --eval_mode fast

Programmatic usage (from train.py):
    from scripts.03_training.evaluate import run_evaluation
    results = run_evaluation(model_path, backend="causal", eval_mode="fast")
"""

import json
import os
import re
import subprocess
import sys
from pathlib import Path
from typing import Optional


ROOT = Path(__file__).resolve().parent.parent.parent
EVAL_PIPELINE_DIR = ROOT / "evaluation-pipeline-2025"
EVAL_DATA_DIR = EVAL_PIPELINE_DIR / "evaluation_data"


# ═══════════════════════════════════════════════════════════════════════
# Task definitions
# ═══════════════════════════════════════════════════════════════════════

FAST_EVAL_TASKS = [
    # (task_name, data_subdir, module, extra_args)
    ("blimp", "fast_eval/blimp_fast", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "blimp"]),
    ("supplement", "fast_eval/supplement_fast", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "blimp"]),
    ("ewok", "fast_eval/ewok_fast", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "ewok"]),
    ("entity_tracking", "fast_eval/entity_tracking_fast", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "entity_tracking"]),
    ("wug_past", "fast_eval/wug_past_tense", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "wug_past"]),
    ("wug_adj", "fast_eval/wug_adj_nominalization", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "wug_adj"]),
    ("reading", "fast_eval/reading/reading_data.csv", "evaluation_pipeline.reading.run",
     []),
]

FULL_EVAL_TASKS = [
    ("blimp", "full_eval/blimp_filtered", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "blimp"]),
    ("supplement", "full_eval/supplement_filtered", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "blimp"]),
    ("ewok", "full_eval/ewok_filtered", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "ewok"]),
    ("entity_tracking", "full_eval/entity_tracking", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "entity_tracking"]),
    ("wug_past", "full_eval/wug_past_tense", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "wug_past"]),
    ("wug_adj", "full_eval/wug_adj_nominalization", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "wug_adj"]),
    ("comps", "full_eval/comps", "evaluation_pipeline.sentence_zero_shot.run",
     ["--task", "comps"]),
    ("reading", "full_eval/reading/reading_data.csv", "evaluation_pipeline.reading.run",
     []),
]

# GLUE finetuning tasks: (task_name, train_data, valid_data, num_labels, batch_size, epochs, metric_for_valid)
GLUE_TASKS = [
    ("boolq",   "glue_filtered/boolq.train.jsonl",   "glue_filtered/boolq.valid.jsonl",   2, 16, 10, "accuracy"),
    ("multirc", "glue_filtered/multirc.train.jsonl", "glue_filtered/multirc.valid.jsonl", 2, 16, 10, "accuracy"),
    ("rte",     "glue_filtered/rte.train.jsonl",     "glue_filtered/rte.valid.jsonl",     2, 32, 10, "accuracy"),
    ("wsc",     "glue_filtered/wsc.train.jsonl",     "glue_filtered/wsc.valid.jsonl",     2, 32, 30, "accuracy"),
    ("mrpc",    "glue_filtered/mrpc.train.jsonl",    "glue_filtered/mrpc.valid.jsonl",    2, 32, 10, "f1"),
    ("qqp",     "glue_filtered/qqp.train.jsonl",     "glue_filtered/qqp.valid.jsonl",     2, 32, 10, "f1"),
    ("mnli",    "glue_filtered/mnli.train.jsonl",    "glue_filtered/mnli.valid.jsonl",    3, 32, 10, "accuracy"),
]


# ═══════════════════════════════════════════════════════════════════════
# Parse eval output
# ═══════════════════════════════════════════════════════════════════════

def _parse_accuracy_from_output(output: str, task_name: str) -> Optional[float]:
    """Parse the average accuracy from eval script stdout.

    The eval pipeline prints lines like:
        1.0     72.35

    (temperature, accuracy) and then a detailed report with:
        ### AVERAGE ACCURACY
        72.35

    We extract the first temperature line (which is the best accuracy).
    For reading tasks, we look for "EYE TRACKING SCORE:" and "SELF-PACED READING SCORE:".
    """
    if task_name == "reading":
        # Look for eye tracking and self-paced reading scores
        scores = {}
        for line in output.split("\n"):
            if "EYE TRACKING SCORE:" in line:
                match = re.search(r"EYE TRACKING SCORE:\s*([-\d.]+)", line)
                if match:
                    scores["eye_tracking"] = float(match.group(1))
            elif "SELF-PACED READING SCORE:" in line:
                match = re.search(r"SELF-PACED READING SCORE:\s*([-\d.]+)", line)
                if match:
                    scores["self_paced_reading"] = float(match.group(1))
        return scores if scores else None

    # For standard tasks, look for temperature lines: "1.0\t72.35"
    best_acc = None
    for line in output.split("\n"):
        parts = line.strip().split("\t")
        if len(parts) == 2:
            try:
                _temp = float(parts[0])
                acc = float(parts[1])
                if best_acc is None or acc > best_acc:
                    best_acc = acc
            except ValueError:
                continue

    # Also look for "### AVERAGE ACCURACY" block
    if best_acc is None:
        match = re.search(r"### AVERAGE (?:ACCURACY|SPEARMAN'S RHO)\s*\n\s*([-\d.]+)", output)
        if match:
            best_acc = float(match.group(1))

    return best_acc


# ═══════════════════════════════════════════════════════════════════════
# Run evaluation
# ═══════════════════════════════════════════════════════════════════════

def run_single_task(
    model_path: str,
    backend: str,
    task_name: str,
    data_path: str,
    module: str,
    extra_args: list,
    results_dir: str = "results",
) -> dict:
    """Run a single evaluation task via subprocess.

    Returns:
        dict with keys: task, accuracy (or scores), status, output
    """
    # Determine backend for reading tasks
    if module == "evaluation_pipeline.reading.run":
        if "enc_dec" in backend:
            read_backend = "enc_dec"
        else:
            read_backend = backend
        cmd = [
            sys.executable, "-m", module,
            "--model_path_or_name", str(model_path),
            "--backend", read_backend,
            "--data_path", str(data_path),
        ]
    else:
        cmd = [
            sys.executable, "-m", module,
            "--model_path_or_name", str(model_path),
            "--backend", backend,
            "--data_path", str(data_path),
            "--save_predictions",
            *extra_args,
        ]

    print(f"  Running {task_name}...", end=" ", flush=True)

    try:
        result = subprocess.run(
            cmd,
            cwd=str(EVAL_PIPELINE_DIR),
            capture_output=True,
            text=True,
            timeout=1800,  # 30 minutes per task
        )

        combined_output = result.stdout + "\n" + result.stderr
        accuracy = _parse_accuracy_from_output(combined_output, task_name)

        if result.returncode != 0:
            print(f"FAILED (rc={result.returncode})")
            # Print first few lines of stderr for debugging
            stderr_lines = result.stderr.strip().split("\n")
            for line in stderr_lines[-5:]:
                print(f"    {line}")
            return {
                "task": task_name,
                "accuracy": None,
                "status": "failed",
                "returncode": result.returncode,
                "stderr_tail": "\n".join(stderr_lines[-5:]),
            }

        if isinstance(accuracy, dict):
            print(f"OK ({accuracy})")
        elif accuracy is not None:
            print(f"OK ({accuracy:.2f})")
        else:
            print("OK (score not parsed)")

        return {
            "task": task_name,
            "accuracy": accuracy,
            "status": "completed",
        }

    except subprocess.TimeoutExpired:
        print("TIMEOUT")
        return {"task": task_name, "accuracy": None, "status": "timeout"}
    except Exception as e:
        print(f"ERROR: {e}")
        return {"task": task_name, "accuracy": None, "status": "error", "error": str(e)}


def run_glue_task(model_path: str, task_name: str, train_data: str, valid_data: str,
                  num_labels: int, batch_size: int, epochs: int, metric_for_valid: str,
                  results_dir: str, lr: float = 3e-5, seed: int = 42) -> dict:
    """Run a single GLUE finetuning task."""
    train_path = EVAL_DATA_DIR / "full_eval" / train_data
    valid_path = EVAL_DATA_DIR / "full_eval" / valid_data

    if not train_path.exists() or not valid_path.exists():
        print(f"  GLUE {task_name}: data not found, skipping")
        return {"task": f"glue_{task_name}", "accuracy": None, "status": "data_missing"}

    cmd = [
        sys.executable, "-m", "evaluation_pipeline.finetune.run",
        "--model_name_or_path", str(model_path),
        "--train_data", str(train_path),
        "--valid_data", str(valid_path),
        "--predict_data", str(valid_path),
        "--task", task_name,
        "--num_labels", str(num_labels),
        "--batch_size", str(batch_size),
        "--learning_rate", str(lr),
        "--num_epochs", str(epochs),
        "--sequence_length", "512",
        "--results_dir", results_dir,
        "--save",
        "--save_dir", results_dir,
        "--metrics", "accuracy", "f1", "mcc",
        "--metric_for_valid", metric_for_valid,
        "--seed", str(seed),
        "--verbose",
    ]

    print(f"  Running GLUE/{task_name}...", end=" ", flush=True)
    try:
        result = subprocess.run(
            cmd, cwd=str(EVAL_PIPELINE_DIR),
            capture_output=True, text=True, timeout=3600,
        )
        combined = result.stdout + "\n" + result.stderr

        # Parse accuracy from output: look for "Best valid accuracy: X.XX" or similar
        acc = None
        for line in combined.split("\n"):
            # finetune prints: "Valid accuracy: 0.7234" or "Test accuracy: 0.7234"
            match = re.search(r"(?:valid|test|best).*?(?:accuracy|f1).*?:\s*([\d.]+)", line, re.IGNORECASE)
            if match:
                acc = float(match.group(1))
            # Also look for JSON results
            match2 = re.search(r'"accuracy":\s*([\d.]+)', line)
            if match2:
                acc = float(match2.group(1))

        # Check results JSON file
        if acc is None:
            results_json = Path(results_dir) / f"{task_name}_results.json"
            if results_json.exists():
                with open(results_json) as f:
                    data = json.load(f)
                acc = data.get("accuracy", data.get("f1"))

        if result.returncode != 0:
            print(f"FAILED")
            return {"task": f"glue_{task_name}", "accuracy": None, "status": "failed",
                    "stderr_tail": result.stderr.strip().split("\n")[-3:]}

        if acc is not None:
            # Convert to percentage if needed
            if acc < 1.0:
                acc = acc * 100.0
            print(f"OK ({acc:.1f})")
        else:
            print(f"OK (score not parsed)")

        return {"task": f"glue_{task_name}", "accuracy": acc, "status": "completed"}

    except subprocess.TimeoutExpired:
        print("TIMEOUT")
        return {"task": f"glue_{task_name}", "accuracy": None, "status": "timeout"}
    except Exception as e:
        print(f"ERROR: {e}")
        return {"task": f"glue_{task_name}", "accuracy": None, "status": "error"}


def run_aoa(model_path: str, backend: str, results_dir: str) -> dict:
    """Run AoA (Age of Acquisition) evaluation."""
    word_path = EVAL_DATA_DIR / "full_eval" / "aoa" / "cdi_childes.json"
    if not word_path.exists():
        print(f"  AoA: data not found, skipping")
        return {"task": "aoa", "accuracy": None, "status": "data_missing"}

    # Determine track name based on backend
    track_name = "strict_small"

    cmd = [
        sys.executable, "-m", "evaluation_pipeline.AoA_word.run",
        "--model_name", str(model_path),
        "--backend", backend,
        "--track_name", track_name,
        "--word_path", str(word_path),
        "--output_dir", results_dir,
    ]

    print(f"  Running AoA...", end=" ", flush=True)
    try:
        result = subprocess.run(
            cmd, cwd=str(EVAL_PIPELINE_DIR),
            capture_output=True, text=True, timeout=3600,
        )
        combined = result.stdout + "\n" + result.stderr

        # Parse AoA score: look for correlation or score
        score = None
        for line in combined.split("\n"):
            match = re.search(r"(?:correlation|spearman|aoa.*score).*?:\s*([-\d.]+)", line, re.IGNORECASE)
            if match:
                score = float(match.group(1))

        if result.returncode != 0:
            print(f"FAILED")
            return {"task": "aoa", "accuracy": None, "status": "failed",
                    "stderr_tail": result.stderr.strip().split("\n")[-3:]}

        if score is not None:
            print(f"OK ({score:.4f})")
        else:
            print(f"OK (score not parsed)")

        return {"task": "aoa", "accuracy": score, "status": "completed"}

    except subprocess.TimeoutExpired:
        print("TIMEOUT")
        return {"task": "aoa", "accuracy": None, "status": "timeout"}
    except Exception as e:
        print(f"ERROR: {e}")
        return {"task": "aoa", "accuracy": None, "status": "error"}


def run_evaluation(
    model_path: str,
    backend: str = "causal",
    eval_mode: str = "fast",
    results_dir: Optional[str] = None,
) -> dict:
    """Run the full evaluation pipeline and return results.

    Args:
        model_path: path to HF-format model directory
        backend: "causal", "mntp", "mlm", etc.
        eval_mode: "fast" or "full"
        results_dir: where to save detailed results (default: next to model)

    Returns:
        dict with per-task results and summary scores
    """
    model_path = str(Path(model_path).resolve())

    if results_dir is None:
        results_dir = str(Path(model_path) / "eval_results")
    Path(results_dir).mkdir(parents=True, exist_ok=True)

    # Select tasks
    if eval_mode == "fast":
        tasks = FAST_EVAL_TASKS
    else:
        tasks = FULL_EVAL_TASKS

    print(f"\n  Evaluation: {eval_mode} mode, backend={backend}")
    print(f"  Model: {model_path}")
    print(f"  Tasks: {len(tasks)}")

    # Check eval data exists
    if not EVAL_DATA_DIR.exists():
        print(f"  WARNING: Eval data not found at {EVAL_DATA_DIR}")
        return {"status": "no_eval_data", "tasks": {}}

    # Run each task
    task_results = {}
    for task_name, data_subdir, module, extra_args in tasks:
        data_path = EVAL_DATA_DIR / data_subdir
        if not data_path.exists():
            print(f"  Skipping {task_name}: data not found at {data_path}")
            task_results[task_name] = {"task": task_name, "accuracy": None, "status": "data_missing"}
            continue

        result = run_single_task(
            model_path, backend, task_name, str(data_path),
            module, extra_args, results_dir,
        )
        task_results[task_name] = result

    # ── GLUE finetuning (full mode only) ──
    if eval_mode == "full":
        glue_dir = str(Path(results_dir) / "glue")
        Path(glue_dir).mkdir(parents=True, exist_ok=True)
        print(f"\n  GLUE Finetuning ({len(GLUE_TASKS)} tasks):")
        for task_name, train_data, valid_data, num_labels, bsz, epochs, metric in GLUE_TASKS:
            glue_result = run_glue_task(
                model_path, task_name, train_data, valid_data,
                num_labels, bsz, epochs, metric, glue_dir,
            )
            task_results[f"glue_{task_name}"] = glue_result

    # ── AoA evaluation (full mode only) ──
    if eval_mode == "full":
        aoa_dir = str(Path(results_dir) / "aoa")
        Path(aoa_dir).mkdir(parents=True, exist_ok=True)
        print(f"\n  AoA Evaluation:")
        aoa_result = run_aoa(model_path, backend, aoa_dir)
        task_results["aoa"] = aoa_result

    # ── Compute summary ──
    summary = _compute_summary(task_results)

    # ── Save results ──
    all_results = {
        "model_path": model_path,
        "backend": backend,
        "eval_mode": eval_mode,
        "tasks": task_results,
        "summary": summary,
    }

    results_file = Path(results_dir) / "eval_results.json"
    with open(results_file, "w") as f:
        json.dump(all_results, f, indent=2, default=str)
    print(f"\n  Results saved to {results_file}")

    # ── Print summary ──
    _print_summary(summary, task_results)

    return all_results


def _compute_summary(task_results: dict) -> dict:
    """Compute summary scores from per-task results."""
    summary = {}

    # BLiMP score (main metric)
    blimp_acc = task_results.get("blimp", {}).get("accuracy")
    if blimp_acc is not None:
        summary["blimp"] = blimp_acc

    # Supplement
    supplement_acc = task_results.get("supplement", {}).get("accuracy")
    if supplement_acc is not None:
        summary["supplement"] = supplement_acc

    # EWoK
    ewok_acc = task_results.get("ewok", {}).get("accuracy")
    if ewok_acc is not None:
        summary["ewok"] = ewok_acc

    # Entity tracking
    et_acc = task_results.get("entity_tracking", {}).get("accuracy")
    if et_acc is not None:
        summary["entity_tracking"] = et_acc

    # WUG tasks (Spearman's rho, not accuracy)
    wug_past = task_results.get("wug_past", {}).get("accuracy")
    if wug_past is not None:
        summary["wug_past"] = wug_past

    wug_adj = task_results.get("wug_adj", {}).get("accuracy")
    if wug_adj is not None:
        summary["wug_adj"] = wug_adj

    # Reading
    reading = task_results.get("reading", {}).get("accuracy")
    if reading is not None:
        summary["reading"] = reading

    # COMPS
    comps_acc = task_results.get("comps", {}).get("accuracy")
    if comps_acc is not None:
        summary["comps"] = comps_acc

    # GLUE scores
    glue_scores = []
    for task_name in ["boolq", "multirc", "rte", "wsc", "mrpc", "qqp", "mnli"]:
        glue_key = f"glue_{task_name}"
        acc = task_results.get(glue_key, {}).get("accuracy")
        if acc is not None:
            summary[glue_key] = acc
            glue_scores.append(acc)
    if glue_scores:
        summary["glue_avg"] = sum(glue_scores) / len(glue_scores)

    # AoA
    aoa_score = task_results.get("aoa", {}).get("accuracy")
    if aoa_score is not None:
        summary["aoa"] = aoa_score

    # Average of available zero-shot numeric scores (excluding reading which is a dict)
    numeric_scores = []
    for key in ["blimp", "supplement", "ewok", "entity_tracking"]:
        if key in summary and isinstance(summary[key], (int, float)):
            numeric_scores.append(summary[key])

    if numeric_scores:
        summary["avg_zero_shot"] = sum(numeric_scores) / len(numeric_scores)

    return summary


def _print_summary(summary: dict, task_results: dict):
    """Print a human-readable summary."""
    print(f"\n  {'='*50}")
    print(f"  Evaluation Summary")
    print(f"  {'='*50}")

    for task_name, result in task_results.items():
        acc = result.get("accuracy")
        status = result.get("status", "?")
        if status != "completed":
            print(f"  {task_name:<20s}: {status}")
        elif isinstance(acc, dict):
            for k, v in acc.items():
                print(f"  {task_name}/{k:<15s}: {v:.2f}")
        elif acc is not None:
            print(f"  {task_name:<20s}: {acc:.2f}")
        else:
            print(f"  {task_name:<20s}: (no score)")

    avg = summary.get("avg_zero_shot")
    if avg is not None:
        print(f"  {'─'*50}")
        print(f"  {'AVG ZERO-SHOT':<20s}: {avg:.2f}")
    print(f"  {'='*50}")


# ═══════════════════════════════════════════════════════════════════════
# CLI entry point
# ═══════════════════════════════════════════════════════════════════════

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Run BabyLM evaluation")
    parser.add_argument("--model_path", required=True, help="Path to HF model directory")
    parser.add_argument("--backend", default="causal", choices=["causal", "mntp", "mlm"])
    parser.add_argument("--eval_mode", default="fast", choices=["fast", "full"],
                        help="fast=zero-shot only (~15min); full=zero-shot+GLUE+AoA (~60-90min)")
    parser.add_argument("--results_dir", default=None, help="Where to save results")
    args = parser.parse_args()

    results = run_evaluation(
        model_path=args.model_path,
        backend=args.backend,
        eval_mode=args.eval_mode,
        results_dir=args.results_dir,
    )

    # Exit with non-zero if any task failed
    failed = sum(1 for t in results.get("tasks", {}).values() if t.get("status") != "completed")
    if failed:
        print(f"\n  {failed} task(s) failed or were skipped")


if __name__ == "__main__":
    main()