File size: 10,541 Bytes
f908c84
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
"""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()