Upload folder using huggingface_hub

src/mathcompose/eval/prmbench.py

@@ -83,3 +83,73 @@ def evaluate_prmbench(
             for c, d in sorted(per_cat.items())
         },
     }
+
+
+def evaluate_batched(
+    runner,
+    dataset_id: str = "hitsmy/PRMBench_Preview",
+    split: str = "train",
+    limit: Optional[int] = None,
+    n: int = 1,
+    batch_size: int = 16,
+    temperature: float = 0.0,
+    max_new_tokens: int = 640,
+    error_steps_base: int = 1,
+    include_original: bool = True,
+) -> dict:
+    """Fast path: batched generation via an HFRunner."""
+    from collections import defaultdict as _dd
+
+    from datasets import load_dataset
+
+    from .runners import verify_batch_predictions
+
+    ds = load_dataset(dataset_id, split=split)
+    if limit is not None:
+        ds = ds.select(range(min(limit, len(ds))))
+
+    mod_items, mod_gold, mod_cat = [], [], []
+    orig_items = []
+    for row in ds:
+        steps_mod = list(row.get("modified_process") or [])
+        gold0 = {int(e) - error_steps_base for e in (row.get("error_steps") or [])}
+        if steps_mod and gold0:
+            mod_items.append((row.get("modified_question") or row.get("question") or "", steps_mod))
+            mod_gold.append(gold0)
+            mod_cat.append(row.get("classification") or "unknown")
+        if include_original:
+            steps_orig = list(row.get("original_process") or [])
+            if steps_orig:
+                orig_items.append((row.get("original_question") or row.get("question") or "", steps_orig))
+
+    mod_preds = verify_batch_predictions(runner, mod_items, n=n, batch_size=batch_size,
+                                         temperature=temperature, max_new_tokens=max_new_tokens,
+                                         desc="PRMBench/modified")
+    per_cat: dict[str, dict] = _dd(lambda: {"n": 0, "detected": 0})
+    detected = localized = 0
+    for pred, gold0, cat in zip(mod_preds, mod_gold, mod_cat):
+        is_det = pred != -1
+        detected += 1 if is_det else 0
+        localized += 1 if pred in gold0 else 0
+        per_cat[cat]["n"] += 1
+        per_cat[cat]["detected"] += 1 if is_det else 0
+
+    false_pos = None
+    if include_original and orig_items:
+        orig_preds = verify_batch_predictions(runner, orig_items, n=n, batch_size=batch_size,
+                                              temperature=temperature, max_new_tokens=max_new_tokens,
+                                              desc="PRMBench/original")
+        false_pos = sum(1 for p in orig_preds if p != -1) / len(orig_preds)
+
+    n_mod = len(mod_items)
+    return {
+        "dataset": dataset_id,
+        "n_modified": n_mod,
+        "detection": detected / n_mod if n_mod else 0.0,
+        "localization": localized / n_mod if n_mod else 0.0,
+        "false_positive_rate": false_pos,
+        "per_category_detection": {
+            c: {"n": d["n"], "detection": d["detected"] / d["n"] if d["n"] else 0.0}
+            for c, d in sorted(per_cat.items())
+        },
+    }

src/mathcompose/eval/processbench.py

@@ -114,3 +114,39 @@ def evaluate(
         "average_f1": avg_f1,
         "gpt4o_reference_f1": 61.9,
     }
+
+
+def evaluate_batched(
+    runner,
+    splits: Sequence[str] = DEFAULT_SPLITS,
+    limit: Optional[int] = None,
+    n: int = 1,
+    batch_size: int = 16,
+    temperature: float = 0.0,
+    max_new_tokens: int = 640,
+    dataset_id: str = "Qwen/ProcessBench",
+) -> dict:
+    """Fast path: batched generation via an HFRunner (vs the per-example verify_fn)."""
+    from datasets import load_dataset
+
+    from .runners import verify_batch_predictions
+
+    per_subset: list[SubsetScore] = []
+    for split in splits:
+        ds = load_dataset(dataset_id, split=split)
+        if limit is not None:
+            ds = ds.select(range(min(limit, len(ds))))
+        items = [(row["problem"], list(row["steps"])) for row in ds]
+        golds = [int(row["label"]) for row in ds]
+        preds = verify_batch_predictions(
+            runner, items, n=n, batch_size=batch_size, temperature=temperature,
+            max_new_tokens=max_new_tokens, desc=f"ProcessBench/{split}",
+        )
+        per_subset.append(score_subset(split, golds, preds))
+
+    avg_f1 = sum(s.f1 for s in per_subset) / len(per_subset) if per_subset else 0.0
+    return {
+        "subsets": {s.subset: s.as_dict() for s in per_subset},
+        "average_f1": avg_f1,
+        "gpt4o_reference_f1": 61.9,
+    }

