""" FRANKENSTALLM 3B — Re-evaluation Pipeline ========================================== Re-runs Phase 2 benchmarks with corrected task names and English benchmarks, then regenerates reports with the fixed report_generator. Reuses: - HF checkpoint from previous eval run (Phase 0 skip) - phase1_results.json from previous eval run (Phase 1 skip) Usage: python eval/reeval_pipeline.py python eval/reeval_pipeline.py --dry-run python eval/reeval_pipeline.py --skip-phase2 # regenerate reports only python eval/reeval_pipeline.py --prev-run eval/outputs/3b_full_eval_20260305_0318 """ from __future__ import annotations import argparse import json import logging import multiprocessing as mp import shutil import sys import time import traceback from datetime import datetime from pathlib import Path from typing import Any, Dict, List, Optional # --------------------------------------------------------------------------- # Project root # --------------------------------------------------------------------------- _PROJECT_ROOT = Path(__file__).resolve().parent.parent if str(_PROJECT_ROOT) not in sys.path: sys.path.insert(0, str(_PROJECT_ROOT)) from eval.full_eval_pipeline import ( _bar, _build_phase2_tasks, _fmt_seconds, _print_banner, _print_phase_header, _save_json, _spawn_phase2_batch, _spawn_task, _wait_and_collect, CHECKPOINT, TOKENIZER_PATH, ) # --------------------------------------------------------------------------- # Logging # --------------------------------------------------------------------------- logging.basicConfig( level=logging.INFO, format="%(asctime)s [%(levelname)s] %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) logger = logging.getLogger("reeval") # --------------------------------------------------------------------------- # Defaults # --------------------------------------------------------------------------- _DEFAULT_PREV_RUN = _PROJECT_ROOT / "eval" / "outputs" / "3b_full_eval_20260305_0318" def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser( description="FRANKENSTALLM 3B — Re-evaluation Pipeline", formatter_class=argparse.RawDescriptionHelpFormatter, ) parser.add_argument( "--prev-run", type=str, default=str(_DEFAULT_PREV_RUN), help="Previous eval run directory (for HF checkpoint and phase1 results).", ) parser.add_argument( "--output-dir", type=str, default=None, help="Output directory. Default: eval/outputs/3b_reeval_YYYYMMDD_HHMM/", ) parser.add_argument( "--skip-phase2", action="store_true", help="Skip Phase 2 (re-run reports only from existing phase2 results).", ) parser.add_argument( "--dry-run", action="store_true", help="Print task distribution without running.", ) parser.add_argument( "--gpus", type=str, default=None, help="Comma-separated GPU IDs. Default: all 8 GPUs (0-7).", ) return parser.parse_args() def _find_hf_checkpoint(prev_run: Path) -> Optional[Path]: """Locate HF checkpoint in previous run directory.""" candidates = list(prev_run.glob("hf_3b_*")) if candidates: return candidates[0] # Search one level up candidates = list(prev_run.parent.glob("**/hf_3b_*")) return candidates[0] if candidates else None def _copy_phase1_artifacts(prev_run: Path, output_dir: Path) -> Dict[str, Any]: """Copy phase1_results.json and generation_samples.json from previous run.""" phase1_src = prev_run / "phase1_results.json" phase1_dst = output_dir / "phase1_results.json" if not phase1_src.exists(): raise FileNotFoundError(f"phase1_results.json not found in {prev_run}") shutil.copy2(phase1_src, phase1_dst) logger.info(" Copied phase1_results.json from previous run") # Also copy generation_samples if available gen_src = prev_run / "generation_samples.json" if gen_src.exists(): shutil.copy2(gen_src, output_dir / "generation_samples.json") logger.info(" Copied generation_samples.json") with open(phase1_dst, encoding="utf-8") as f: return json.load(f) def _copy_existing_reports(prev_run: Path, output_dir: Path) -> None: """Copy existing report files that won't change (perplexity, calibration, generation).""" prev_reports = prev_run / "reports" if not prev_reports.exists(): return dst_reports = output_dir / "reports" dst_reports.mkdir(parents=True, exist_ok=True) # These are unchanged — copy as backup reference for fname in ["01_perplexity_report.md", "02_calibration_report.md", "03_generation_quality_report.md"]: src = prev_reports / fname if src.exists(): shutil.copy2(src, dst_reports / fname) def run_phase2_reeval( hf_model_path: Path, output_dir: Path, gpu_ids: List[int], ) -> Dict[str, Any]: """Run Phase 2 benchmarks with per-GPU pipelining. Korean GPUs run 0-shot then 5-shot in the SAME process (model loaded once). English GPUs run 0-shot only. All GPUs are spawned in a single batch so there is no barrier between 0-shot and 5-shot — early-finishing Korean GPUs start 5-shot immediately while slow GPUs (e.g. MMLU-EN) are still running their 0-shot. """ gpu_task_list = _build_phase2_tasks(gpu_ids) _KO_PREFIXES = ("kobest", "haerae", "global_mmlu_ko") processes: list = [] for gpu_id, task_names, label in gpu_task_list: is_korean = any(t.startswith(_KO_PREFIXES) for t in task_names) out_path = output_dir / f"phase2_gpu{gpu_id}_pipeline_reeval.json" if is_korean: extra_args = { "--hf-model-path": str(hf_model_path), "--lm-eval-tasks": ",".join(task_names), "--fewshot-list": "0,5", } spawn_label = f"[pipeline 0+5shot] {label}" else: extra_args = { "--hf-model-path": str(hf_model_path), "--lm-eval-tasks": ",".join(task_names), "--num-fewshot": "0", } spawn_label = f"[0-shot] {label}" proc_info = _spawn_task( task_name="lm_eval", gpu_id=gpu_id, output_path=out_path, label=spawn_label, extra_args=extra_args, ) processes.append(proc_info) logger.info( " Spawned %d GPUs (Korean GPUs run 0+5-shot pipeline, EN GPUs 0-shot only).", len(gpu_task_list), ) raw_results = _wait_and_collect(processes) # --- Reorganise into the expected output format --- # Pipeline results come as {"0shot": {...}, "5shot": {...}} # Non-pipeline results come as a flat dict (single fewshot). all_results: Dict[str, Any] = {} five_shot_bucket: Dict[str, Any] = {} for label, data in raw_results.items(): if isinstance(data, dict) and "error" not in data and "0shot" in data: # Pipeline result — split into 0-shot and 5-shot zero_label = label.replace("[pipeline 0+5shot]", "[0-shot]") all_results[zero_label] = data["0shot"] if "5shot" in data and "error" not in data.get("5shot", {}): five_label = label.replace("[pipeline 0+5shot]", "[5-shot]") five_shot_bucket[five_label] = data["5shot"] else: all_results[label] = data if five_shot_bucket: all_results["5shot"] = five_shot_bucket _save_json(all_results, output_dir / "phase2_results.json") logger.info(" Phase 2 results saved: %s", output_dir / "phase2_results.json") return all_results def run_phase3_reeval( phase1_results: Dict[str, Any], phase2_results: Dict[str, Any], output_dir: Path, total_elapsed_sec: float = 0.0, ) -> Optional[Path]: """Generate reports using the fixed report_generator.""" try: from eval.report_generator import generate_report # Extract generation samples from phase1_results gen_samples = [] for label, result in phase1_results.items(): if isinstance(result, dict) and "Generation" in label: if "samples" in result: gen_samples = result["samples"] break generate_report( phase1_results=phase1_results, phase2_results=phase2_results, generation_samples=gen_samples, output_dir=output_dir, checkpoint_name=Path(CHECKPOINT).name, total_elapsed_sec=total_elapsed_sec, ) report_path = output_dir / "full_eval_report.md" logger.info(" Report saved: %s", report_path) logger.info(" Individual reports: %s", output_dir / "reports") return report_path except Exception: logger.error(" Report generation failed:\n%s", traceback.format_exc()) return None def main() -> None: try: mp.set_start_method("spawn", force=True) except RuntimeError: pass args = parse_args() # Parse GPU IDs if args.gpus: gpu_ids = sorted([int(g.strip()) for g in args.gpus.split(",")]) else: gpu_ids = list(range(8)) # Resolve paths prev_run = Path(args.prev_run) if args.output_dir: output_dir = Path(args.output_dir) else: timestamp = datetime.now().strftime("%Y%m%d_%H%M") output_dir = _PROJECT_ROOT / "eval" / "outputs" / f"3b_reeval_{timestamp}" output_dir.mkdir(parents=True, exist_ok=True) # Find HF checkpoint hf_model_path = _find_hf_checkpoint(prev_run) # --------------------------------------------------------------------------- # Dry run # --------------------------------------------------------------------------- if args.dry_run: _print_banner("DRY RUN — FRANKENSTALLM 3B Re-evaluation Pipeline") logger.info(" Previous run : %s", prev_run) logger.info(" HF checkpoint: %s", hf_model_path or "NOT FOUND") logger.info(" Output dir : %s", output_dir) logger.info(" GPUs : %s", gpu_ids) _print_phase_header("Phase 2", f"Corrected Benchmarks ({len(gpu_ids)} GPUs)") gpu_task_list = _build_phase2_tasks(gpu_ids) logger.info(" %-6s %-60s", "GPU", "Tasks") logger.info(" %s %s", "-" * 6, "-" * 60) for gpu_id, tasks, label in gpu_task_list: logger.info(" cuda:%-2d %s (%d tasks)", gpu_id, label, len(tasks)) for t in tasks: logger.info(" - %s", t) total_tasks = sum(len(tasks) for _, tasks, _ in gpu_task_list) logger.info("") logger.info(" Total benchmark tasks: %d", total_tasks) logger.info(" Estimated time: ~80 min (0-shot ~40 min + 5-shot ~40 min)") logger.info("") logger.info(" Dry run complete.") sys.exit(0) # --------------------------------------------------------------------------- # Banner # --------------------------------------------------------------------------- _print_banner("FRANKENSTALLM 3B — Re-evaluation Pipeline") logger.info(" Previous run : %s", prev_run) logger.info(" HF checkpoint: %s", hf_model_path) logger.info(" Output dir : %s", output_dir) logger.info(" GPUs : %s", gpu_ids) logger.info(" Skip Phase 2 : %s", args.skip_phase2) pipeline_start = time.time() # --------------------------------------------------------------------------- # Phase 1 — Copy from previous run # --------------------------------------------------------------------------- _print_phase_header("PHASE 1", "Copy from Previous Run") try: phase1_results = _copy_phase1_artifacts(prev_run, output_dir) except FileNotFoundError as exc: logger.error(" %s", exc) sys.exit(1) # --------------------------------------------------------------------------- # Phase 2 — Corrected Benchmarks # --------------------------------------------------------------------------- phase2_results: Dict[str, Any] = {} _print_phase_header("PHASE 2", f"Corrected Benchmarks — {len(gpu_ids)} GPU Parallel") if args.skip_phase2: logger.info(" Skipping Phase 2.") phase2_out = output_dir / "phase2_results.json" if phase2_out.exists(): with open(phase2_out, encoding="utf-8") as f: phase2_results = json.load(f) logger.info(" Loaded existing Phase 2 results: %s", phase2_out) else: # Try previous run prev_p2 = prev_run / "phase2_results.json" if prev_p2.exists(): shutil.copy2(prev_p2, phase2_out) with open(phase2_out, encoding="utf-8") as f: phase2_results = json.load(f) logger.info(" Copied Phase 2 results from previous run") elif hf_model_path is None: logger.error(" HF checkpoint not found in %s — cannot run Phase 2", prev_run) sys.exit(1) else: t0 = time.time() try: phase2_results = run_phase2_reeval(hf_model_path, output_dir, gpu_ids) logger.info(" Phase 2 complete in %s.", _fmt_seconds(time.time() - t0)) except Exception: logger.error(" Phase 2 FAILED:\n%s", traceback.format_exc()) # --------------------------------------------------------------------------- # Phase 3 — Report Generation # --------------------------------------------------------------------------- _print_phase_header("PHASE 3", "Report Generation (Fixed)") t0 = time.time() report_path = run_phase3_reeval( phase1_results, phase2_results, output_dir, total_elapsed_sec=time.time() - pipeline_start, ) logger.info(" Phase 3 complete in %s.", _fmt_seconds(time.time() - t0)) # --------------------------------------------------------------------------- # Final Summary # --------------------------------------------------------------------------- total_elapsed = time.time() - pipeline_start _print_banner("RE-EVALUATION COMPLETE") logger.info(" Total time : %s", _fmt_seconds(total_elapsed)) logger.info(" Output dir : %s", output_dir) logger.info(" Phase 1 results : %s", output_dir / "phase1_results.json") logger.info(" Phase 2 results : %s", output_dir / "phase2_results.json") logger.info(" Report : %s", report_path or "N/A") logger.info(" Reports dir : %s", output_dir / "reports") if phase2_results: p2_entries = {k: v for k, v in phase2_results.items() if k != "5shot"} p2_ok = sum(1 for v in p2_entries.values() if not (isinstance(v, dict) and "error" in v)) p2_fail = len(p2_entries) - p2_ok logger.info(" Phase 2 (0-shot): %d OK / %d failed", p2_ok, p2_fail) five_shot = phase2_results.get("5shot", {}) if isinstance(five_shot, dict) and "error" not in five_shot: fs_ok = sum(1 for v in five_shot.values() if not (isinstance(v, dict) and "error" in v)) logger.info(" Phase 2 (5-shot): %d OK", fs_ok) logger.info(_bar()) if __name__ == "__main__": main()