"""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()