src/mathcompose/eval/run.py

@@ -34,20 +34,22 @@ def _runner(base_id, adapter):
 
 
 def cmd_processbench(a):
-    from .processbench import evaluate
-    from .runners import make_verify_fn
+    from .processbench import evaluate_batched
     r = _runner(a.base_id, a.adapter)
-    verify_fn = make_verify_fn(r, n=a.maj_k, temperature=(0.6 if a.maj_k > 1 else 0.0))
-    res = evaluate(verify_fn, splits=tuple(a.splits), limit=a.limit)
+    bs = max(1, a.batch_size // max(1, a.maj_k))  # cap concurrent seqs = prompts * k
+    res = evaluate_batched(r, splits=tuple(a.splits), limit=a.limit, n=a.maj_k,
+                           batch_size=bs, temperature=(0.6 if a.maj_k > 1 else 0.0),
+                           max_new_tokens=a.max_new_tokens)
     _save("processbench", a.tag, res)
 
 
 def cmd_prmbench(a):
-    from .prmbench import evaluate_prmbench
-    from .runners import make_verify_fn
+    from .prmbench import evaluate_batched
     r = _runner(a.base_id, a.adapter)
-    verify_fn = make_verify_fn(r, n=a.maj_k, temperature=(0.6 if a.maj_k > 1 else 0.0))
-    res = evaluate_prmbench(verify_fn, limit=a.limit)
+    bs = max(1, a.batch_size // max(1, a.maj_k))
+    res = evaluate_batched(r, limit=a.limit, n=a.maj_k, batch_size=bs,
+                           temperature=(0.6 if a.maj_k > 1 else 0.0),
+                           max_new_tokens=a.max_new_tokens)
     _save("prmbench", a.tag, res)
 
 
@@ -105,6 +107,8 @@ def main() -> None:
     p.add_argument("--tag", default="tuned")
     p.add_argument("--maj-k", type=int, default=1)
     p.add_argument("--limit", type=int, default=None)
+    p.add_argument("--batch-size", type=int, default=16, help="max concurrent sequences (÷ maj-k)")
+    p.add_argument("--max-new-tokens", type=int, default=640)
     p.add_argument("--splits", nargs="+", default=["gsm8k", "math", "olympiadbench", "omnimath"])
     p.set_defaults(func=cmd_processbench)
 
@@ -113,6 +117,8 @@ def main() -> None:
     p.add_argument("--tag", default="tuned")
     p.add_argument("--maj-k", type=int, default=1)
     p.add_argument("--limit", type=int, default=None)
+    p.add_argument("--batch-size", type=int, default=16, help="max concurrent sequences (÷ maj-k)")
+    p.add_argument("--max-new-tokens", type=int, default=640)
     p.set_defaults(func=cmd_prmbench)
 
     p = sub.add_parser("math")

src/mathcompose/eval/runners.py

@@ -79,7 +79,7 @@ class HFRunner:
         input_ids = input_ids.to(self.device)
         attn = attn.to(self.device) if attn is not None else None
 
-        do_sample = bool(temperature and temperature > 0)
+        do_sample = bool(temperature and temperature > 0) or n > 1  # n>1 needs sampling
         gen_kwargs = dict(
             max_new_tokens=max_new_tokens,
             do_sample=do_sample,
@@ -88,7 +88,8 @@ class HFRunner:
             repetition_penalty=repetition_penalty,
         )
         if do_sample:
-            gen_kwargs.update(temperature=temperature, top_p=0.95)
+            gen_kwargs.update(temperature=(temperature if temperature and temperature > 0 else 0.7),
+                              top_p=0.95)
 
         with torch.no_grad():
             out = self.model.generate(input_ids, attention_mask=attn, **gen_kwargs)
@@ -96,6 +97,52 @@ class HFRunner:
         return self.tokenizer.batch_decode(gen, skip_special_tokens=True)
 
 
+    def chat_many(
+        self,
+        messages_list: list[list[dict]],
+        n: int = 1,
+        temperature: float = 0.0,
+        max_new_tokens: int = 640,
+        repetition_penalty: float = 1.1,
+    ) -> list[list[str]]:
+        """Batched generation across many prompts (left-padded). Returns, per input
+        prompt, a list of n decoded completions. Far faster than looping chat()."""
+        import torch
+
+        tok = self.tokenizer
+        old_side = tok.padding_side
+        tok.padding_side = "left"  # required for correct batched generation
+        try:
+            rendered = [
+                tok.apply_chat_template(m, add_generation_prompt=True, tokenize=False)
+                for m in messages_list
+            ]
+            enc = tok(
+                rendered, return_tensors="pt", padding=True, truncation=True,
+                max_length=max(1, self.max_context - max_new_tokens),
+                add_special_tokens=False,
+            )
+        finally:
+            tok.padding_side = old_side
+
+        input_ids = enc["input_ids"].to(self.device)
+        attn = enc["attention_mask"].to(self.device)
+        do_sample = bool(temperature and temperature > 0) or n > 1  # n>1 needs sampling
+        gen_kwargs = dict(
+            max_new_tokens=max_new_tokens, do_sample=do_sample, num_return_sequences=n,
+            pad_token_id=tok.pad_token_id, repetition_penalty=repetition_penalty,
+        )
+        if do_sample:
+            gen_kwargs.update(temperature=(temperature if temperature and temperature > 0 else 0.7),
+                              top_p=0.95)
+        with torch.no_grad():
+            out = self.model.generate(input_ids, attention_mask=attn, **gen_kwargs)
+        gen = out[:, input_ids.shape[1]:]
+        texts = tok.batch_decode(gen, skip_special_tokens=True)
+        # generate returns batch*n rows, grouped by input then sample
+        return [texts[i * n:(i + 1) * n] for i in range(len(messages_list))]
+
+
 def make_verify_fn(runner: HFRunner, n: int = 1, temperature: float = 0.0,
                    max_new_tokens: int = 1024) -> Callable[[str, list[str]], list[str]]:
     def verify_fn(problem: str, steps: list[str]) -> list[str]:
@@ -110,3 +157,27 @@ def make_solve_fn(runner: HFRunner, n: int = 1, temperature: float = 0.0,
         return runner.chat(generator_messages(problem), n=n,
                            temperature=temperature, max_new_tokens=max_new_tokens)
     return solve_fn
+
+
+def verify_batch_predictions(runner: HFRunner, items, n: int = 1, batch_size: int = 16,
+                             temperature: float = 0.0, max_new_tokens: int = 640,
+                             desc: str = "verify") -> list[int]:
+    """items: list of (problem, steps). Returns the Maj@n first-error index per item,
+    using batched generation. This is the fast path used by ProcessBench/PRMBench."""
+    from .parse import majority_index
+
+    try:
+        from tqdm.auto import tqdm
+    except Exception:  # pragma: no cover
+        def tqdm(x, **k):
+            return x
+
+    preds: list[int] = []
+    chunks = list(range(0, len(items), batch_size))
+    for i in tqdm(chunks, desc=desc):
+        chunk = items[i:i + batch_size]
+        msgs = [verifier_messages(p, s) for p, s in chunk]
+        samples = runner.chat_many(msgs, n=n, temperature=temperature,
+                                   max_new_tokens=max_new_tokens)
+        preds.extend(majority_index(s) for s in samples)
+    return preds

src/mathcompose/train/colab_train.ipynb

@@ -4,10 +4,10 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "# mathcompose \u2014 QLoRA training on Colab (free T4)\n",
+    "# mathcompose \u2014 QLoRA training on Colab\n",
     "\n",
-    "Pulls the code **and** the Model-V dataset from HF Hub with just your token \u2014 no\n",
-    "GitHub needed. **Runtime \u2192 Change runtime type \u2192 T4 GPU** before running."
+    "Pulls code+data from HF with just your token. **Runtime \u2192 GPU** before running.\n",
+    "Training auto-tunes to the detected GPU (`--gpu-preset auto`)."
    ]
   },
   {
@@ -16,7 +16,7 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "# 1. GPU deps (versions this project was built against) + 4-bit kernels\n",
+    "# 1. GPU deps (+ 4-bit kernels)\n",
     "!pip -q install \"transformers>=5.13,<6\" \"trl>=1.7,<2\" \"peft>=0.19,<0.20\" \\\n",
     "    \"accelerate>=1.14\" \"datasets>=5.0\" \"huggingface_hub>=1.22\" sympy bitsandbytes"
    ]
@@ -27,16 +27,16 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "# 2. Token -> pull the code+data snapshot -> install\n",
+    "# 2. Token -> pull code+data snapshot -> install\n",
     "import os\n",
     "from huggingface_hub import HfApi, snapshot_download, login\n",
-    "os.environ[\"HF_TOKEN\"] = \"hf_xxx\"   # paste, or Colab Secrets: userdata.get('HF_TOKEN')\n",
+    "os.environ[\"HF_TOKEN\"] = \"hf_xxx\"   # paste, or Colab Secrets\n",
     "login(os.environ[\"HF_TOKEN\"])\n",
     "USER = HfApi().whoami()[\"name\"]\n",
     "proj = snapshot_download(f\"{USER}/mathcompose\", repo_type=\"model\")\n",
     "%cd {proj}\n",
     "!pip -q install -e .\n",
-    "print(\"project at\", proj, \"| data:\", os.listdir(\"data/verifier\"))"
+    "print(\"data:\", os.listdir(\"data/verifier\"))"
    ]
   },
   {
@@ -53,7 +53,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 4. Train Model V (verifier) \u2014 QLoRA on the balanced dataset"
+    "## 4. Train Model V (auto-tunes batch/seq/precision to the GPU)"
    ]
   },
   {
@@ -69,8 +69,9 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 5. Eval V on ProcessBench \u2014 base vs tuned (Maj@8) + PRMBench\n",
-    "Use `--limit` for a fast read; drop it for the full 3,400-example score."
+    "## 5. Eval V on ProcessBench \u2014 base vs tuned (now batched \u2192 fast)\n",
+    "Fast first read below (1 split, small limit, Maj@4). For the headline number,\n",
+    "drop `--splits`/raise `--limit` and use `--maj-k 8`."
    ]
   },
   {
@@ -79,9 +80,13 @@
    "execution_count": null,
    "outputs": [],
    "source": [
-    "!python -m mathcompose.eval.run processbench --tag base  --maj-k 1 --limit 200\n",
-    "!python -m mathcompose.eval.run processbench --adapter runs/verifier_v --tag tuned --maj-k 8 --limit 200\n",
-    "!python -m mathcompose.eval.run prmbench    --adapter runs/verifier_v --tag tuned --limit 300\n",
+    "# fast sanity read (~a few min on A100):\n",
+    "!python -m mathcompose.eval.run processbench --tag base  --maj-k 1 --limit 100 --splits gsm8k\n",
+    "!python -m mathcompose.eval.run processbench --adapter runs/verifier_v --tag tuned --maj-k 4 --limit 100 --splits gsm8k\n",
+    "# full run (all 4 splits, Maj@8) \u2014 uncomment when the fast read looks right:\n",
+    "# !python -m mathcompose.eval.run processbench --tag base  --maj-k 1\n",
+    "# !python -m mathcompose.eval.run processbench --adapter runs/verifier_v --tag tuned --maj-k 8\n",
+    "!python -m mathcompose.eval.run prmbench --adapter runs/verifier_v --tag tuned --maj-k 1 --limit 200\n",
     "!python -m mathcompose.eval.run report"
    ]
   },
