Upload 6 files

f908c84 verified 6 months ago

10.5 kB

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

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