"""Local validation: run a checkpoint against validation_samples/multilingual.jsonl. Mirrors the nightly CI's inference + scoring logic so you can catch issues before pushing to HuggingFace. Checks performed: 1. Model loads cleanly with vLLM. 2. All 10 validation problems receive n completions. 3. Every completion contains at least one \\boxed{LETTER} pattern (boxed_rate check). 4. pass@1 metric is reported (multilingual benchmark uses pass@1, not pass@8). Usage on the cluster: python -m trainingHelena.validate \\ --checkpoint /scratch/multilingual_model_sft/merged \\ [--samples validation_samples/multilingual.jsonl] \\ [--output /scratch/multilingual_val_results.json] \\ [--n 4] """ from __future__ import annotations import argparse import json import sys from dataclasses import dataclass from pathlib import Path # --------------------------------------------------------------------------- # Config # --------------------------------------------------------------------------- @dataclass class ValidateConfig: checkpoint: str samples_file: str = "validation_samples/multilingual.jsonl" output: str | None = None n_completions: int = 4 # 4 is enough for a quick smoke test; use 8 for CI parity max_tokens: int = 512 # MC answers are short temperature: float = 0.7 top_p: float = 0.95 top_k: int = 20 max_model_len: int = 4096 dtype: str = "bfloat16" # --------------------------------------------------------------------------- # Answer extraction helper # --------------------------------------------------------------------------- def _extract_boxed_letter(text: str) -> str | None: """Return the content of the last \\boxed{} if it is a single letter, else None.""" idx = text.rfind("\\boxed") if idx < 0: return None i, depth, right = idx, 0, None while i < len(text): if text[i] == "{": depth += 1 elif text[i] == "}": depth -= 1 if depth == 0: right = i break i += 1 if right is None: return None inner = text[idx + len("\\boxed{"):right].strip() if len(inner) == 1 and inner.upper() in "ABCDEFGHIJKLMNOPQRST": return inner.upper() return None # --------------------------------------------------------------------------- # Scorer (pass@1 for MC) # --------------------------------------------------------------------------- def _compute_pass1(problems: list[dict], completions_per_problem: list[list[str]]) -> dict: """Compute pass@1 for multiple-choice: is the first completion correct? For multilingual the CI uses pass@1 (single attempt), so we average whether completion[0] contains the right letter. """ correct = 0 details = [] for prob, comps in zip(problems, completions_per_problem): gold = prob["answer"].strip().upper() pred = _extract_boxed_letter(comps[0]) if comps else None is_correct = (pred == gold) correct += int(is_correct) details.append({ "prompt": prob["prompt"][:120] + "…", "gold": gold, "predicted": pred, "correct": is_correct, }) n = len(problems) return { "metrics": {"pass@1": correct / n if n else 0.0}, "n_problems": n, "n_completions": 1, "benchmark_method": "choice", "details": details, } # --------------------------------------------------------------------------- # Validator # --------------------------------------------------------------------------- class LocalValidator: """Run vLLM inference on multilingual.jsonl and score pass@1. The system prompt is applied exactly as it was during SFT training (see trainingHelena/data.py SYSTEM_PROMPT). """ SYSTEM_PROMPT = ( "You are a multilingual assistant. " "Read the question carefully, reason step by step, and place the letter " "of your final answer inside \\boxed{}. " "For example, if the answer is option B, your response must end with \\boxed{B}." ) def __init__(self, cfg: ValidateConfig): self.cfg = cfg def run(self) -> dict: problems = self._load_problems() completions_per_problem = self._generate(problems) results = self._score(problems, completions_per_problem) self._report(results) if self.cfg.output: self._save(results) return results # ------------------------------------------------------------------ def _load_problems(self) -> list[dict]: path = Path(self.cfg.samples_file) if not path.exists(): raise FileNotFoundError( f"Validation samples not found: {path}\n" "Run from the repo root so the relative path resolves correctly." ) problems = [ json.loads(line) for line in path.read_text().splitlines() if line.strip() ] print(f"Loaded {len(problems)} validation problems from {path}") return problems def _generate(self, problems: list[dict]) -> list[list[str]]: try: from vllm import LLM, SamplingParams except ImportError: raise ImportError( "vLLM is not installed. Run validation on the GPU cluster." ) from transformers import AutoTokenizer tokenizer = AutoTokenizer.from_pretrained(self.cfg.checkpoint) print(f"Loading model from {self.cfg.checkpoint} …") llm = LLM( model=self.cfg.checkpoint, dtype=self.cfg.dtype, max_model_len=self.cfg.max_model_len, ) params = SamplingParams( temperature=self.cfg.temperature, top_p=self.cfg.top_p, top_k=self.cfg.top_k, max_tokens=self.cfg.max_tokens, n=self.cfg.n_completions, ) # Apply the same chat template the CI will use. # The system prompt is embedded so the model knows to output \\boxed{letter}. formatted_prompts = [ tokenizer.apply_chat_template( [ {"role": "system", "content": self.SYSTEM_PROMPT}, {"role": "user", "content": p["prompt"]}, ], tokenize=False, add_generation_prompt=True, ) for p in problems ] print(f"Generating {self.cfg.n_completions} completion(s) per problem …") outputs = llm.generate(formatted_prompts, params) completions_per_problem = [ [c.text for c in out.outputs] for out in outputs ] # Boxed-rate sanity check — every completion should have \\boxed{LETTER} all_ok = True for i, (comps, prob) in enumerate(zip(completions_per_problem, problems)): has_boxed = sum( 1 for c in comps if _extract_boxed_letter(c) is not None ) ok = has_boxed == len(comps) if not ok: all_ok = False print( f" [{i:2d}] gold={prob['answer']!r:4s} " f"boxed_letter_rate={has_boxed}/{len(comps)}" f"{' ✓' if ok else ' ✗ MISSING BOXED'}" ) if not all_ok: print( "\nWARNING: some completions are missing \\boxed{{LETTER}}. " "Do NOT push until this is fixed." ) else: print("\nAll completions have \\boxed{{LETTER}} — safe to push.") return completions_per_problem def _score( self, problems: list[dict], completions_per_problem: list[list[str]], ) -> dict: return _compute_pass1(problems, completions_per_problem) def _report(self, results: dict) -> None: m = results["metrics"] print( f"\n{'='*55}\n" f" Multilingual validation results\n" f" pass@1 = {m['pass@1']:.4f}\n" f" (n_problems={results['n_problems']}, " f"method={results['benchmark_method']})\n" f"{'='*55}" ) # Print per-problem breakdown for d in results["details"]: tick = "✓" if d["correct"] else "✗" print(f" {tick} gold={d['gold']} pred={d['predicted']} | {d['prompt']}") def _save(self, results: dict) -> None: out_path = Path(self.cfg.output) out_path.parent.mkdir(parents=True, exist_ok=True) out_path.write_text(json.dumps(results, ensure_ascii=False, indent=2)) print(f"\nDetailed results written to {out_path}") # --------------------------------------------------------------------------- # CLI # --------------------------------------------------------------------------- if __name__ == "__main__": parser = argparse.ArgumentParser( description="Local multilingual validation (mirrors CI)" ) parser.add_argument("--checkpoint", required=True, help="Path to the model checkpoint to evaluate") parser.add_argument("--samples", default="validation_samples/multilingual.jsonl", dest="samples_file", help="Path to the validation JSONL file") parser.add_argument("--output", default=None, help="Optional path to write detailed JSON results") parser.add_argument("--n", type=int, default=4, dest="n_completions", help="Number of completions per problem (1 is enough for pass@1)") parser.add_argument("--max-tokens", type=int, default=512) parser.add_argument("--temperature", type=float, default=0.7) parser.add_argument("--top-p", type=float, default=0.95) parser.add_argument("--top-k", type=int, default=20) args = parser.parse_args() cfg = ValidateConfig( checkpoint=args.checkpoint, samples_file=args.samples_file, output=args.output, n_completions=args.n_completions, max_tokens=args.max_tokens, temperature=args.temperature, top_p=args.top_p, top_k=args.top_k, ) validator = LocalValidator(cfg) validator.run()