@@ -109,9 +114,8 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## 7. (Later) Model G + the composition\n",
-    "Build G's dataset (teacher=claude-opus-4-8), train, then run V-reranks-G \u2014 the\n",
-    "headline generator-verifier-gap number."
+    "## 7. (Later) Model G + composition\n",
+    "Build G's dataset (teacher=claude-opus-4-8), train, then V-reranks-G."
    ]
   },
   {

tests/test_eval_plumbing.py

@@ -31,3 +31,17 @@ def test_eval_pipeline_runs():
     s = score_subset("smoke", golds, preds)   # tiny model -> garbage preds, but scoring runs
     assert 0.0 <= s.f1 <= 1.0
     assert s.n_error + s.n_correct == 3
+
+
+@pytest.mark.slow
+@pytest.mark.network
+def test_batched_processbench_eval_runs():
+    """The fast batched path: load -> chat_many -> parse -> score over a tiny slice."""
+    pytest.importorskip("torch")
+    from mathcompose.eval.runners import HFRunner
+    from mathcompose.eval.processbench import evaluate_batched
+
+    runner = HFRunner(TINY, max_context=512)
+    res = evaluate_batched(runner, splits=["gsm8k"], limit=4, n=1, batch_size=2, max_new_tokens=6)
+    assert 0.0 <= res["average_f1"] <= 1.0
+    assert "gsm8k" in res["subsets"]