validate.py

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