From 296718c4f7815cf773401b7e862f6c335f3a5c3e Mon Sep 17 00:00:00 2001 From: karthikreddy-02 Date: Sun, 9 Aug 2026 11:42:32 -0400 Subject: [PATCH 1/2] =?UTF-8?q?judge:=20trained=20ETP=20relation=20judge?= =?UTF-8?q?=20=E2=80=94=20library,=20CLI,=20training=20and=20eval?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Replaces the prompted-LLM judge with a trained artifact: known training data, a known split, and weights we own and can probe. Labels come from oracle/'s Lean-verified implication matrix, so the training signal has zero label noise and no prover runs at training time. The split is by equivalence class, never by pair (1,064/140/211 of 1,415). Laws inside a class are mutually derivable, so a random pair split lets a lookup table score perfectly. The x = y class alone is 31.9% of the catalogue and is pinned into training. pairs_val minus classes_val is therefore the memorization measurement: 1.91 pp on Gemma-2-2B, 1.25 pp on Llama-3.1-8B. Both trained a judge, not a lookup table. judge.py exposes both a CLI (compare / batch / rank) and a Judge class, with the live PeftModel reachable as .model so the interpretability work can attach hooks rather than shell out. Scoring restricts logits to the four label tokens before the argmax, so an invalid answer is impossible by construction rather than merely unlikely. Results and ablations across two model families are in judge/README.md, including two findings that did NOT survive replication and should not be claimed: the attention/MLP asymmetry is a Gemma artifact, and the early-vs-middle layer ordering flips between families. What does replicate: rank saturates at or below 16, and late-layer-only adaptation collapses (5-6x worse than any other layer choice) in both models. Neither base model can be contaminated: ETP launched 2024-09-25, while Gemma-2 shipped 2024-06-27 and Llama-3.1 on 2024-07-23. Generated data and trained weights are gitignored; both are regenerable. --- .gitignore | 5 + judge/README.md | 275 +++++++++++++++++++++++++++++ judge/build_pair_dataset.py | 250 ++++++++++++++++++++++++++ judge/eval_judge.py | 320 ++++++++++++++++++++++++++++++++++ judge/infra/INFRA.md | 84 +++++++++ judge/infra/launch_thunder.sh | 82 +++++++++ judge/judge.py | 243 ++++++++++++++++++++++++++ judge/sweep.sh | 116 ++++++++++++ judge/train_judge.py | 245 ++++++++++++++++++++++++++ 9 files changed, 1620 insertions(+) create mode 100644 judge/README.md create mode 100644 judge/build_pair_dataset.py create mode 100644 judge/eval_judge.py create mode 100644 judge/infra/INFRA.md create mode 100755 judge/infra/launch_thunder.sh create mode 100755 judge/judge.py create mode 100755 judge/sweep.sh create mode 100644 judge/train_judge.py diff --git a/.gitignore b/.gitignore index 7a60b85..5e88b90 100644 --- a/.gitignore +++ b/.gitignore @@ -1,2 +1,7 @@ __pycache__/ *.pyc + +# judge/: generated dataset and trained weights (regenerable / too large) +judge/data/ +judge/runs/ +*.tar.gz diff --git a/judge/README.md b/judge/README.md new file mode 100644 index 0000000..d6c8dfb --- /dev/null +++ b/judge/README.md @@ -0,0 +1,275 @@ +# `judge` — a trained relation judge for the ETP fragment + +Classifies the logical relation between two equational laws: + +``` +equivalent | weaker | stronger | incomparable +``` + +The label describes **B relative to A**, matching `oracle/`'s convention +(`weaker` = B drops constraints). + +This replaces a stack of prompted LLMs with a single measurable artifact: known +training data, a known split, and weights we own and can probe. It is trained on +the ETP implication matrix, where every label is Lean-verified rather than judged. + +**No GPU required to use it.** One comparison is a couple of seconds on CPU; +CUDA and Apple-Silicon MPS are used automatically when available. + +--- + +## Quickstart + +```bash +pip install torch transformers peft +huggingface-cli login # Gemma is gated: accept the license first + +export JUDGE_ADAPTER=/path/to/judge-v0/adapter-final +python3 judge.py compare "x * y = y * x" "a * b = b * a" +``` + +``` +A: x * y = y * x +B: a * b = b * a + + equivalent (confidence 1.000) + + equivalent 1.000 ######################################## + weaker 0.000 + stronger 0.000 + incomparable 0.000 +``` + +First run downloads the ~5 GB base model. The adapter itself is ~113 MB. + +### Commands + +| command | purpose | +|---|---| +| `compare A B` | score one pair; `--json` for machine output | +| `batch pairs.jsonl` | score a JSONL of `{text_a, text_b}` rows; reports accuracy if `label` present | +| `rank intended candidates.txt` | order candidates by `P(equivalent)` — what best-of-n selection actually does | + +`rank` is the one to reach for when using the judge as a *checker*: + +``` +$ python3 judge.py rank "x * y = y * x" candidates.txt +intended: x * y = y * x + + 1. p(equiv)=0.9997 [equivalent ] p * q = q * p + 2. p(equiv)=0.9996 [equivalent ] a * b = b * a + 3. p(equiv)=0.0383 [weaker ] x = x + 4. p(equiv)=0.0017 [incomparable] x * (y * z) = (x * y) * z +``` + +### As a library + +Downstream interpretability tooling should import rather than shell out: + +```python +from judge import Judge + +j = Judge("path/to/adapter-final") # base model read from adapter_config.json +v = j.compare("x * y = y * x", "a * b = b * a") +v.label, v.confidence, v.probs + +j.probs([(a, b), ...]) # (n, 4) array, no argmax applied +j.rank(intended, candidates) # [(candidate, p_equivalent, pred), ...] +j.model # the live PeftModel, for hooks/probing +``` + +`Judge.model` is exposed deliberately: activation capture, SAE work, and +ablations attach there. + +--- + +## How it works + +Prompt is `"A: {eq}\nB: {eq}\nRelation:"`. The label is a **single next token** +with loss masked to that token alone. At inference the logits are **restricted to +the four label tokens** before the argmax, so the judge cannot emit an invalid +answer — a wrong prediction is always one of the other three, never gibberish. + +`unknown` is not a class. The 190 unresolved matrix cells are never sampled, so +the judge cannot abstain. If you need abstention, threshold on `confidence` — +but note calibration is mediocre (see Caveats). + +--- + +## Training data + +Built by `build_pair_dataset.py` from `oracle/`'s matrix. Labels are two byte +reads from the Lean-verified implication matrix — **no prover runs at training +time**; the certification is a frozen artifact of the +`2024-11-10-outcomes` snapshot. + +```bash +cd judge +python3 build_pair_dataset.py --out data/ +``` + +Requires `oracle/data/matrix.bin` — see `oracle/README.md` to build it. +The generated `data/` is gitignored (39 MB, fully regenerable). + +### The split is the important part + +**By equivalence class, never by pair.** 1,415 classes → **1,064 train / 140 val +/ 211 test**. Laws inside a class are mutually derivable, so a random pair split +puts the same law on both sides and a lookup table scores perfectly. + +The `x = y` class alone is **1,496 laws — 31.9% of the catalogue** — and is +*pinned into training*. A random split could otherwise have removed a third of +all laws from training. + +Four splits, and the contrast between two of them is the measurement: + +| split | composition | measures | +|---|---|---| +| `train` | 240,000 pairs from train classes | — | +| `pairs_val` | new pairings of **familiar** laws | interpolation | +| `classes_val` | ≥1 law from held-out classes | novel theories | +| `test` | ≥1 law from test classes, never selected on | final | + +**`pairs_val − classes_val` is the memorization measurement.** Everything else is +held constant; only law familiarity varies. A pair-level memoriser fails +`pairs_val`; a law-level lookup passes it and fails `classes_val`. Only a +relational model passes both. + +--- + +## Results + +Trained models, `test_acc` on 211 held-out equivalence classes (chance = 0.25): + +| config | Gemma-2-2B | Llama-3.1-8B | +|---|---|---| +| **baseline** (r16, all layers, all modules) | **0.9619** | **0.9649** | +| rank 4 | 0.9566 | 0.9596 | +| rank 64 | 0.9625 | 0.9647 | +| early layers only | 0.9414 | 0.9619 | +| middle layers only | 0.9494 | 0.9457 | +| **late layers only** | **0.8625** | **0.8491** | +| attention only | 0.9523 | 0.9629 | +| MLP only | 0.9601 | 0.9631 | + +Memorization gap: **1.91 pp** (Gemma), **1.25 pp** (Llama). Both trained a judge, +not a lookup table. + +**Neither base model can be contaminated.** ETP launched 2024-09-25; Gemma-2 was +released 2024-06-27 and Llama-3.1 on 2024-07-23. Both predate the project's +existence, which is a stronger guarantee than any stated data cutoff. + +### What replicates across both families + +- **Rank saturates at or below 16.** Quadrupling the adapter moves accuracy + +0.06 pp / −0.01 pp. Cutting to r=4 costs exactly −0.53 pp on both. +- **Late-only collapses** — 5–6× worse than any other layer choice. The relation + is built in the front two-thirds, not decided at the end. Note `late` has + *more* adapted layers than `early` (10 vs 8 on Gemma), so this is positional, + not a capacity effect. +- **MLP-only is nearly free** — −0.18 pp in both. + +### What does NOT replicate — do not claim + +- **The attention/MLP asymmetry is a Gemma artifact.** Gemma: attention-only cost + 5× more than MLP-only. Llama: they are equivalent. "The computation lives in + the MLP" does not survive replication. +- **Early-vs-middle ordering flips** between families. + +--- + +## Retraining + +Any CUDA box with **≥ 40 GB** works. Memory is dominated by the vocabulary-sized +logits plus retained activations for every adapted layer. + +```bash +python3 train_judge.py --data data --out runs/judge-v0 \ + --micro-batch 64 --grad-accum 2 --eval-rows 2000 --keep-checkpoints 0 +``` + +Measured, so you don't rediscover it: + +| model | micro-batch | peak VRAM | outcome | +|---|---|---|---| +| Gemma-2-2B | 128 | > 47 GB | **OOM** on a 48 GB A6000 | +| Gemma-2-2B | 64 | 28.6 GB | fine, ~50 min | +| Llama-3.1-8B | 64 | 43.8 GB | fine, ~50 min on an A100 80 GB | + +Gemma-2's 256k vocabulary is what makes it expensive — the logits tensor alone is +~13.7 GB at micro-batch 128 (bf16 + fp32 upcast + gradient). Llama's 128k vocab +halves that. + +Keep `--grad-accum` × `--micro-batch` = 128 so runs stay comparable: gradient +accumulation is mathematically equivalent at a fixed effective batch, since each +example contributes exactly one unmasked loss token. + +`--keep-checkpoints 0` retains every checkpoint (10 per run at the default +`--eval-every 200`), which the interpretability work needs for trajectories. + +### Ablation sweep + +`sweep.sh` reproduces the table above, one variable per run, seed fixed: + +```bash +MODEL=meta-llama/Llama-3.1-8B OUT=runs/sweep-llama MB=64 GA=2 \ + SKIP_RUNS=baseline bash sweep.sh run +bash sweep.sh summary +``` + +Finished runs are skipped, so it is safe to re-invoke after a crash. + +--- + +## Evaluation beyond accuracy + +`eval_judge.py` reports what training-time accuracy cannot — confusion matrix, +precision (training reports recall only), the false-equivalence rate, the +`equivalent` tier split, symmetry, and best-of-n ranking: + +```bash +python3 eval_judge.py --adapter runs/judge-v0/adapter-final --data data +``` + +Findings for judge-v0 (Gemma): + +- **False-equivalence rate 0.0077** — the dangerous error, where a drifted + candidate is accepted. Missed-equivalence runs at 0.044, so the judge errs + toward caution by **5.7 : 1**. That is the right asymmetry for a checker. +- **`equivalent` is 50% trivial.** Same-node rows (the identical law re-rendered) + score **1.000**; cross-node rows score **0.912**. The reported 0.956 is their + average — **quote 0.912** for genuine equivalence detection. +- **Dominant failure is over-predicting `incomparable`**: recall 0.977 but + precision 0.899, absorbing 69% of all errors. +- **Symmetry 0.978** — `label(A,B)=weaker` implies `label(B,A)=stronger`. + +--- + +## Caveats + +- **Single seed.** No variance estimate on any number here. +- **Calibration is mediocre** — 39.7% of errors are held above 0.9 confidence, so + a confidence threshold will not cleanly filter dangerous errors. +- **The best-of-n exploitation rate is not yet valid.** The current test uses + same-node faithful candidates, which the judge never misses (1.000). It must be + rebuilt with cross-node candidates (0.912) before being quoted. +- **`classes_val` is only 27% both-novel**; 73% pair one novel law with a + familiar one, so the memorization gap is a *lower bound*. +- **Out-of-fragment behaviour is untrained and unknown.** Training contains only + the 4,694 catalogued laws. Order > 4 identities, multi-operation expressions, + and chained equalities never appear. +- **Layer results show where adaptation is *needed*, not where computation + *lives*.** A frozen late layer may still be doing essential work. + +--- + +## Files + +| file | purpose | +|---|---| +| `judge.py` | library + CLI (`compare`, `batch`, `rank`) | +| `train_judge.py` | training; LoRA rank / layer / module flags | +| `build_pair_dataset.py` | builds the splits from `oracle/`'s matrix | +| `eval_judge.py` | post-hoc evaluation | +| `sweep.sh` | ablation sweep driver | +| `infra/` | optional GPU-rental helpers — see `infra/INFRA.md` | diff --git a/judge/build_pair_dataset.py b/judge/build_pair_dataset.py new file mode 100644 index 0000000..22c923e --- /dev/null +++ b/judge/build_pair_dataset.py @@ -0,0 +1,250 @@ +#!/usr/bin/env python3 +"""Build the judge-LM pair dataset from the ETP implication matrix. + +Reads the oracle's matrix.bin + equations.txt, computes equivalence classes, +splits BY CLASS (never by pair), samples balanced relation pairs with surface +augmentation, and writes train/pairs_val/classes_val/test JSONL. + +Labels (4-way): equivalent | weaker | stronger | incomparable + For a pair (A, B) the label describes B relative to A, matching the + oracle's convention (weaker = B drops constraints). + 'unknown' is not a trainable class (190 cells in the snapshot); unresolved + pairs are simply never sampled. + +Splits: + - classes with > PIN_SIZE members are pinned to train (the x = y monster + class would otherwise remove a third of the catalogue from training) + - remaining classes: TEST_FRAC to test, VAL_FRAC to classes_val, rest train + - pairs_val: unseen PAIRS of train-class laws (interpolation measure) + - classes_val/test: pairs where at least one law is from a held-out class + (novel-theory measure). The pairs_val vs classes_val accuracy gap is the + memorization measurement. + +Usage: + ETP_EQUATIONS=... python3 build_pair_dataset.py --out data/ + (defaults assume the semantic-diffchecking-repo oracle layout) +""" +from __future__ import annotations + +import argparse +import json +import os +import random +import sys +from collections import Counter, defaultdict + +ORACLE_DIR = os.path.expanduser("~/Documents/semantic-diffchecking-repo/oracle") +sys.path.insert(0, ORACLE_DIR) +os.environ.setdefault("ETP_EQUATIONS", os.path.join(ORACLE_DIR, "data", "equations.txt")) + +from normalizer import INFIX_OPS, Equation, Op, Var, parse_equation # noqa: E402 + +MATRIX_BIN = os.path.join(ORACLE_DIR, "data", "matrix.bin") +META_JSON = os.path.join(ORACLE_DIR, "data", "matrix_meta.json") + +PROOF_FALSE, PROOF_TRUE = 0, 1 + +# Only symbols the oracle's own tokenizer accepts (NOT ascii '.', which it +# rejects) — keeps every augmented rendering round-trippable through mapper. +OP_SYMBOLS = ["*", "·", "@", "+"] +assert all(s in INFIX_OPS for s in OP_SYMBOLS) +VAR_POOLS = [ + list("xyzwuvst"), + list("abcdefgh"), + list("pqrstuvw"), + list("mnkljihg"), +] + + +def render(term, mapping, op, top=True): + if isinstance(term, Var): + return mapping[term.name] + s = f"{render(term.left, mapping, op, False)} {op} {render(term.right, mapping, op, False)}" + return s if top else f"({s})" + + +def augment(eq: Equation, rng: random.Random) -> str: + """Render a parsed law with random variable names, op symbol, orientation.""" + varnames: list = [] + eq.lhs.variables(varnames) + eq.rhs.variables(varnames) + pool = rng.choice(VAR_POOLS)[:] + rng.shuffle(pool) + mapping = {v: pool[i] for i, v in enumerate(varnames)} + op = rng.choice(OP_SYMBOLS) + lhs = render(eq.lhs, mapping, op) + rhs = render(eq.rhs, mapping, op) + if rng.random() < 0.5: + lhs, rhs = rhs, lhs + return f"{lhs} = {rhs}" + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--out", default="data") + ap.add_argument("--per-class-train", type=int, default=60_000, + help="examples per relation label in train") + ap.add_argument("--per-class-eval", type=int, default=2_000, + help="examples per relation label in each eval split") + ap.add_argument("--pin-size", type=int, default=50, + help="classes larger than this are pinned to train") + ap.add_argument("--test-frac", type=float, default=0.15) + ap.add_argument("--val-frac", type=float, default=0.10) + ap.add_argument("--seed", type=int, default=0) + args = ap.parse_args() + rng = random.Random(args.seed) + + with open(META_JSON, encoding="utf-8") as f: + n = json.load(f)["n"] + with open(MATRIX_BIN, "rb") as f: + buf = f.read() + with open(os.environ["ETP_EQUATIONS"], encoding="utf-8") as f: + laws = [ln.strip() for ln in f.read().splitlines() if ln.strip()] + assert len(laws) == n + asts = [parse_equation(law) for law in laws] + + # --- equivalence classes (union-find over mutual proof_true) ------------- + parent = list(range(n)) + + def find(x: int) -> int: + while parent[x] != x: + parent[x] = parent[parent[x]] + x = parent[x] + return x + + for i in range(n): + base = i * n + for j in range(i + 1, n): + if buf[base + j] == PROOF_TRUE and buf[j * n + i] == PROOF_TRUE: + ri, rj = find(i), find(j) + if ri != rj: + parent[rj] = ri + members = defaultdict(list) + for i in range(n): + members[find(i)].append(i) + classes = list(members.values()) + print(f"{n} laws, {len(classes)} equivalence classes") + + # --- split by class ------------------------------------------------------ + pinned = [c for c in classes if len(c) > args.pin_size] + free = [c for c in classes if len(c) <= args.pin_size] + rng.shuffle(free) + n_test = int(len(free) * args.test_frac) + n_val = int(len(free) * args.val_frac) + test_classes = free[:n_test] + val_classes = free[n_test:n_test + n_val] + train_classes = free[n_test + n_val:] + pinned + + def lawset(cs): + return sorted(i for c in cs for i in c) + + train_laws = lawset(train_classes) + val_laws = lawset(val_classes) + test_laws = lawset(test_classes) + print(f"classes train/val/test: {len(train_classes)}/{len(val_classes)}/{len(test_classes)}" + f" (pinned to train: {len(pinned)})") + print(f"laws train/val/test: {len(train_laws)}/{len(val_laws)}/{len(test_laws)}") + + cls_of = {} + for c in classes: + r = find(c[0]) + for i in c: + cls_of[i] = r + + def relation(a: int, b: int): + """Label of b relative to a, or None if unresolved/uninformative.""" + fwd, bwd = buf[a * n + b], buf[b * n + a] + if fwd == PROOF_TRUE and bwd == PROOF_TRUE: + return "equivalent" + if fwd == PROOF_TRUE and bwd == PROOF_FALSE: + return "weaker" + if fwd == PROOF_FALSE and bwd == PROOF_TRUE: + return "stronger" + if fwd == PROOF_FALSE and bwd == PROOF_FALSE: + return "incomparable" + return None # conjecture / open: never sampled + + def example(a: int, b: int, label: str, tier: str): + return { + "text_a": augment(asts[a], rng), + "text_b": augment(asts[b], rng), + "label": label, + "node_a": a + 1, + "node_b": b + 1, + "tier": tier, + } + + def sample_split(pool_a, pool_b, per_class, cross_split): + """Sample `per_class` examples per label. cross_split: at least one + side must come from pool_b (the held-out laws).""" + out = [] + eq_classes_here = [c for c in classes + if len(c) >= 2 and all(i in pool_set for i in c)] + pool_set_a, pool_set_b = set(pool_a), set(pool_b) + # equivalent: half same-node augmented (easy), half cross-node (hard) + n_easy = per_class // 2 + easy_src = pool_b if cross_split else pool_a + for _ in range(n_easy): + a = rng.choice(easy_src) + out.append(example(a, a, "equivalent", "same-node")) + made, guard = 0, 0 + while made < per_class - n_easy and guard < per_class * 200: + guard += 1 + c = rng.choice(multi_classes) + a, b = rng.sample(c, 2) + if cross_split and (a not in pool_set_b and b not in pool_set_b): + continue + if not cross_split and (a not in pool_set_a or b not in pool_set_a): + continue + out.append(example(a, b, "equivalent", "cross-node")) + made += 1 + # directional + incomparable by rejection sampling + want = {"weaker": per_class, "stronger": per_class, "incomparable": per_class} + got = Counter() + guard = 0 + while any(got[k] < want[k] for k in want) and guard < 5_000_000: + guard += 1 + if cross_split: + a = rng.choice(pool_b if rng.random() < 0.5 else pool_a) + b = rng.choice(pool_b if a in pool_set_a else (pool_a + pool_b)) + else: + a, b = rng.choice(pool_a), rng.choice(pool_a) + if a == b: + continue + lab = relation(a, b) + if lab in want and got[lab] < want[lab]: + out.append(example(a, b, lab, "graph")) + got[lab] += 1 + rng.shuffle(out) + return out + + global multi_classes, pool_set + multi_classes = [c for c in classes if len(c) >= 2] + pool_set = set(range(n)) + + os.makedirs(args.out, exist_ok=True) + splits = { + "train": sample_split(train_laws, train_laws, args.per_class_train, False), + "pairs_val": sample_split(train_laws, train_laws, args.per_class_eval, False), + "classes_val": sample_split(train_laws, val_laws, args.per_class_eval, True), + "test": sample_split(train_laws, test_laws, args.per_class_eval, True), + } + manifest = {"seed": args.seed, "n_laws": n, "n_classes": len(classes), + "train_classes": len(train_classes), "val_classes": len(val_classes), + "test_classes": len(test_classes), "pin_size": args.pin_size, + "source": "2024-11-10-outcomes snapshot via oracle/build_matrix.py"} + for name, rows in splits.items(): + path = os.path.join(args.out, f"{name}.jsonl") + with open(path, "w", encoding="utf-8") as f: + for r in rows: + f.write(json.dumps(r, ensure_ascii=False) + "\n") + dist = Counter(r["label"] for r in rows) + tiers = Counter(r["tier"] for r in rows if r["label"] == "equivalent") + print(f"{name:12s} {len(rows):7d} rows {dict(dist)} equiv tiers {dict(tiers)}") + with open(os.path.join(args.out, "manifest.json"), "w") as f: + json.dump(manifest, f, indent=2) + print("wrote", args.out) + + +if __name__ == "__main__": + main() diff --git a/judge/eval_judge.py b/judge/eval_judge.py new file mode 100644 index 0000000..948de76 --- /dev/null +++ b/judge/eval_judge.py @@ -0,0 +1,320 @@ +#!/usr/bin/env python3 +"""Post-hoc evaluation of a trained ETP relation judge. + +Answers what training-time accuracy cannot: + 1. confusion matrix + precision (training reports recall only) + 2. false-equivalence rate — the error that accepts a drifted translation + 3. `equivalent` split by tier: same-node (surface rewrite) vs cross-node (real) + 4. symmetry: label(A,B)=weaker must imply label(B,A)=stronger + 5. best-of-n ranking — does the judge rank a faithful candidate top? + 6. both-novel slice: rows where BOTH laws come from held-out classes + +Runs on a laptop against the adapter; no GPU required. + + python3 eval_judge.py --adapter runs/judge-v0/adapter-final --data data +""" +from __future__ import annotations + +import argparse +import json +import os +import random +import sys +from collections import Counter, defaultdict + +LABELS = ["equivalent", "weaker", "stronger", "incomparable"] +SWAP = {"equivalent": "equivalent", "weaker": "stronger", + "stronger": "weaker", "incomparable": "incomparable"} +ORACLE_DIR = os.path.expanduser("~/Documents/semantic-diffchecking-repo/oracle") + + +def load_rows(path): + with open(path, encoding="utf-8") as f: + return [json.loads(l) for l in f] + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--adapter", default="runs/judge-v0/adapter-final") + ap.add_argument("--base", default="google/gemma-2-2b") + ap.add_argument("--data", default="data") + ap.add_argument("--split", default="test") + ap.add_argument("--batch", type=int, default=64) + ap.add_argument("--limit", type=int, default=0, help="0 = all rows") + ap.add_argument("--bon-anchors", type=int, default=200) + ap.add_argument("--bon-n", type=int, default=32) + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--out", default=None, help="write metrics JSON here") + args = ap.parse_args() + + import torch + from peft import PeftModel + from transformers import AutoModelForCausalLM, AutoTokenizer + + dev = "mps" if torch.backends.mps.is_available() else ( + "cuda" if torch.cuda.is_available() else "cpu") + print(f"device: {dev}") + + tok = AutoTokenizer.from_pretrained(args.adapter) + label_ids = [tok(" " + l, add_special_tokens=False).input_ids[0] for l in LABELS] + assert len(set(label_ids)) == 4, "label first-tokens collide; head is not 4-way" + + try: + base = AutoModelForCausalLM.from_pretrained( + args.base, dtype=torch.bfloat16, attn_implementation="eager") + except TypeError: # transformers < 5 spells it torch_dtype + base = AutoModelForCausalLM.from_pretrained( + args.base, torch_dtype=torch.bfloat16, attn_implementation="eager") + model = PeftModel.from_pretrained(base, args.adapter).to(dev).eval() + + @torch.no_grad() + def score(pairs): + """pairs: [(text_a, text_b)] -> probs over LABELS, shape (n, 4).""" + out = [] + for i in range(0, len(pairs), args.batch): + chunk = pairs[i:i + args.batch] + enc = [tok(f"A: {a}\nB: {b}\nRelation:", add_special_tokens=True).input_ids + for a, b in chunk] + m = max(len(e) for e in enc) + pad = tok.pad_token_id if tok.pad_token_id is not None else tok.eos_token_id + ids = torch.tensor([e + [pad] * (m - len(e)) for e in enc], device=dev) + att = torch.tensor([[1] * len(e) + [0] * (m - len(e)) for e in enc], device=dev) + last = torch.tensor([len(e) - 1 for e in enc], device=dev) + lg = model(input_ids=ids, attention_mask=att).logits + lg = lg[torch.arange(lg.size(0), device=dev), last][:, label_ids] + out.append(torch.softmax(lg.float(), -1).cpu()) + print(f"\r scored {min(i + args.batch, len(pairs))}/{len(pairs)}", + end="", flush=True) + print() + return torch.cat(out) + + rows = load_rows(os.path.join(args.data, f"{args.split}.jsonl")) + if args.limit: + rows = rows[:args.limit] + train_nodes = {r["node_a"] for r in load_rows(os.path.join(args.data, "train.jsonl"))} | \ + {r["node_b"] for r in load_rows(os.path.join(args.data, "train.jsonl"))} + print(f"{len(rows)} rows from {args.split}.jsonl") + + print("scoring forward direction...") + probs = score([(r["text_a"], r["text_b"]) for r in rows]) + pred = probs.argmax(-1).tolist() + gold = [LABELS.index(r["label"]) for r in rows] + conf = probs.max(-1).values.tolist() + + R = {} + n = len(rows) + R["n"] = n + R["accuracy"] = sum(p == g for p, g in zip(pred, gold)) / n + + # ---- confusion matrix, precision, recall ------------------------------- + cm = [[0] * 4 for _ in range(4)] + for p, g in zip(pred, gold): + cm[g][p] += 1 + R["confusion"] = {LABELS[g]: {LABELS[p]: cm[g][p] for p in range(4)} for g in range(4)} + R["per_label"] = {} + for i, lab in enumerate(LABELS): + tp = cm[i][i] + fn = sum(cm[i]) - tp + fp = sum(cm[g][i] for g in range(4)) - tp + R["per_label"][lab] = { + "recall": tp / (tp + fn) if tp + fn else None, + "precision": tp / (tp + fp) if tp + fp else None, + "support": sum(cm[i]), + } + + # ---- the dangerous error ------------------------------------------------ + non_eq = [i for i in range(n) if gold[i] != 0] + false_eq = [i for i in non_eq if pred[i] == 0] + R["false_equivalence"] = { + "rate": len(false_eq) / len(non_eq), + "count": len(false_eq), + "of_non_equivalent": len(non_eq), + "by_true_label": dict(Counter(LABELS[gold[i]] for i in false_eq)), + "mean_confidence": sum(conf[i] for i in false_eq) / len(false_eq) if false_eq else None, + } + miss_eq = [i for i in range(n) if gold[i] == 0 and pred[i] != 0] + R["missed_equivalence"] = { + "rate": len(miss_eq) / max(1, sum(1 for g in gold if g == 0)), + "into": dict(Counter(LABELS[pred[i]] for i in miss_eq)), + } + + # ---- tier split: is `equivalent` inflated by same-node rows? ------------ + R["by_tier"] = {} + for tier in ["same-node", "cross-node", "graph"]: + idx = [i for i in range(n) if rows[i]["tier"] == tier] + if idx: + R["by_tier"][tier] = { + "n": len(idx), + "accuracy": sum(pred[i] == gold[i] for i in idx) / len(idx), + } + + # ---- novelty slice ------------------------------------------------------ + R["by_novelty"] = {} + buckets = defaultdict(list) + for i, r in enumerate(rows): + a_new = r["node_a"] not in train_nodes + b_new = r["node_b"] not in train_nodes + buckets["both novel" if a_new and b_new else + ("one novel" if a_new or b_new else "neither novel")].append(i) + for k, idx in buckets.items(): + R["by_novelty"][k] = {"n": len(idx), + "accuracy": sum(pred[i] == gold[i] for i in idx) / len(idx)} + + # ---- calibration -------------------------------------------------------- + right = [conf[i] for i in range(n) if pred[i] == gold[i]] + wrong = [conf[i] for i in range(n) if pred[i] != gold[i]] + R["calibration"] = { + "mean_conf_correct": sum(right) / len(right) if right else None, + "mean_conf_wrong": sum(wrong) / len(wrong) if wrong else None, + "wrong_above_0.9_conf": sum(1 for c in wrong if c > 0.9) / len(wrong) if wrong else None, + } + + # ---- symmetry: label(B,A) must be the mirror of label(A,B) ------------- + print("scoring swapped direction...") + sprobs = score([(r["text_b"], r["text_a"]) for r in rows]) + spred = sprobs.argmax(-1).tolist() + agree = sum(1 for i in range(n) if LABELS[spred[i]] == SWAP[LABELS[pred[i]]]) + R["symmetry"] = { + "consistent_rate": agree / n, + "note": "fraction where the swapped prediction mirrors the forward prediction", + "swapped_accuracy": sum( + 1 for i in range(n) if LABELS[spred[i]] == SWAP[rows[i]["label"]]) / n, + } + + # ---- best-of-n ranking: the referee test -------------------------------- + bon = run_best_of_n(args, rows, train_nodes, score) + if bon: + R["best_of_n"] = bon + + report(R) + out = args.out or os.path.join(os.path.dirname(args.adapter) or ".", "eval_report.json") + with open(out, "w") as f: + json.dump(R, f, indent=2) + print(f"\nwrote {out}") + + +def run_best_of_n(args, rows, train_nodes, score): + """Simulate best-of-n selection: one faithful candidate among n-1 drifted. + + The judge picks argmax P(equivalent). An exploitation event is a pick that + is not actually equivalent to the anchor -- the thing top-1 accuracy on + independent pairs cannot measure. + """ + sys.path.insert(0, ORACLE_DIR) + os.environ.setdefault("ETP_EQUATIONS", os.path.join(ORACLE_DIR, "data", "equations.txt")) + try: + from normalizer import parse_equation # noqa: F401 + sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) + from build_pair_dataset import augment + meta = json.load(open(os.path.join(ORACLE_DIR, "data", "matrix_meta.json"))) + nn = meta["n"] + buf = open(os.path.join(ORACLE_DIR, "data", "matrix.bin"), "rb").read() + laws = [l.strip() for l in + open(os.environ["ETP_EQUATIONS"], encoding="utf-8").read().splitlines() if l.strip()] + asts = [parse_equation(l) for l in laws] + except Exception as e: + print(f"\nbest-of-n skipped (oracle data unavailable): {str(e)[:120]}") + return None + + rng = random.Random(args.seed) + novel = sorted({r["node_a"] for r in rows if r["node_a"] not in train_nodes} | + {r["node_b"] for r in rows if r["node_b"] not in train_nodes}) + anchors = novel[:args.bon_anchors] if len(novel) >= args.bon_anchors else novel + print(f"\nbest-of-n: {len(anchors)} anchors x {args.bon_n} candidates") + + flat, meta_rows = [], [] + for a in anchors: + ai = a - 1 + anchor_txt = augment(asts[ai], rng) + cands = [(augment(asts[ai], rng), True)] # the faithful one + tries = 0 + while len(cands) < args.bon_n and tries < args.bon_n * 40: + tries += 1 + bi = rng.randrange(nn) + if bi == ai: + continue + fwd, bwd = buf[ai * nn + bi], buf[bi * nn + ai] + if fwd == 1 and bwd == 1: # genuinely equivalent, skip as a distractor + continue + if fwd > 1 or bwd > 1: # unresolved cell + continue + cands.append((augment(asts[bi], rng), False)) + for txt, faithful in cands: + flat.append((anchor_txt, txt)) + meta_rows.append((a, faithful)) + + p = score(flat)[:, 0].tolist() # P(equivalent) + groups = defaultdict(list) + for (anchor, faithful), pe in zip(meta_rows, p): + groups[anchor].append((pe, faithful)) + + top1, mrr, margins = 0, 0.0, [] + for g in groups.values(): + g.sort(key=lambda x: -x[0]) + rank = next(i for i, (_, f) in enumerate(g) if f) + 1 + top1 += rank == 1 + mrr += 1 / rank + faith = max(pe for pe, f in g if f) + best_drift = max((pe for pe, f in g if not f), default=0.0) + margins.append(faith - best_drift) + k = len(groups) + return { + "anchors": k, + "n_candidates": args.bon_n, + "faithful_ranked_top1": top1 / k, + "exploitation_rate": 1 - top1 / k, + "mrr": mrr / k, + "mean_margin": sum(margins) / k, + "negative_margin_rate": sum(1 for m in margins if m < 0) / k, + } + + +def report(R): + p = print + p("\n" + "=" * 66) + p(f"OVERALL ACCURACY {R['accuracy']:.4f} (n={R['n']}, chance=0.25)") + p("=" * 66) + p("\nCONFUSION MATRIX rows = truth, cols = predicted") + p(f"{'':14s}" + "".join(f"{l[:9]:>11s}" for l in LABELS)) + for g in LABELS: + p(f"{g:14s}" + "".join(f"{R['confusion'][g][x]:>11d}" for x in LABELS)) + p("\nPER-LABEL") + p(f"{'label':14s}{'recall':>10s}{'precision':>11s}{'support':>9s}") + for l in LABELS: + d = R["per_label"][l] + p(f"{l:14s}{d['recall']:>10.4f}{d['precision']:>11.4f}{d['support']:>9d}") + fe = R["false_equivalence"] + p(f"\nFALSE-EQUIVALENCE (the dangerous error)") + p(f" rate {fe['rate']:.4f} ({fe['count']} of {fe['of_non_equivalent']} non-equivalent rows)") + p(f" by true label: {fe['by_true_label']}") + if fe["mean_confidence"]: + p(f" mean confidence when it does this: {fe['mean_confidence']:.3f}") + p(f" missed equivalence (opposite error): {R['missed_equivalence']['rate']:.4f} " + f"-> {R['missed_equivalence']['into']}") + p("\nBY TIER (same-node = identical law re-rendered)") + for k, d in R["by_tier"].items(): + p(f" {k:12s} n={d['n']:<6d} acc={d['accuracy']:.4f}") + p("\nBY NOVELTY") + for k, d in R["by_novelty"].items(): + p(f" {k:14s} n={d['n']:<6d} acc={d['accuracy']:.4f}") + c = R["calibration"] + p(f"\nCALIBRATION") + p(f" mean confidence when correct {c['mean_conf_correct']:.3f} / when wrong " + f"{c['mean_conf_wrong']:.3f}") + p(f" wrong answers held above 0.9 confidence: {c['wrong_above_0.9_conf']:.3f}") + s = R["symmetry"] + p(f"\nSYMMETRY label(A,B) vs label(B,A)") + p(f" self-consistent {s['consistent_rate']:.4f} swapped-direction accuracy " + f"{s['swapped_accuracy']:.4f}") + if "best_of_n" in R: + b = R["best_of_n"] + p(f"\nBEST-OF-{b['n_candidates']} RANKING ({b['anchors']} novel anchors)") + p(f" faithful candidate ranked #1 : {b['faithful_ranked_top1']:.4f}") + p(f" EXPLOITATION RATE : {b['exploitation_rate']:.4f}") + p(f" mean reciprocal rank : {b['mrr']:.4f}") + p(f" mean margin (faithful - best drifted): {b['mean_margin']:+.4f}") + p() + + +if __name__ == "__main__": + main() diff --git a/judge/infra/INFRA.md b/judge/infra/INFRA.md new file mode 100644 index 0000000..c845ee0 --- /dev/null +++ b/judge/infra/INFRA.md @@ -0,0 +1,84 @@ +# `infra/` — optional GPU-rental helpers + +**Nothing here is required.** Using the judge needs no GPU, and retraining works +on any CUDA box via `train_judge.py` directly. These scripts exist only because +the runs behind `judge/README.md` were done on rented GPUs, and reproducing them +exactly is easier with the same glue. + +If you have your own GPU, ignore this directory. + +--- + +## `launch_thunder.sh` + +Laptop-side wrapper around the [Thunder Compute](https://www.thundercompute.com) +CLI (`tnr`), which rents per-minute GPUs. + +```bash +DISK=100 ./launch_thunder.sh up # create an instance +./launch_thunder.sh push # upload code + data +./launch_thunder.sh token # copy the local HF token up +./launch_thunder.sh pull # bring runs/ back +./launch_thunder.sh down # delete (this is what stops billing) +``` + +Override `GPU` (`a6000` / `a100` / `h100` / `l40`), `VCPUS`, `DISK`, `TEMPLATE` +by environment variable. + +Two behaviours worth knowing, both learned the hard way: + +- **`push` stashes `runs/` for the duration of the sync**, with a trap to restore + it even on failure. Without this, rsync uploads tens of GB of checkpoints the + instance never needs — and worse, can overwrite a log file a live training job + is writing to. +- **Billing stops on `delete`, not on disconnect.** An idle instance still + charges. `tnr status` should read `No instances found` when you are done. + +--- + +## What was deliberately left out + +Three scripts used during the original runs are **not** in the repo, because they +are operational scaffolding rather than research artifacts. Recreating them is a +few lines if you want them; what mattered is documented below. + +### `remote_run.sh` — instance-side driver + +Ran on the GPU box and wrapped four things: + +1. **`setup`** — `pip install "transformers>=4.44,<5" peft datasets accelerate`. + The upper pin matters: transformers 5.x renamed `torch_dtype=` to `dtype=`, + which breaks `train_judge.py`'s model loading. +2. **Preflight** — assert CUDA is visible, assert the gated Gemma download works, + and assert all four label words are single distinct tokens under the + tokenizer. That last check is now an assertion inside `train_judge.py`, since + a first-token collision would silently degrade the 4-way head to 3-way. +3. **`train`** — launch under `nohup` so a dropped SSH session does not kill a + multi-hour run. +4. **`status`** — parse the `\r`-laden progress log into something readable. + +### `watch.sh` / `watch-llama.sh` — progress monitors + +Polled the instance and emitted one line per state change: run started, run +finished with `test_acc`, OOM, traceback, or process death. Useful while +babysitting a six-hour sweep; useless afterwards. + +The one design note worth keeping: a monitor that greps only for the success +marker stays **silent** through a crashloop, and silence is indistinguishable +from "still training". Any replacement should match every terminal state, not +just the happy path. + +--- + +## Reproducing the runs without Thunder + +Nothing above is load-bearing. On any machine with a suitable GPU: + +```bash +python3 build_pair_dataset.py --out data/ +python3 train_judge.py --data data --out runs/judge-v0 \ + --micro-batch 64 --grad-accum 2 --eval-rows 2000 --keep-checkpoints 0 +``` + +See the memory table in `judge/README.md` for VRAM requirements per model and +micro-batch — `≥ 40 GB` is the practical floor at micro-batch 64. diff --git a/judge/infra/launch_thunder.sh b/judge/infra/launch_thunder.sh new file mode 100755 index 0000000..c995582 --- /dev/null +++ b/judge/infra/launch_thunder.sh @@ -0,0 +1,82 @@ +#!/usr/bin/env bash +# Provision a Thunder Compute instance and stage the judge finetune on it. +# Laptop side. Run `tnr login` once before the first use. +# +# ./launch_thunder.sh up # create the instance, print its id +# ./launch_thunder.sh push # upload finetune/ (scripts + 39 MB data) +# ./launch_thunder.sh token # copy the local HF token up (never printed) +# ./launch_thunder.sh pull # download runs/ back to the laptop +# ./launch_thunder.sh down # delete the instance (stops billing) +# +# Training itself runs on the instance: `tnr connect ` then +# `bash ~/finetune/remote_run.sh smoke` and `... remote_run.sh train`. +set -euo pipefail + +GPU="${GPU:-a6000}" # a6000 (48 GB) is ample for 2B + LoRA; a100 if you want headroom +VCPUS="${VCPUS:-8}" +DISK="${DISK:-200}" # merged bf16 model (~5 GB) + checkpoints + HF cache +TEMPLATE="${TEMPLATE:-base}" + +HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" + +# tnr lives in ~/.local/bin, which only non-login shells that source .zshrc see. +command -v tnr >/dev/null 2>&1 || export PATH="$HOME/.local/bin:$PATH" + +die() { echo "error: $*" >&2; exit 1; } + +case "${1:-}" in + up) + # tnr >= 2.0.71 dropped --mode and renamed --disk-size-gb to --disk. + tnr create --gpu "$GPU" --vcpus "$VCPUS" \ + --template "$TEMPLATE" --disk "$DISK" -y + echo + echo "Instance list:" + tnr status --no-wait + ;; + + push) + ID="${2:?usage: $0 push }" + # runs/ holds tens of GB the instance never needs, and syncing it has twice + # caused trouble: once uploading a 5.8 GB archive, once overwriting a log a + # live job was writing. Stash it for the duration of the sync. + STASH="" + if [ -d "$HERE/runs" ]; then + STASH="$(dirname "$HERE")/.runs_stash_$$" + mv "$HERE/runs" "$STASH" + trap 'mv "$STASH" "$HERE/runs" 2>/dev/null' EXIT INT TERM + fi + tnr scp "$HERE" "$ID:~/" + rc=$? + if [ -n "$STASH" ]; then + trap - EXIT INT TERM + mv "$STASH" "$HERE/runs" + fi + [ "$rc" -eq 0 ] || die "push failed (rc=$rc); runs/ restored" + echo "uploaded $HERE -> $ID:~/finetune (runs/ excluded)" + ;; + + token) + ID="${2:?usage: $0 token }" + TOKFILE="$HOME/.cache/huggingface/token" + [ -f "$TOKFILE" ] || die "no HF token at $TOKFILE — run 'hf auth login' first" + tnr scp "$TOKFILE" "$ID:~/hf_token" + echo "token copied to $ID:~/hf_token (remote_run.sh reads it from there)" + ;; + + pull) + ID="${2:?usage: $0 pull }" + mkdir -p "$HERE/runs" + tnr scp "$ID:~/finetune/runs" "$HERE/" + echo "downloaded runs/ -> $HERE/runs" + ;; + + down) + ID="${2:?usage: $0 down }" + tnr delete "$ID" + ;; + + *) + sed -n '2,12p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//' + exit 1 + ;; +esac diff --git a/judge/judge.py b/judge/judge.py new file mode 100755 index 0000000..2cdcf41 --- /dev/null +++ b/judge/judge.py @@ -0,0 +1,243 @@ +#!/usr/bin/env python3 +"""ETP relation judge — library and command-line interface. + +Classifies the logical relation between two equational laws: + + equivalent | weaker | stronger | incomparable + +The label describes B relative to A, matching the oracle's convention +(`weaker` = B drops constraints). Scoring is a single forward pass whose +logits are restricted to the four label tokens, so the judge cannot emit an +invalid answer and every prediction carries a calibrated probability. + +CLI +--- + judge.py compare "x * y = y * x" "a * b = b * a" + judge.py compare -A "..." -B "..." --json + judge.py batch pairs.jsonl --out scored.jsonl + judge.py rank intended.txt candidates.txt + +Library +------- + from judge import Judge + j = Judge("adapters/judge-v0") + j.compare("x * y = y * x", "a * b = b * a") # -> Verdict + j.probs([(a, b), ...]) # -> (n, 4) ndarray + +No GPU required. A single comparison takes a couple of seconds on CPU; +CUDA and Apple-Silicon MPS are used automatically when present. +""" +from __future__ import annotations + +import argparse +import json +import os +import sys +from dataclasses import dataclass, asdict + +LABELS = ("equivalent", "weaker", "stronger", "incomparable") +DEFAULT_BASE = "google/gemma-2-2b" + + +@dataclass +class Verdict: + label: str + confidence: float + probs: dict # label -> probability + text_a: str + text_b: str + + def __str__(self) -> str: + return f"{self.label} ({self.confidence:.3f})" + + +class Judge: + """A loaded relation judge. Construct once, score many times.""" + + def __init__(self, adapter: str, base: str | None = None, + device: str | None = None, dtype: str = "bfloat16"): + import torch + from peft import PeftModel + from transformers import AutoModelForCausalLM, AutoTokenizer + + self.torch = torch + if device is None: + device = ("cuda" if torch.cuda.is_available() + else "mps" if torch.backends.mps.is_available() else "cpu") + self.device = device + # bf16 is unreliable on CPU; fall back to fp32 there. + td = getattr(torch, dtype) if device != "cpu" else torch.float32 + + # The adapter records which base model it was trained against, so the + # caller does not have to remember. An explicit --base still wins. + if base is None: + cfg = os.path.join(adapter, "adapter_config.json") + if os.path.isfile(cfg): + with open(cfg) as f: + base = json.load(f).get("base_model_name_or_path") + base = base or DEFAULT_BASE + self.base_name = base + + # The tokenizer is saved beside the adapter; fall back to the base. + try: + self.tok = AutoTokenizer.from_pretrained(adapter) + except Exception: + self.tok = AutoTokenizer.from_pretrained(base) + + self.label_ids = [self.tok(" " + l, add_special_tokens=False).input_ids[0] + for l in LABELS] + if len(set(self.label_ids)) != len(LABELS): + raise SystemExit( + f"label first-tokens collide under {base}; the 4-way head would " + "silently degrade. Pick single-token label synonyms and retrain.") + + # Gemma-2 requires eager attention for its logit soft-capping. + attn = "eager" if "gemma-2" in base.lower() else "sdpa" + try: + model = AutoModelForCausalLM.from_pretrained( + base, dtype=td, attn_implementation=attn) + except TypeError: # transformers < 5 spells it torch_dtype + model = AutoModelForCausalLM.from_pretrained( + base, torch_dtype=td, attn_implementation=attn) + self.model = PeftModel.from_pretrained(model, adapter).to(device).eval() + + @staticmethod + def prompt(text_a: str, text_b: str) -> str: + return f"A: {text_a}\nB: {text_b}\nRelation:" + + def probs(self, pairs, batch_size: int = 32): + """pairs: iterable of (text_a, text_b). Returns an (n, 4) float array.""" + torch = self.torch + pairs = list(pairs) + out = [] + with torch.no_grad(): + for i in range(0, len(pairs), batch_size): + chunk = pairs[i:i + batch_size] + enc = [self.tok(self.prompt(a, b), add_special_tokens=True).input_ids + for a, b in chunk] + m = max(len(e) for e in enc) + pad = self.tok.pad_token_id + if pad is None: + pad = self.tok.eos_token_id + ids = torch.tensor([e + [pad] * (m - len(e)) for e in enc], + device=self.device) + att = torch.tensor([[1] * len(e) + [0] * (m - len(e)) for e in enc], + device=self.device) + last = torch.tensor([len(e) - 1 for e in enc], device=self.device) + lg = self.model(input_ids=ids, attention_mask=att).logits + # Take the logits at each row's true final position, then keep + # only the four label tokens -- this is what makes an invalid + # answer impossible rather than merely unlikely. + lg = lg[torch.arange(lg.size(0), device=self.device), last][:, self.label_ids] + out.append(torch.softmax(lg.float(), -1).cpu()) + return torch.cat(out).numpy() if out else None + + def compare(self, text_a: str, text_b: str) -> Verdict: + p = self.probs([(text_a, text_b)])[0] + k = int(p.argmax()) + return Verdict(label=LABELS[k], confidence=float(p[k]), + probs={l: float(v) for l, v in zip(LABELS, p)}, + text_a=text_a, text_b=text_b) + + def rank(self, intended: str, candidates): + """Order candidates by P(equivalent to `intended`), best first. + + This is the operation best-of-n selection performs, so it is the one + that matters when the judge is used as a checker under optimisation + pressure -- not top-1 accuracy on independent pairs. + """ + cands = list(candidates) + p = self.probs([(intended, c) for c in cands]) + scored = [(c, float(p[i][0]), LABELS[int(p[i].argmax())]) + for i, c in enumerate(cands)] + return sorted(scored, key=lambda t: -t[1]) + + +# --------------------------------------------------------------------------- CLI + +def _add_model_args(p): + p.add_argument("--adapter", default=os.environ.get("JUDGE_ADAPTER"), + help="path to a trained adapter dir (or set JUDGE_ADAPTER)") + p.add_argument("--base", default=None, + help="base model id; read from adapter_config.json if omitted") + p.add_argument("--device", default=None, help="cuda | mps | cpu (auto by default)") + + +def main() -> None: + ap = argparse.ArgumentParser( + prog="judge", description="Classify the relation between two equational laws.") + sub = ap.add_subparsers(dest="cmd", required=True) + + c = sub.add_parser("compare", help="score a single pair") + c.add_argument("a", nargs="?"); c.add_argument("b", nargs="?") + c.add_argument("-A", dest="a_flag"); c.add_argument("-B", dest="b_flag") + c.add_argument("--json", action="store_true") + _add_model_args(c) + + b = sub.add_parser("batch", help="score a JSONL file of {text_a, text_b} rows") + b.add_argument("infile") + b.add_argument("--out", default=None, help="write JSONL here (default: stdout)") + b.add_argument("--batch-size", type=int, default=32) + _add_model_args(b) + + r = sub.add_parser("rank", help="rank candidates by P(equivalent to intended)") + r.add_argument("intended", help="the intended law, or a file containing it") + r.add_argument("candidates", help="file with one candidate per line") + r.add_argument("--json", action="store_true") + _add_model_args(r) + + args = ap.parse_args() + if not args.adapter: + ap.error("no adapter given; pass --adapter PATH or set JUDGE_ADAPTER") + + j = Judge(args.adapter, base=args.base, device=args.device) + + if args.cmd == "compare": + a = args.a_flag or args.a + b_ = args.b_flag or args.b + if not a or not b_: + ap.error("need two equations: judge.py compare 'A' 'B'") + v = j.compare(a, b_) + if args.json: + print(json.dumps(asdict(v), indent=2)) + else: + print(f"A: {a}\nB: {b_}\n\n {v.label} (confidence {v.confidence:.3f})\n") + for lab, p in sorted(v.probs.items(), key=lambda t: -t[1]): + bar = "#" * int(round(p * 40)) + print(f" {lab:14s} {p:6.3f} {bar}") + + elif args.cmd == "batch": + rows = [json.loads(l) for l in open(args.infile, encoding="utf-8") if l.strip()] + p = j.probs([(r["text_a"], r["text_b"]) for r in rows], batch_size=args.batch_size) + sink = open(args.out, "w", encoding="utf-8") if args.out else sys.stdout + agree = total = 0 + for row, pr in zip(rows, p): + k = int(pr.argmax()) + row["pred"] = LABELS[k] + row["confidence"] = float(pr[k]) + row["probs"] = {l: float(v) for l, v in zip(LABELS, pr)} + if "label" in row: # gold present -> score it + total += 1 + agree += row["label"] == row["pred"] + sink.write(json.dumps(row, ensure_ascii=False) + "\n") + if args.out: + sink.close() + if total: + print(f"accuracy {agree}/{total} = {agree / total:.4f}", file=sys.stderr) + + elif args.cmd == "rank": + intended = (open(args.intended, encoding="utf-8").read().strip() + if os.path.isfile(args.intended) else args.intended) + cands = [l.strip() for l in open(args.candidates, encoding="utf-8") if l.strip()] + ranked = j.rank(intended, cands) + if args.json: + print(json.dumps([{"candidate": c, "p_equivalent": p, "pred": lab} + for c, p, lab in ranked], indent=2)) + else: + print(f"intended: {intended}\n") + for i, (c, p, lab) in enumerate(ranked, 1): + print(f" {i:3d}. p(equiv)={p:.4f} [{lab:12s}] {c}") + + +if __name__ == "__main__": + main() diff --git a/judge/sweep.sh b/judge/sweep.sh new file mode 100755 index 0000000..875a583 --- /dev/null +++ b/judge/sweep.sh @@ -0,0 +1,116 @@ +#!/usr/bin/env bash +# LoRA hyperparameter sweep — runs on the Thunder instance. +# +# bash ~/finetune/sweep.sh smoke # 8 steps of every config (~10 min) +# bash ~/finetune/sweep.sh run # the real sweep, sequential (~6 h) +# bash ~/finetune/sweep.sh summary # table of finished runs +# +# One variable per run, seed fixed, everything else frozen. Each run is a fresh +# start from the base model -- LoRA ranks and layer sets have different shapes, +# so a run cannot be branched off another run's checkpoint. +# +# Already-finished runs are skipped, so this is safe to re-invoke after a crash. +set -uo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")" +# Overridable so the same sweep runs on a second model family: +# MODEL=meta-llama/Llama-3.1-8B OUT=runs/sweep-llama MB=64 GA=2 \ +# SKIP_RUNS=baseline bash sweep.sh run +MODEL="${MODEL:-google/gemma-2-2b}" +OUT="${OUT:-runs/sweep}" +MB="${MB:-128}" +GA="${GA:-1}" +SKIP_RUNS="${SKIP_RUNS:-}" +COMMON="--data data --seed 0 --model $MODEL --micro-batch $MB --grad-accum $GA --eval-rows 2000 --keep-checkpoints 0" + +# name extra flags varies +# layers=all + modules=all is the memory peak: Gemma-2's 256k vocab puts ~13.7 GB +# in the logits region, and every adapted layer retains its input for backward. +# micro-batch 128 OOMs on a 48 GB A6000, so these three drop to 64 x 2. Effective +# batch stays 128, so they remain mathematically comparable with the rest. +CONFIGS=( + "baseline|--rank 16 --micro-batch 64 --grad-accum 2|reference point" + "rank4|--rank 4 --micro-batch 64 --grad-accum 2 |rank down" + "rank64|--rank 64 --micro-batch 64 --grad-accum 2|rank up" + "early|--rank 16 --lora-layers early |layer position" + "middle|--rank 16 --lora-layers middle |layer position" + "late|--rank 16 --lora-layers late |layer position" + # attn peaked at 45.3/47.4 GB at micro-batch 128 -- survived, but with no + # margin. mlp adapts down_proj, whose input is 9216-dim vs 2304, so it stores + # ~1.3 GB more and would almost certainly OOM. Both drop to 64 x 2. + "attn|--rank 16 --lora-modules attn --micro-batch 64 --grad-accum 2|module type" + "mlp|--rank 16 --lora-modules mlp --micro-batch 64 --grad-accum 2|module type" +) + +load_token() { + if [ -z "${HF_TOKEN:-}" ] && [ -f "$HOME/hf_token" ]; then + HF_TOKEN="$(tr -d '[:space:]' < "$HOME/hf_token")"; export HF_TOKEN + fi + [ -n "${HF_TOKEN:-}" ] || { echo "error: no HF_TOKEN and no ~/hf_token" >&2; exit 1; } + export HUGGING_FACE_HUB_TOKEN="$HF_TOKEN" +} + +case "${1:-}" in + smoke) + load_token + for c in "${CONFIGS[@]}"; do + name="${c%%|*}"; rest="${c#*|}"; flags="${rest%%|*}" + echo "=== smoke: $name ===" + python3 train_judge.py $COMMON $flags --dry-run --out "$OUT/_smoke_$name" 2>&1 \ + | grep -E "LoRA:|trainable params|error|Error|Traceback" || true + done + echo "smoke done -- if every config printed trainable params, the sweep will run" + ;; + + run) + load_token + mkdir -p "$OUT" + START=$(date +%s) + for c in "${CONFIGS[@]}"; do + name="${c%%|*}"; rest="${c#*|}"; flags="${rest%%|*}"; note="${rest##*|}" + case " $SKIP_RUNS " in *" $name "*) + echo "=== skip $name (in SKIP_RUNS) ==="; continue ;; + esac + if [ -f "$OUT/$name/final_metrics.json" ]; then + echo "=== skip $name (already finished) ==="; continue + fi + echo "=== $name [$note] $(date -u +%H:%M:%S) ===" + python3 -u train_judge.py $COMMON $flags --out "$OUT/$name" > "$OUT/$name.log" 2>&1 + if [ -f "$OUT/$name/final_metrics.json" ]; then + acc=$(python3 -c "import json;print(json.load(open('$OUT/$name/final_metrics.json'))['test_acc'])" 2>/dev/null) + echo " done test_acc=$acc" + else + echo " FAILED -- see $OUT/$name.log"; tail -5 "$OUT/$name.log" + fi + done + echo "sweep finished in $(( ($(date +%s) - START) / 60 )) min" + bash "$0" summary + ;; + + summary) + python3 - <<'PY' +import json, os, glob +rows = [] +for f in sorted(glob.glob("runs/sweep/*/final_metrics.json")): + name = os.path.basename(os.path.dirname(f)) + m = json.load(open(f)) + cfg = json.load(open(os.path.join(os.path.dirname(f), "run_config.json"))) + ck = len(glob.glob(os.path.join(os.path.dirname(f), "checkpoint-*"))) + rows.append((name, cfg["rank"], cfg["lora_layers"], cfg["lora_modules"], + m["test_acc"], m.get("test_acc_equivalent"), ck)) +if not rows: + print("no finished runs yet"); raise SystemExit +base = next((r[4] for r in rows if r[0] == "baseline"), None) +print(f"\n{'run':10s}{'r':>4s}{'layers':>9s}{'modules':>9s}{'test_acc':>10s}" + f"{'vs base':>9s}{'equiv':>8s}{'ckpts':>7s}") +for n, r, l, mo, a, e, ck in sorted(rows, key=lambda x: -x[4]): + d = f"{(a - base) * 100:+.2f}" if base and n != "baseline" else "--" + print(f"{n:10s}{r:>4d}{l:>9s}{mo:>9s}{a:>10.4f}{d:>9s}" + f"{(e or 0):>8.4f}{ck:>7d}") +PY + ;; + + *) + sed -n '2,12p' "${BASH_SOURCE[0]}" | sed 's/^# \{0,1\}//' + exit 1 ;; +esac diff --git a/judge/train_judge.py b/judge/train_judge.py new file mode 100644 index 0000000..b0784e3 --- /dev/null +++ b/judge/train_judge.py @@ -0,0 +1,245 @@ +#!/usr/bin/env python3 +"""Finetune the ETP relation judge (Gemma-2-2B base + LoRA). + +Runs on a single GPU (A100 recommended; fits a T4 with --load-4bit). +Reads the JSONL splits from build_pair_dataset.py. Label = single next token +after "Relation:", loss masked to that token only. + +Smoke-test first (5 min): + python3 train_judge.py --data data --dry-run + +Full run: + python3 train_judge.py --data data --out runs/judge-v0 \ + --push-to-hub yourname/etp-judge-gemma2-2b-v0 + +Env: HF_TOKEN must be set (Gemma is a gated model; accept the license on the +model page first). All hyperparameters are flags with sane defaults. +""" +from __future__ import annotations + +import argparse +import json +import os + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--data", default="data") + ap.add_argument("--out", default="runs/judge-v0") + ap.add_argument("--model", default="google/gemma-2-2b") + ap.add_argument("--push-to-hub", default=None, + help="private HF repo id to upload adapter + merged model") + ap.add_argument("--save-merged", action="store_true", + help="also write the merged bf16 model locally (~5 GB); " + "skip it and merge from the adapter later instead") + ap.add_argument("--lr", type=float, default=1e-4) + ap.add_argument("--rank", type=int, default=16) + ap.add_argument("--lora-layers", default="all", + help="all | early | middle | late | a-b (inclusive, 0-based)") + ap.add_argument("--lora-modules", default="all", choices=["all", "attn", "mlp"]) + ap.add_argument("--micro-batch", type=int, default=32) + ap.add_argument("--grad-accum", type=int, default=4) + ap.add_argument("--max-steps", type=int, default=2000) + ap.add_argument("--eval-every", type=int, default=200) + ap.add_argument("--eval-rows", type=int, default=0, + help="subsample the pairs_val/classes_val CURVE evals to N rows " + "(0 = full). The final test eval always uses the full split.") + ap.add_argument("--keep-checkpoints", type=int, default=0, + help="save_total_limit; 0 = keep every checkpoint (needed for " + "the interp trajectory)") + ap.add_argument("--seed", type=int, default=0) + ap.add_argument("--load-4bit", action="store_true", + help="QLoRA mode for 16GB GPUs (T4); skip on A100") + ap.add_argument("--dry-run", action="store_true", + help="200 train rows, 8 steps, no upload") + args = ap.parse_args() + + import torch + from datasets import load_dataset + from peft import LoraConfig, get_peft_model + from transformers import (AutoModelForCausalLM, AutoTokenizer, + Trainer, TrainingArguments, set_seed) + + set_seed(args.seed) + LABELS = ["equivalent", "weaker", "stronger", "incomparable"] + + tok = AutoTokenizer.from_pretrained(args.model) + # Each label must be exactly one token after "Relation:" — verify, don't + # assume: leading-space piece differs per tokenizer. + label_ids = {} + for lab in LABELS: + ids = tok(" " + lab, add_special_tokens=False).input_ids + label_ids[lab] = ids[0] + if len(ids) > 1: + print(f"note: ' {lab}' is {len(ids)} tokens; training on first, " + "eval compares first-token logits (still a valid 4-way head)") + # Multi-token labels are fine ONLY while their first tokens stay distinct. + # If two collided the head would silently degrade to 3-way and every number + # downstream would be quiet garbage, so fail loudly instead of printing. + if len(set(label_ids.values())) != len(LABELS): + collisions = {tok.convert_ids_to_tokens([i])[0]: [l for l in LABELS if label_ids[l] == i] + for i in set(label_ids.values())} + raise SystemExit(f"label first-tokens collide under {args.model}: {collisions}\n" + "pick single-token label synonyms before training.") + + def fmt(row): + prompt = f"A: {row['text_a']}\nB: {row['text_b']}\nRelation:" + ids = tok(prompt, add_special_tokens=True).input_ids + input_ids = ids + [label_ids[row["label"]]] + labels = [-100] * len(ids) + [label_ids[row["label"]]] + return {"input_ids": input_ids, "labels": labels} + + files = {s: os.path.join(args.data, f"{s}.jsonl") + for s in ["train", "pairs_val", "classes_val", "test"]} + ds = load_dataset("json", data_files=files) + if args.dry_run: + ds["train"] = ds["train"].select(range(200)) + for s in ["pairs_val", "classes_val", "test"]: + ds[s] = ds[s].select(range(100)) + keep = ["input_ids", "labels"] + ds = ds.map(fmt, remove_columns=[c for c in ds["train"].column_names if c not in keep]) + # Curve evals are 62% as much forward compute as training at full size. + # Subsampling them costs ~0.2pp of standard error and buys back ~35% of + # wall clock; `test` is deliberately left whole for the headline number. + curve = {s: ds[s] for s in ["pairs_val", "classes_val"]} + if args.eval_rows and not args.dry_run: + for s in curve: + curve[s] = ds[s].select(range(min(args.eval_rows, len(ds[s])))) + + # Gemma-2 needs eager attention for its logit soft-capping; everything else + # is faster on sdpa. + attn = "eager" if "gemma-2" in args.model.lower() else "sdpa" + model_kwargs = dict(torch_dtype=torch.bfloat16, attn_implementation=attn) + if args.load_4bit: + from transformers import BitsAndBytesConfig + model_kwargs["quantization_config"] = BitsAndBytesConfig( + load_in_4bit=True, bnb_4bit_compute_dtype=torch.bfloat16, + bnb_4bit_quant_type="nf4") + model = AutoModelForCausalLM.from_pretrained(args.model, **model_kwargs) + + ATTN = ["q_proj", "k_proj", "v_proj", "o_proj"] + MLP = ["gate_proj", "up_proj", "down_proj"] + mods = {"all": ATTN + MLP, "attn": ATTN, "mlp": MLP}[args.lora_modules] + missing = [m for m in mods if not any(m in n for n, _ in model.named_modules())] + if missing: + raise SystemExit(f"{args.model} has no modules named {missing}; " + "check the architecture's projection names.") + + n_layers = model.config.num_hidden_layers + third = n_layers // 3 + spans = {"all": None, + "early": list(range(0, third)), + "middle": list(range(third, 2 * third)), + "late": list(range(2 * third, n_layers))} + if args.lora_layers in spans: + layers = spans[args.lora_layers] + else: + lo, hi = (int(x) for x in args.lora_layers.split("-")) + layers = list(range(lo, hi + 1)) + print(f"LoRA: r={args.rank} modules={args.lora_modules} " + f"layers={args.lora_layers} -> {'all ' + str(n_layers) if layers is None else layers}") + + lora = LoraConfig(r=args.rank, lora_alpha=2 * args.rank, lora_dropout=0.05, + target_modules=mods, task_type="CAUSAL_LM", + layers_to_transform=layers, + layers_pattern="layers" if layers is not None else None) + model = get_peft_model(model, lora) + model.print_trainable_parameters() + + def collate(batch): + maxlen = max(len(b["input_ids"]) for b in batch) + pad = tok.pad_token_id or tok.eos_token_id + input_ids, labels, attn = [], [], [] + for b in batch: + k = maxlen - len(b["input_ids"]) + input_ids.append(b["input_ids"] + [pad] * k) + labels.append(b["labels"] + [-100] * k) + attn.append([1] * len(b["input_ids"]) + [0] * k) + return {"input_ids": torch.tensor(input_ids), + "labels": torch.tensor(labels), + "attention_mask": torch.tensor(attn)} + + id_list = [label_ids[lab] for lab in LABELS] + + def compute_metrics(pred): + import numpy as np + logits, labels = pred + # logits arrive already argmaxed over the label-token position via + # preprocess_logits_for_metrics below: shape (n, 4) restricted scores + rows, gold = [], [] + for lg, lb in zip(logits, labels): + pos = np.where(lb != -100)[0] + if len(pos) == 0: + continue + rows.append(lg) + gold.append(id_list.index(lb[pos[0]])) + rows = np.array(rows); gold = np.array(gold) + pred4 = rows.argmax(-1) + acc = float((pred4 == gold).mean()) + per = {f"acc_{LABELS[i]}": float((pred4[gold == i] == i).mean()) + for i in range(4) if (gold == i).any()} + return {"acc": acc, **per} + + def preprocess_logits(logits, labels): + # keep only the 4 label-token logits at the label position + import torch as t + pos = (labels != -100).float().argmax(-1) - 1 # logits index preceding label + idx = pos.clamp(min=0).long() + gathered = logits[t.arange(logits.size(0)), idx][:, id_list] + return gathered + + steps = 8 if args.dry_run else args.max_steps + targs = TrainingArguments( + output_dir=args.out, + per_device_train_batch_size=args.micro_batch, + per_device_eval_batch_size=args.micro_batch, + gradient_accumulation_steps=args.grad_accum, + learning_rate=args.lr, lr_scheduler_type="cosine", + warmup_ratio=0.03, max_steps=steps, logging_steps=10, + eval_strategy="steps", eval_steps=args.eval_every if not args.dry_run else 4, + save_steps=args.eval_every, + save_total_limit=args.keep_checkpoints or None, + bf16=True, report_to="none", seed=args.seed, + dataloader_num_workers=2, remove_unused_columns=False, + ) + trainer = Trainer( + model=model, args=targs, + train_dataset=ds["train"], + eval_dataset={"pairs": curve["pairs_val"], "classes": curve["classes_val"]}, + data_collator=collate, compute_metrics=compute_metrics, + preprocess_logits_for_metrics=preprocess_logits, + ) + trainer.train() + + print("\nFinal test-split eval (held-out classes):") + test_metrics = trainer.evaluate(ds["test"], metric_key_prefix="test") + print(test_metrics) + + # Adapter + tokenizer are the durable artifact: ~80 MB, and the merged + # model is reconstructible from them plus the base weights at any time. + adapter_dir = os.path.join(args.out, "adapter-final") + trainer.save_model(adapter_dir) + tok.save_pretrained(adapter_dir) + print(f"saved adapter + tokenizer to {adapter_dir}") + + with open(os.path.join(args.out, "run_config.json"), "w") as f: + json.dump(vars(args), f, indent=2, default=str) + with open(os.path.join(args.out, "final_metrics.json"), "w") as f: + json.dump(test_metrics, f, indent=2, default=str) + + if args.save_merged and not args.dry_run: + merged_dir = os.path.join(args.out, "merged-bf16") + model.merge_and_unload().save_pretrained(merged_dir) + tok.save_pretrained(merged_dir) + print(f"saved merged bf16 model to {merged_dir} (~5 GB)") + + if args.push_to_hub and not args.dry_run: + model.push_to_hub(args.push_to_hub, private=True) + tok.push_to_hub(args.push_to_hub, private=True) + merged = model.merge_and_unload() + merged.push_to_hub(args.push_to_hub + "-merged", private=True) + print("pushed adapter and merged model to the Hub") + + +if __name__ == "__main__": + main() From 3956e64b5a4763e0aaa7c063f67085b9c1c82dde Mon Sep 17 00:00:00 2001 From: karthikreddy-02 Date: Wed, 23 Sep 2026 18:38:37 -0600 Subject: [PATCH 2/2] Expose callable judges and a portable configuration-driven pipeline --- .github/workflows/judge-tests.yml | 28 +++ .gitignore | 1 + README.md | 3 + configs/judge-v0-baseline.yaml | 7 + configs/rep-lean.yaml | 9 + judge/HISTORICAL_RESULTS.md | 294 +++++++++++++++++++++++++++++ judge/README.md | 303 ++++++------------------------ judge/__init__.py | 5 + judge/__main__.py | 4 + judge/build_pair_dataset.py | 14 +- judge/judge.py | 67 +++++-- judge/requirements.txt | 6 + judge/train_judge.py | 27 ++- pipeline/README.md | 80 ++++++++ pipeline/__init__.py | 4 + pipeline/config.py | 241 ++++++++++++++++++++++++ pipeline/metrics.py | 43 +++++ pipeline/representations.py | 211 +++++++++++++++++++++ pipeline/requirements.txt | 4 + pipeline/run.py | 189 +++++++++++++++++++ tests/test_callable_interfaces.py | 195 +++++++++++++++++++ tests/test_pipeline_contract.py | 107 +++++++++++ 22 files changed, 1586 insertions(+), 256 deletions(-) create mode 100644 .github/workflows/judge-tests.yml create mode 100644 configs/judge-v0-baseline.yaml create mode 100644 configs/rep-lean.yaml create mode 100644 judge/HISTORICAL_RESULTS.md create mode 100644 judge/__init__.py create mode 100644 judge/__main__.py create mode 100644 judge/requirements.txt create mode 100644 pipeline/README.md create mode 100644 pipeline/__init__.py create mode 100644 pipeline/config.py create mode 100644 pipeline/metrics.py create mode 100644 pipeline/representations.py create mode 100644 pipeline/requirements.txt create mode 100644 pipeline/run.py create mode 100644 tests/test_callable_interfaces.py create mode 100644 tests/test_pipeline_contract.py diff --git a/.github/workflows/judge-tests.yml b/.github/workflows/judge-tests.yml new file mode 100644 index 0000000..094c2ed --- /dev/null +++ b/.github/workflows/judge-tests.yml @@ -0,0 +1,28 @@ +name: Judge interfaces + +on: + pull_request: + paths: + - 'judge/**' + - 'pipeline/**' + - 'configs/**' + - 'tests/**' + - '.github/workflows/judge-tests.yml' + workflow_dispatch: + +permissions: + contents: read + +jobs: + offline-tests: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - uses: actions/setup-python@v5 + with: + python-version: '3.11' + - run: python -m pip install -r pipeline/requirements.txt + - run: python -W error::ResourceWarning -m unittest discover -s tests -v + - run: python -m pipeline.run --config configs/judge-v0-baseline.yaml --stage all --dry-run + - run: python -m pipeline.run --config configs/rep-lean.yaml --stage all --dry-run + - run: python -m judge --help diff --git a/.gitignore b/.gitignore index 5e88b90..47284ba 100644 --- a/.gitignore +++ b/.gitignore @@ -3,5 +3,6 @@ __pycache__/ # judge/: generated dataset and trained weights (regenerable / too large) judge/data/ +judge/data-*/ judge/runs/ *.tar.gz diff --git a/README.md b/README.md index cdafa2e..ebc9901 100644 --- a/README.md +++ b/README.md @@ -38,6 +38,9 @@ The suggested week-by-week plan: [ONBOARDING.md](ONBOARDING.md). ## Instruments +- [`judge/`](judge/README.md) — callable learned relation judge and CLI. +- [`pipeline/`](pipeline/README.md) — configuration-driven data, training, and evaluation functions. + - [`oracle/`](oracle/README.md) — the formal semantic oracle: normalizes a (possibly messy, LLM-emitted) magma equation, maps it to its node in the [Equational Theories Project](https://github.com/teorth/equational_theories) catalogue, and labels an (intended, generated) pair `equivalent / weaker / stronger / incomparable / unknown` using only proven implications. See its README for setup and usage. diff --git a/configs/judge-v0-baseline.yaml b/configs/judge-v0-baseline.yaml new file mode 100644 index 0000000..bdd1917 --- /dev/null +++ b/configs/judge-v0-baseline.yaml @@ -0,0 +1,7 @@ +# Existing label-balanced sampler and original infix input format. +name: judge-v0-baseline +seed: 0 +data: + out_dir: judge/data +train: + out_dir: judge/runs/judge-v0 diff --git a/configs/rep-lean.yaml b/configs/rep-lean.yaml new file mode 100644 index 0000000..5044d57 --- /dev/null +++ b/configs/rep-lean.yaml @@ -0,0 +1,9 @@ +# Separate data and output paths prevent accidental reuse of infix examples. +name: judge-lean +seed: 0 +data: + out_dir: judge/data-lean +render: + kind: lean +train: + out_dir: judge/runs/judge-lean diff --git a/judge/HISTORICAL_RESULTS.md b/judge/HISTORICAL_RESULTS.md new file mode 100644 index 0000000..d44fbc5 --- /dev/null +++ b/judge/HISTORICAL_RESULTS.md @@ -0,0 +1,294 @@ +# Historical judge report + +This is the original single-seed report. Its prose has not been revalidated. +Use the current README for supported interfaces and limitations. + +Known qualifications: + +- `pairs_val` can repeat underlying training pairs; its gap does not prove absence of memorization. +- Most class-held-out pairs have only one novel endpoint. +- The original sampler contains many same-law equivalences and a shape/label association. +- Release dates exclude exposure to later ETP artifacts, not exposure to the underlying mathematics. +- LoRA placement experiments measure adaptation choices, not the full location of computation. +- Restricted label scores are not calibrated probabilities of correctness. +- Historical accuracy does not establish reliability on generated translations. + +The original report follows for provenance. Do not repeat its stronger claims without new evidence. + +--- + +# `judge` — a trained relation judge for the ETP fragment + +Classifies the logical relation between two equational laws: + +``` +equivalent | weaker | stronger | incomparable +``` + +The label describes **B relative to A**, matching `oracle/`'s convention +(`weaker` = B drops constraints). + +This replaces a stack of prompted LLMs with a single measurable artifact: known +training data, a known split, and weights we own and can probe. It is trained on +the ETP implication matrix, where every label is Lean-verified rather than judged. + +**No GPU required to use it.** One comparison is a couple of seconds on CPU; +CUDA and Apple-Silicon MPS are used automatically when available. + +--- + +## Quickstart + +```bash +pip install torch transformers peft +huggingface-cli login # Gemma is gated: accept the license first + +export JUDGE_ADAPTER=/path/to/judge-v0/adapter-final +python3 judge.py compare "x * y = y * x" "a * b = b * a" +``` + +``` +A: x * y = y * x +B: a * b = b * a + + equivalent (confidence 1.000) + + equivalent 1.000 ######################################## + weaker 0.000 + stronger 0.000 + incomparable 0.000 +``` + +First run downloads the ~5 GB base model. The adapter itself is ~113 MB. + +### Commands + +| command | purpose | +|---|---| +| `compare A B` | score one pair; `--json` for machine output | +| `batch pairs.jsonl` | score a JSONL of `{text_a, text_b}` rows; reports accuracy if `label` present | +| `rank intended candidates.txt` | order candidates by `P(equivalent)` — what best-of-n selection actually does | + +`rank` is the one to reach for when using the judge as a *checker*: + +``` +$ python3 judge.py rank "x * y = y * x" candidates.txt +intended: x * y = y * x + + 1. p(equiv)=0.9997 [equivalent ] p * q = q * p + 2. p(equiv)=0.9996 [equivalent ] a * b = b * a + 3. p(equiv)=0.0383 [weaker ] x = x + 4. p(equiv)=0.0017 [incomparable] x * (y * z) = (x * y) * z +``` + +### As a library + +Downstream interpretability tooling should import rather than shell out: + +```python +from judge import Judge + +j = Judge("path/to/adapter-final") # base model read from adapter_config.json +v = j.compare("x * y = y * x", "a * b = b * a") +v.label, v.confidence, v.probs + +j.probs([(a, b), ...]) # (n, 4) array, no argmax applied +j.rank(intended, candidates) # [(candidate, p_equivalent, pred), ...] +j.model # the live PeftModel, for hooks/probing +``` + +`Judge.model` is exposed deliberately: activation capture, SAE work, and +ablations attach there. + +--- + +## How it works + +Prompt is `"A: {eq}\nB: {eq}\nRelation:"`. The label is a **single next token** +with loss masked to that token alone. At inference the logits are **restricted to +the four label tokens** before the argmax, so the judge cannot emit an invalid +answer — a wrong prediction is always one of the other three, never gibberish. + +`unknown` is not a class. The 190 unresolved matrix cells are never sampled, so +the judge cannot abstain. If you need abstention, threshold on `confidence` — +but note calibration is mediocre (see Caveats). + +--- + +## Training data + +Built by `build_pair_dataset.py` from `oracle/`'s matrix. Labels are two byte +reads from the Lean-verified implication matrix — **no prover runs at training +time**; the certification is a frozen artifact of the +`2024-11-10-outcomes` snapshot. + +```bash +cd judge +python3 build_pair_dataset.py --out data/ +``` + +Requires `oracle/data/matrix.bin` — see `oracle/README.md` to build it. +The generated `data/` is gitignored (39 MB, fully regenerable). + +### The split is the important part + +**By equivalence class, never by pair.** 1,415 classes → **1,064 train / 140 val +/ 211 test**. Laws inside a class are mutually derivable, so a random pair split +puts the same law on both sides and a lookup table scores perfectly. + +The `x = y` class alone is **1,496 laws — 31.9% of the catalogue** — and is +*pinned into training*. A random split could otherwise have removed a third of +all laws from training. + +Four splits, and the contrast between two of them is the measurement: + +| split | composition | measures | +|---|---|---| +| `train` | 240,000 pairs from train classes | — | +| `pairs_val` | new pairings of **familiar** laws | interpolation | +| `classes_val` | ≥1 law from held-out classes | novel theories | +| `test` | ≥1 law from test classes, never selected on | final | + +**`pairs_val − classes_val` is the memorization measurement.** Everything else is +held constant; only law familiarity varies. A pair-level memoriser fails +`pairs_val`; a law-level lookup passes it and fails `classes_val`. Only a +relational model passes both. + +--- + +## Results + +Trained models, `test_acc` on 211 held-out equivalence classes (chance = 0.25): + +| config | Gemma-2-2B | Llama-3.1-8B | +|---|---|---| +| **baseline** (r16, all layers, all modules) | **0.9619** | **0.9649** | +| rank 4 | 0.9566 | 0.9596 | +| rank 64 | 0.9625 | 0.9647 | +| early layers only | 0.9414 | 0.9619 | +| middle layers only | 0.9494 | 0.9457 | +| **late layers only** | **0.8625** | **0.8491** | +| attention only | 0.9523 | 0.9629 | +| MLP only | 0.9601 | 0.9631 | + +Memorization gap: **1.91 pp** (Gemma), **1.25 pp** (Llama). Both trained a judge, +not a lookup table. + +**Neither base model can be contaminated.** ETP launched 2024-09-25; Gemma-2 was +released 2024-06-27 and Llama-3.1 on 2024-07-23. Both predate the project's +existence, which is a stronger guarantee than any stated data cutoff. + +### What replicates across both families + +- **Rank saturates at or below 16.** Quadrupling the adapter moves accuracy + +0.06 pp / −0.01 pp. Cutting to r=4 costs exactly −0.53 pp on both. +- **Late-only collapses** — 5–6× worse than any other layer choice. The relation + is built in the front two-thirds, not decided at the end. Note `late` has + *more* adapted layers than `early` (10 vs 8 on Gemma), so this is positional, + not a capacity effect. +- **MLP-only is nearly free** — −0.18 pp in both. + +### What does NOT replicate — do not claim + +- **The attention/MLP asymmetry is a Gemma artifact.** Gemma: attention-only cost + 5× more than MLP-only. Llama: they are equivalent. "The computation lives in + the MLP" does not survive replication. +- **Early-vs-middle ordering flips** between families. + +--- + +## Retraining + +Any CUDA box with **≥ 40 GB** works. Memory is dominated by the vocabulary-sized +logits plus retained activations for every adapted layer. + +```bash +python3 train_judge.py --data data --out runs/judge-v0 \ + --micro-batch 64 --grad-accum 2 --eval-rows 2000 --keep-checkpoints 0 +``` + +Measured, so you don't rediscover it: + +| model | micro-batch | peak VRAM | outcome | +|---|---|---|---| +| Gemma-2-2B | 128 | > 47 GB | **OOM** on a 48 GB A6000 | +| Gemma-2-2B | 64 | 28.6 GB | fine, ~50 min | +| Llama-3.1-8B | 64 | 43.8 GB | fine, ~50 min on an A100 80 GB | + +Gemma-2's 256k vocabulary is what makes it expensive — the logits tensor alone is +~13.7 GB at micro-batch 128 (bf16 + fp32 upcast + gradient). Llama's 128k vocab +halves that. + +Keep `--grad-accum` × `--micro-batch` = 128 so runs stay comparable: gradient +accumulation is mathematically equivalent at a fixed effective batch, since each +example contributes exactly one unmasked loss token. + +`--keep-checkpoints 0` retains every checkpoint (10 per run at the default +`--eval-every 200`), which the interpretability work needs for trajectories. + +### Ablation sweep + +`sweep.sh` reproduces the table above, one variable per run, seed fixed: + +```bash +MODEL=meta-llama/Llama-3.1-8B OUT=runs/sweep-llama MB=64 GA=2 \ + SKIP_RUNS=baseline bash sweep.sh run +bash sweep.sh summary +``` + +Finished runs are skipped, so it is safe to re-invoke after a crash. + +--- + +## Evaluation beyond accuracy + +`eval_judge.py` reports what training-time accuracy cannot — confusion matrix, +precision (training reports recall only), the false-equivalence rate, the +`equivalent` tier split, symmetry, and best-of-n ranking: + +```bash +python3 eval_judge.py --adapter runs/judge-v0/adapter-final --data data +``` + +Findings for judge-v0 (Gemma): + +- **False-equivalence rate 0.0077** — the dangerous error, where a drifted + candidate is accepted. Missed-equivalence runs at 0.044, so the judge errs + toward caution by **5.7 : 1**. That is the right asymmetry for a checker. +- **`equivalent` is 50% trivial.** Same-node rows (the identical law re-rendered) + score **1.000**; cross-node rows score **0.912**. The reported 0.956 is their + average — **quote 0.912** for genuine equivalence detection. +- **Dominant failure is over-predicting `incomparable`**: recall 0.977 but + precision 0.899, absorbing 69% of all errors. +- **Symmetry 0.978** — `label(A,B)=weaker` implies `label(B,A)=stronger`. + +--- + +## Caveats + +- **Single seed.** No variance estimate on any number here. +- **Calibration is mediocre** — 39.7% of errors are held above 0.9 confidence, so + a confidence threshold will not cleanly filter dangerous errors. +- **The best-of-n exploitation rate is not yet valid.** The current test uses + same-node faithful candidates, which the judge never misses (1.000). It must be + rebuilt with cross-node candidates (0.912) before being quoted. +- **`classes_val` is only 27% both-novel**; 73% pair one novel law with a + familiar one, so the memorization gap is a *lower bound*. +- **Out-of-fragment behaviour is untrained and unknown.** Training contains only + the 4,694 catalogued laws. Order > 4 identities, multi-operation expressions, + and chained equalities never appear. +- **Layer results show where adaptation is *needed*, not where computation + *lives*.** A frozen late layer may still be doing essential work. + +--- + +## Files + +| file | purpose | +|---|---| +| `judge.py` | library + CLI (`compare`, `batch`, `rank`) | +| `train_judge.py` | training; LoRA rank / layer / module flags | +| `build_pair_dataset.py` | builds the splits from `oracle/`'s matrix | +| `eval_judge.py` | post-hoc evaluation | +| `sweep.sh` | ablation sweep driver | +| `infra/` | optional GPU-rental helpers — see `infra/INFRA.md` | diff --git a/judge/README.md b/judge/README.md index d6c8dfb..b1e3bf7 100644 --- a/judge/README.md +++ b/judge/README.md @@ -1,275 +1,96 @@ -# `judge` — a trained relation judge for the ETP fragment +# ETP relation judge: Python API and CLI -Classifies the logical relation between two equational laws: +Load a trained adapter once, then score equation pairs from another Python program. +The output describes **B relative to A**: equivalent, weaker, stronger, or incomparable. -``` -equivalent | weaker | stronger | incomparable -``` - -The label describes **B relative to A**, matching `oracle/`'s convention -(`weaker` = B drops constraints). - -This replaces a stack of prompted LLMs with a single measurable artifact: known -training data, a known split, and weights we own and can probe. It is trained on -the ETP implication matrix, where every label is Lean-verified rather than judged. - -**No GPU required to use it.** One comparison is a couple of seconds on CPU; -CUDA and Apple-Silicon MPS are used automatically when available. - ---- - -## Quickstart - -```bash -pip install torch transformers peft -huggingface-cli login # Gemma is gated: accept the license first - -export JUDGE_ADAPTER=/path/to/judge-v0/adapter-final -python3 judge.py compare "x * y = y * x" "a * b = b * a" -``` - -``` -A: x * y = y * x -B: a * b = b * a - - equivalent (confidence 1.000) - - equivalent 1.000 ######################################## - weaker 0.000 - stronger 0.000 - incomparable 0.000 -``` - -First run downloads the ~5 GB base model. The adapter itself is ~113 MB. - -### Commands - -| command | purpose | -|---|---| -| `compare A B` | score one pair; `--json` for machine output | -| `batch pairs.jsonl` | score a JSONL of `{text_a, text_b}` rows; reports accuracy if `label` present | -| `rank intended candidates.txt` | order candidates by `P(equivalent)` — what best-of-n selection actually does | +## Python interface -`rank` is the one to reach for when using the judge as a *checker*: - -``` -$ python3 judge.py rank "x * y = y * x" candidates.txt -intended: x * y = y * x - - 1. p(equiv)=0.9997 [equivalent ] p * q = q * p - 2. p(equiv)=0.9996 [equivalent ] a * b = b * a - 3. p(equiv)=0.0383 [weaker ] x = x - 4. p(equiv)=0.0017 [incomparable] x * (y * z) = (x * y) * z -``` - -### As a library - -Downstream interpretability tooling should import rather than shell out: +Run from the repository root: ```python from judge import Judge -j = Judge("path/to/adapter-final") # base model read from adapter_config.json -v = j.compare("x * y = y * x", "a * b = b * a") -v.label, v.confidence, v.probs +judge = Judge("judge/runs/judge-v0/adapter-final") +verdict = judge("x * y = y * x", "a * b = b * a") +print(verdict.label, verdict.confidence) +print(verdict.to_dict()) -j.probs([(a, b), ...]) # (n, 4) array, no argmax applied -j.rank(intended, candidates) # [(candidate, p_equivalent, pred), ...] -j.model # the live PeftModel, for hooks/probing +results = judge.compare_many([ + ("x * y = y * x", "a * b = b * a"), + ("x * y = y * x", "x = x"), +], batch_size=2) ``` -`Judge.model` is exposed deliberately: activation capture, SAE work, and -ablations attach there. - ---- - -## How it works - -Prompt is `"A: {eq}\nB: {eq}\nRelation:"`. The label is a **single next token** -with loss masked to that token alone. At inference the logits are **restricted to -the four label tokens** before the argmax, so the judge cannot emit an invalid -answer — a wrong prediction is always one of the other three, never gibberish. - -`unknown` is not a class. The 190 unresolved matrix cells are never sampled, so -the judge cannot abstain. If you need abstention, threshold on `confidence` — -but note calibration is mediocre (see Caveats). - ---- +`compare(a, b)` and `judge(a, b)` are equivalent. +`compare_many` preserves input order. `probs` returns an `(n, 4)` NumPy array. +Its label order is `equivalent`, `weaker`, `stronger`, `incomparable`. +Empty inputs return an empty list or `(0, 4)` array. +`rank(intended, candidates)` sorts candidates by their equivalence score. -## Training data +Use a Llama adapter directory to select the Llama backbone. +The loader reads `base_model_name_or_path` from the adapter configuration. +It supports CPU, CUDA, and MPS. Runtime and memory depend on the model and hardware. -Built by `build_pair_dataset.py` from `oracle/`'s matrix. Labels are two byte -reads from the Lean-verified implication matrix — **no prover runs at training -time**; the certification is a frozen artifact of the -`2024-11-10-outcomes` snapshot. +## Installation and artifacts ```bash -cd judge -python3 build_pair_dataset.py --out data/ +python -m pip install -r judge/requirements.txt -r pipeline/requirements.txt ``` -Requires `oracle/data/matrix.bin` — see `oracle/README.md` to build it. -The generated `data/` is gitignored (39 MB, fully regenerable). +Base weights and trained adapters are not committed. +Obtain an adapter directory from the experiment owner or train one using the pipeline. +A usable adapter includes `adapter_config.json` and `adapter_model.safetensors`. +Base models may require Hugging Face authentication and acceptance of their licenses. +First load downloads uncached base weights. -### The split is the important part - -**By equivalence class, never by pair.** 1,415 classes → **1,064 train / 140 val -/ 211 test**. Laws inside a class are mutually derivable, so a random pair split -puts the same law on both sides and a lookup table scores perfectly. - -The `x = y` class alone is **1,496 laws — 31.9% of the catalogue** — and is -*pinned into training*. A random split could otherwise have removed a third of -all laws from training. - -Four splits, and the contrast between two of them is the measurement: - -| split | composition | measures | -|---|---|---| -| `train` | 240,000 pairs from train classes | — | -| `pairs_val` | new pairings of **familiar** laws | interpolation | -| `classes_val` | ≥1 law from held-out classes | novel theories | -| `test` | ≥1 law from test classes, never selected on | final | - -**`pairs_val − classes_val` is the memorization measurement.** Everything else is -held constant; only law familiarity varies. A pair-level memoriser fails -`pairs_val`; a law-level lookup passes it and fails `classes_val`. Only a -relational model passes both. - ---- - -## Results - -Trained models, `test_acc` on 211 held-out equivalence classes (chance = 0.25): - -| config | Gemma-2-2B | Llama-3.1-8B | -|---|---|---| -| **baseline** (r16, all layers, all modules) | **0.9619** | **0.9649** | -| rank 4 | 0.9566 | 0.9596 | -| rank 64 | 0.9625 | 0.9647 | -| early layers only | 0.9414 | 0.9619 | -| middle layers only | 0.9494 | 0.9457 | -| **late layers only** | **0.8625** | **0.8491** | -| attention only | 0.9523 | 0.9629 | -| MLP only | 0.9601 | 0.9631 | - -Memorization gap: **1.91 pp** (Gemma), **1.25 pp** (Llama). Both trained a judge, -not a lookup table. - -**Neither base model can be contaminated.** ETP launched 2024-09-25; Gemma-2 was -released 2024-06-27 and Llama-3.1 on 2024-07-23. Both predate the project's -existence, which is a stronger guarantee than any stated data cutoff. - -### What replicates across both families - -- **Rank saturates at or below 16.** Quadrupling the adapter moves accuracy - +0.06 pp / −0.01 pp. Cutting to r=4 costs exactly −0.53 pp on both. -- **Late-only collapses** — 5–6× worse than any other layer choice. The relation - is built in the front two-thirds, not decided at the end. Note `late` has - *more* adapted layers than `early` (10 vs 8 on Gemma), so this is positional, - not a capacity effect. -- **MLP-only is nearly free** — −0.18 pp in both. - -### What does NOT replicate — do not claim - -- **The attention/MLP asymmetry is a Gemma artifact.** Gemma: attention-only cost - 5× more than MLP-only. Llama: they are equivalent. "The computation lives in - the MLP" does not survive replication. -- **Early-vs-middle ordering flips** between families. - ---- - -## Retraining - -Any CUDA box with **≥ 40 GB** works. Memory is dominated by the vocabulary-sized -logits plus retained activations for every adapted layer. +## Command line ```bash -python3 train_judge.py --data data --out runs/judge-v0 \ - --micro-batch 64 --grad-accum 2 --eval-rows 2000 --keep-checkpoints 0 +python -m judge compare "x * y = y * x" "x = x" --adapter judge/runs/judge-v0/adapter-final --json +python -m judge batch pairs.jsonl --adapter judge/runs/judge-v0/adapter-final --out scored.jsonl +python -m judge rank intended.txt candidates.txt --adapter judge/runs/judge-v0/adapter-final ``` -Measured, so you don't rediscover it: - -| model | micro-batch | peak VRAM | outcome | -|---|---|---|---| -| Gemma-2-2B | 128 | > 47 GB | **OOM** on a 48 GB A6000 | -| Gemma-2-2B | 64 | 28.6 GB | fine, ~50 min | -| Llama-3.1-8B | 64 | 43.8 GB | fine, ~50 min on an A100 80 GB | +The original `python judge/judge.py ...` invocation remains supported. -Gemma-2's 256k vocabulary is what makes it expensive — the logits tensor alone is -~13.7 GB at micro-batch 128 (bf16 + fp32 upcast + gradient). Llama's 128k vocab -halves that. +## Training and evaluation -Keep `--grad-accum` × `--micro-batch` = 128 so runs stay comparable: gradient -accumulation is mathematically equivalent at a fixed effective batch, since each -example contributes exactly one unmasked loss token. +See [the configurable pipeline](../pipeline/README.md). +It exposes `run(config, stage=..., dry_run=...)` and the equivalent CLI. +The original scripts remain available for data building, training, detailed evaluation, and LoRA sweeps. -`--keep-checkpoints 0` retains every checkpoint (10 per run at the default -`--eval-every 200`), which the interpretability work needs for trajectories. +New training saves `prompt_template.json` with the final adapter and each training checkpoint. +It also saves checkpoint tokenizers, resolved configuration, final metrics, and training-state logs. +Older adapters without prompt metadata use `A: {a}\nB: {b}\nRelation:`. +Change representations during data generation, then train an adapter for the chosen input format. +The inference API does not translate arbitrary Lean or Python inputs into infix. -### Ablation sweep +## Interpretability and limitations -`sweep.sh` reproduces the table above, one variable per run, seed fixed: - -```bash -MODEL=meta-llama/Llama-3.1-8B OUT=runs/sweep-llama MB=64 GA=2 \ - SKIP_RUNS=baseline bash sweep.sh run -bash sweep.sh summary -``` +`judge.model` exposes the live PEFT model for hooks and interventions. +`judge.tok` exposes the tokenizer. `judge.format_prompt(a, b)` uses the saved template. -Finished runs are skipped, so it is safe to re-invoke after a crash. +These models predict a relation token; they do not generate chain-of-thought text. +Training uses the distinct first token of each label, with loss masked to the answer position. +Inference normalizes scores over those four label tokens. +Confidence is not a calibrated correctness guarantee. There is no unknown class or automatic oracle check. ---- +The existing sampler balances label counts, not syntax-label associations. +It permits repeated underlying pairs and includes same-law equivalences. +Class-held-out evaluation requires at least one novel endpoint, not necessarily both. +These limitations prevent interpreting accuracy alone as proof of semantic generalization. +Corrected sampling and causal interpretability experiments remain separate research work. -## Evaluation beyond accuracy +The [historical report](HISTORICAL_RESULTS.md) preserves the original numbers with explicit qualifications. +This interface update does not establish new model accuracy, calibration, or causal mechanisms. -`eval_judge.py` reports what training-time accuracy cannot — confusion matrix, -precision (training reports recall only), the false-equivalence rate, the -`equivalent` tier split, symmetry, and best-of-n ranking: +## Verification ```bash -python3 eval_judge.py --adapter runs/judge-v0/adapter-final --data data +python -m unittest discover -s tests -v +python -m pipeline.run --config configs/judge-v0-baseline.yaml --stage all --dry-run ``` -Findings for judge-v0 (Gemma): - -- **False-equivalence rate 0.0077** — the dangerous error, where a drifted - candidate is accepted. Missed-equivalence runs at 0.044, so the judge errs - toward caution by **5.7 : 1**. That is the right asymmetry for a checker. -- **`equivalent` is 50% trivial.** Same-node rows (the identical law re-rendered) - score **1.000**; cross-node rows score **0.912**. The reported 0.956 is their - average — **quote 0.912** for genuine equivalence detection. -- **Dominant failure is over-predicting `incomparable`**: recall 0.977 but - precision 0.899, absorbing 69% of all errors. -- **Symmetry 0.978** — `label(A,B)=weaker` implies `label(B,A)=stronger`. - ---- - -## Caveats - -- **Single seed.** No variance estimate on any number here. -- **Calibration is mediocre** — 39.7% of errors are held above 0.9 confidence, so - a confidence threshold will not cleanly filter dangerous errors. -- **The best-of-n exploitation rate is not yet valid.** The current test uses - same-node faithful candidates, which the judge never misses (1.000). It must be - rebuilt with cross-node candidates (0.912) before being quoted. -- **`classes_val` is only 27% both-novel**; 73% pair one novel law with a - familiar one, so the memorization gap is a *lower bound*. -- **Out-of-fragment behaviour is untrained and unknown.** Training contains only - the 4,694 catalogued laws. Order > 4 identities, multi-operation expressions, - and chained equalities never appear. -- **Layer results show where adaptation is *needed*, not where computation - *lives*.** A frozen late layer may still be doing essential work. - ---- - -## Files - -| file | purpose | -|---|---| -| `judge.py` | library + CLI (`compare`, `batch`, `rank`) | -| `train_judge.py` | training; LoRA rank / layer / module flags | -| `build_pair_dataset.py` | builds the splits from `oracle/`'s matrix | -| `eval_judge.py` | post-hoc evaluation | -| `sweep.sh` | ablation sweep driver | -| `infra/` | optional GPU-rental helpers — see `infra/INFRA.md` | +Tests cover synthetic scoring, batching, representation controls, configuration, metrics, and failure propagation. +They do not validate the trained Gemma or Llama weights. +A real-model inference smoke test and a GPU training smoke test remain outstanding. diff --git a/judge/__init__.py b/judge/__init__.py new file mode 100644 index 0000000..0feaebb --- /dev/null +++ b/judge/__init__.py @@ -0,0 +1,5 @@ +"""Load an ETP judge once, then call it on equation pairs.""" + +from .judge import Judge, LABELS, Verdict + +__all__ = ["Judge", "LABELS", "Verdict"] diff --git a/judge/__main__.py b/judge/__main__.py new file mode 100644 index 0000000..6dfca84 --- /dev/null +++ b/judge/__main__.py @@ -0,0 +1,4 @@ +from .judge import main + +if __name__ == "__main__": + main() diff --git a/judge/build_pair_dataset.py b/judge/build_pair_dataset.py index 22c923e..f505f3f 100644 --- a/judge/build_pair_dataset.py +++ b/judge/build_pair_dataset.py @@ -33,7 +33,8 @@ class would otherwise remove a third of the catalogue from training) import sys from collections import Counter, defaultdict -ORACLE_DIR = os.path.expanduser("~/Documents/semantic-diffchecking-repo/oracle") +ORACLE_DIR = os.environ.get("SDC_ORACLE", os.path.join( + os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "oracle")) sys.path.insert(0, ORACLE_DIR) os.environ.setdefault("ETP_EQUATIONS", os.path.join(ORACLE_DIR, "data", "equations.txt")) @@ -64,7 +65,16 @@ def render(term, mapping, op, top=True): def augment(eq: Equation, rng: random.Random) -> str: - """Render a parsed law with random variable names, op symbol, orientation.""" + """Render a law using the configured representation and surface choices.""" + # Preserve the original RNG sequence when using its default infix recipe. + kind = os.environ.get("SDC_RENDER_KIND", "infix") + rename = os.environ.get("SDC_RENDER_RENAME", "True") == "True" + flip = os.environ.get("SDC_RENDER_FLIP", "True") == "True" + symbol = os.environ.get("SDC_RENDER_OP") or None + if kind != "infix" or not rename or not flip or symbol is not None: + sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) + from pipeline.representations import render as render_view + return render_view(eq, rng, kind=kind, op=symbol, rename=rename, flip=flip) varnames: list = [] eq.lhs.variables(varnames) eq.rhs.variables(varnames) diff --git a/judge/judge.py b/judge/judge.py index 2cdcf41..8fee92f 100755 --- a/judge/judge.py +++ b/judge/judge.py @@ -8,7 +8,7 @@ The label describes B relative to A, matching the oracle's convention (`weaker` = B drops constraints). Scoring is a single forward pass whose logits are restricted to the four label tokens, so the judge cannot emit an -invalid answer and every prediction carries a calibrated probability. +invalid label. Scores are normalized over four labels, not calibrated guarantees. CLI --- @@ -24,8 +24,7 @@ j.compare("x * y = y * x", "a * b = b * a") # -> Verdict j.probs([(a, b), ...]) # -> (n, 4) ndarray -No GPU required. A single comparison takes a couple of seconds on CPU; -CUDA and Apple-Silicon MPS are used automatically when present. +CPU, CUDA, and Apple-Silicon MPS are supported. Runtime depends on the hardware. """ from __future__ import annotations @@ -37,6 +36,7 @@ LABELS = ("equivalent", "weaker", "stronger", "incomparable") DEFAULT_BASE = "google/gemma-2-2b" +DEFAULT_TEMPLATE = "A: {a}\nB: {b}\nRelation:" @dataclass @@ -50,6 +50,10 @@ class Verdict: def __str__(self) -> str: return f"{self.label} ({self.confidence:.3f})" + def to_dict(self) -> dict: + """Return JSON-compatible fields for a notebook or another pipeline.""" + return asdict(self) + class Judge: """A loaded relation judge. Construct once, score many times.""" @@ -101,19 +105,46 @@ def __init__(self, adapter: str, base: str | None = None, base, torch_dtype=td, attn_implementation=attn) self.model = PeftModel.from_pretrained(model, adapter).to(device).eval() + # An adapter records the prompt it was trained with. Serving it under a + # different template loads fine, answers fine, and is wrong. Older + # adapters predate the file and used the default, which is what the + # fallback assumes. + self.template = DEFAULT_TEMPLATE + tpl = os.path.join(adapter, "prompt_template.json") + if os.path.isfile(tpl): + with open(tpl, encoding="utf-8") as f: + self.template = json.load(f).get("template", DEFAULT_TEMPLATE) + if self.template != DEFAULT_TEMPLATE: + print(f"note: this adapter uses a non-default prompt: " + f"{self.template!r}", file=sys.stderr) + @staticmethod def prompt(text_a: str, text_b: str) -> str: - return f"A: {text_a}\nB: {text_b}\nRelation:" + """The DEFAULT prompt. Use `j.format_prompt` to honour an adapter's own + recorded template; this static form is kept for callers that build + prompts without a loaded Judge.""" + return DEFAULT_TEMPLATE.format(a=text_a, b=text_b) + + def format_prompt(self, text_a: str, text_b: str) -> str: + return self.template.format(a=text_a, b=text_b) def probs(self, pairs, batch_size: int = 32): """pairs: iterable of (text_a, text_b). Returns an (n, 4) float array.""" + if not isinstance(batch_size, int) or isinstance(batch_size, bool) or batch_size < 1: + raise ValueError("batch_size must be a positive integer") torch = self.torch pairs = list(pairs) + for pair in pairs: + if not isinstance(pair, (tuple, list)) or len(pair) != 2: + raise ValueError("each pair must contain two equation strings") + if any(not isinstance(text, str) or not text.strip() for text in pair): + raise ValueError("equations must be non-empty strings") out = [] with torch.no_grad(): for i in range(0, len(pairs), batch_size): chunk = pairs[i:i + batch_size] - enc = [self.tok(self.prompt(a, b), add_special_tokens=True).input_ids + enc = [self.tok(self.format_prompt(a, b), + add_special_tokens=True).input_ids for a, b in chunk] m = max(len(e) for e in enc) pad = self.tok.pad_token_id @@ -130,14 +161,28 @@ def probs(self, pairs, batch_size: int = 32): # answer impossible rather than merely unlikely. lg = lg[torch.arange(lg.size(0), device=self.device), last][:, self.label_ids] out.append(torch.softmax(lg.float(), -1).cpu()) - return torch.cat(out).numpy() if out else None + return (torch.cat(out).numpy() if out + else torch.empty((0, len(LABELS)), dtype=torch.float32).numpy()) def compare(self, text_a: str, text_b: str) -> Verdict: - p = self.probs([(text_a, text_b)])[0] - k = int(p.argmax()) - return Verdict(label=LABELS[k], confidence=float(p[k]), - probs={l: float(v) for l, v in zip(LABELS, p)}, - text_a=text_a, text_b=text_b) + return self.compare_many([(text_a, text_b)])[0] + + def __call__(self, text_a: str, text_b: str) -> Verdict: + """Equivalent to compare(A, B); reuses the already loaded model.""" + return self.compare(text_a, text_b) + + def compare_many(self, pairs, batch_size: int = 32) -> list[Verdict]: + """Score pairs in batches and return verdicts in input order.""" + pairs = list(pairs) + probabilities = self.probs(pairs, batch_size=batch_size) + results = [] + for (a, b), p in zip(pairs, probabilities): + k = int(p.argmax()) + results.append(Verdict( + label=LABELS[k], confidence=float(p[k]), + probs={label: float(value) for label, value in zip(LABELS, p)}, + text_a=a, text_b=b)) + return results def rank(self, intended: str, candidates): """Order candidates by P(equivalent to `intended`), best first. diff --git a/judge/requirements.txt b/judge/requirements.txt new file mode 100644 index 0000000..5527b8b --- /dev/null +++ b/judge/requirements.txt @@ -0,0 +1,6 @@ +torch +transformers +peft +accelerate +datasets +numpy diff --git a/judge/train_judge.py b/judge/train_judge.py index b0784e3..407344c 100644 --- a/judge/train_judge.py +++ b/judge/train_judge.py @@ -22,6 +22,14 @@ import os +def save_prompt_metadata(directory, template): + """Keep the inference prompt beside every saved adapter.""" + os.makedirs(directory, exist_ok=True) + with open(os.path.join(directory, "prompt_template.json"), "w", encoding="utf-8") as sink: + json.dump({"template": template, + "labels": ["equivalent", "weaker", "stronger", "incomparable"]}, sink, indent=2) + + def main() -> None: ap = argparse.ArgumentParser() ap.add_argument("--data", default="data") @@ -58,7 +66,7 @@ def main() -> None: from datasets import load_dataset from peft import LoraConfig, get_peft_model from transformers import (AutoModelForCausalLM, AutoTokenizer, - Trainer, TrainingArguments, set_seed) + Trainer, TrainerCallback, TrainingArguments, set_seed) set_seed(args.seed) LABELS = ["equivalent", "weaker", "stronger", "incomparable"] @@ -82,8 +90,14 @@ def main() -> None: raise SystemExit(f"label first-tokens collide under {args.model}: {collisions}\n" "pick single-token label synonyms before training.") + # The prompt is a research variable, so it comes from the config via the + # environment. It MUST match judge.py at inference; both default to the + # same string, and a mismatch would be silent. + TEMPLATE = os.environ.get("SDC_PROMPT_TEMPLATE", "A: {a}\nB: {b}\nRelation:") + print(f"prompt template: {TEMPLATE!r}") + def fmt(row): - prompt = f"A: {row['text_a']}\nB: {row['text_b']}\nRelation:" + prompt = TEMPLATE.format(a=row["text_a"], b=row["text_b"]) ids = tok(prompt, add_special_tokens=True).input_ids input_ids = ids + [label_ids[row["label"]]] labels = [-100] * len(ids) + [label_ids[row["label"]]] @@ -202,12 +216,19 @@ def preprocess_logits(logits, labels): bf16=True, report_to="none", seed=args.seed, dataloader_num_workers=2, remove_unused_columns=False, ) + class SavePrompt(TrainerCallback): + def on_save(self, args, state, control, **kwargs): + directory = os.path.join(args.output_dir, f"checkpoint-{state.global_step}") + save_prompt_metadata(directory, TEMPLATE) + tok.save_pretrained(directory) + trainer = Trainer( model=model, args=targs, train_dataset=ds["train"], eval_dataset={"pairs": curve["pairs_val"], "classes": curve["classes_val"]}, data_collator=collate, compute_metrics=compute_metrics, preprocess_logits_for_metrics=preprocess_logits, + callbacks=[SavePrompt()], ) trainer.train() @@ -220,7 +241,9 @@ def preprocess_logits(logits, labels): adapter_dir = os.path.join(args.out, "adapter-final") trainer.save_model(adapter_dir) tok.save_pretrained(adapter_dir) + save_prompt_metadata(adapter_dir, TEMPLATE) print(f"saved adapter + tokenizer to {adapter_dir}") + print(f"recorded prompt template: {TEMPLATE!r}") with open(os.path.join(args.out, "run_config.json"), "w") as f: json.dump(vars(args), f, indent=2, default=str) diff --git a/pipeline/README.md b/pipeline/README.md new file mode 100644 index 0000000..92022a4 --- /dev/null +++ b/pipeline/README.md @@ -0,0 +1,80 @@ +# Configurable judge pipeline + +One YAML or JSON configuration controls data generation, training, and evaluation. +Call the same workflow from Python or the command line. + +```python +from pipeline.run import run + +plan = run("configs/judge-v0-baseline.yaml", stage="all", dry_run=True) +print(plan["adapter_dir"]) +``` + +`dry_run=True` prints planned commands without building data or launching training. +Remove it only when the data, dependencies, and compute are prepared. + +```bash +python -m pipeline.run --config configs/rep-lean.yaml --stage all --dry-run +``` + +## Configuration + +The defaults are defined in `config.py`. Unknown keys and unsupported metric names fail before work starts. + +| Section | Controls | +|---|---| +| `data` | Output directory, label counts, class split fractions, pinned-class threshold | +| `render` | Infix, Lean, LaTeX, Python lambda, or function-call equation notation; renaming and equality reversal | +| `prompt` | Exact training/inference template; fixed label identifiers | +| `model` | Base checkpoint and LoRA rank, layer placement, and modules | +| `train` | Learning rate, batches, steps, evaluation interval, checkpoint retention, output directory | +| `eval` | Dataset splits and selected accuracy, error-rate, per-label, and per-tier metrics | + +`function_call` is an equation representation, separate from the Python API. +The current public pipeline supports the original sampler only. +Shape-balanced data construction and shape-specific diagnostics from the research workspace are not included in this interface update. +It does not claim to reproduce the later v2/v3 experiments. + +Relative artifact paths resolve against this repository root, independent of the caller's working directory. +Configuration file paths resolve from the caller's working directory. +Data generation defaults to this checkout's `oracle/`; `SDC_ORACLE` can override that directory. +Prepare its catalogue and matrix using [oracle setup](../oracle/README.md). + +## Reuse and checkpoints + +```python +from pipeline.config import load +from pipeline.run import run + +config = load("configs/judge-v0-baseline.yaml") +config.train["out_dir"] = "judge/runs/new-run" +# run(config, stage="train") # Executes actual training when enabled. +``` + +Stages are `data`, `train`, `eval`, and `all`. +The function revalidates Config objects and returns artifact paths after the requested stages succeed. +Failures propagate to the caller. Evaluation failures cannot silently become successful summaries. + +The data stage reuses a complete dataset only when its saved pipeline recipe matches. +Changing representation or sampling settings requires a fresh data directory. +For pre-existing datasets without a recipe, run only `train` or `eval` explicitly. +The recipe records configuration, not a cryptographic identity of the underlying oracle files. +Changing oracle artifacts requires a fresh data directory. + +Training refuses to overwrite an existing final adapter or its saved configuration. +`resolved_config.yaml` can be loaded again with `load()`. +Intermediate checkpoints and the final adapter carry their prompt and tokenizer. +`keep_checkpoints=0` retains all training checkpoints; positive values limit retained checkpoints. + +## Metrics + +`eval_summary.json` contains the selected metrics and sample counts for each requested split. + +- False equivalence: non-equivalent examples predicted equivalent, divided by all non-equivalent examples. +- Missed equivalence: equivalent examples predicted non-equivalent, divided by all equivalent examples. +- Per-label metrics: precision, recall, gold support, and prediction count. +- Per-tier metrics: separate summaries for the dataset's tier labels. + +Undefined rates are `null`, not zero. Every scored row needs a resolved oracle label. +The pipeline also writes the scored JSONL rows for independent analysis. +It does not treat judge predictions as verified labels. diff --git a/pipeline/__init__.py b/pipeline/__init__.py new file mode 100644 index 0000000..9d99c03 --- /dev/null +++ b/pipeline/__init__.py @@ -0,0 +1,4 @@ +"""Configuration-driven judge training and evaluation. + +Import ``run`` from ``pipeline.run`` and ``load`` from ``pipeline.config``. +""" diff --git a/pipeline/config.py b/pipeline/config.py new file mode 100644 index 0000000..373fe53 --- /dev/null +++ b/pipeline/config.py @@ -0,0 +1,241 @@ +#!/usr/bin/env python3 +"""One config file describes a whole run, and the run records the config back. + +Before this, a run was defined by a command line: which flags were passed to +`build_pair_dataset.py`, then which flags were passed to `train_judge.py`, then +whatever was typed to evaluate it. Nothing tied the three together, and a +finished adapter did not carry the sampler settings that produced its data. Two +of this project's real bugs were of exactly that shape -- a flag that silently +did nothing, and a flag that did the opposite of its name -- and neither was +visible from the artifacts afterwards. + +So: every stage reads one file, and the resolved config is written into the run +directory as `resolved_config.yaml`. If you have the run, you have the recipe. + +Sections +-------- + data which pairs get built, and how they are sampled + render how an equation is written down (see representations.py) + prompt the exact string the model is shown + model base model and LoRA shape + train optimisation settings + eval what gets measured afterwards + +Anything absent falls back to the defaults below, which reproduce judge-v0. + +Usage +----- + from config import load + cfg = load("configs/baseline.yaml") + cfg.render.kind # "infix" + cfg.dump(run_dir) # writes resolved_config.yaml +""" +from __future__ import annotations + +import copy +import json +import os +from dataclasses import dataclass, field, asdict, fields + +DEFAULTS = { + "name": "unnamed", + "seed": 0, + "data": { + "sampler": "original", # existing label-balanced sampler + "out_dir": "judge/data", + # original sampler + "per_class_train": 60000, + "per_class_eval": 2000, + # split, shared + "pin_size": 50, + "test_frac": 0.15, + "val_frac": 0.10, + }, + "render": { + "kind": "infix", # infix | function_call | lean | latex | python + "op": None, # None -> draw at random per row + "rename": True, + "flip": True, + }, + "prompt": { + # {a} and {b} are the two rendered equations. The trailing text is what + # the answer token follows, so it must not end in a space: the space + # belongs to the label token itself (" equivalent", not "equivalent"). + "template": "A: {a}\nB: {b}\nRelation:", + "labels": ["equivalent", "weaker", "stronger", "incomparable"], + }, + "model": { + "base": "google/gemma-2-2b", + "lora_rank": 16, + "lora_layers": "all", # all | early | middle | late + "lora_modules": "all", # all | attn | mlp + }, + "train": { + "lr": 1e-4, + "micro_batch": 32, + "grad_accum": 4, # micro_batch * grad_accum = 128, keep it + "max_steps": 2000, + "eval_every": 200, + "eval_rows": 2000, + "keep_checkpoints": 0, # 0 = keep every one + "out_dir": "judge/runs/unnamed", + }, + "eval": { + "splits": ["pairs_val", "classes_val", "test"], + "metrics": ["accuracy", "false_equivalence", "missed_equivalence", + "per_label", "per_tier"], + }, +} + + +class Section(dict): + """A dict that also answers to attribute access, so cfg.render.kind works.""" + + def __getattr__(self, k): + try: + return self[k] + except KeyError as e: + raise AttributeError( + f"no config key {k!r} in this section; have {sorted(self)}") from e + + +@dataclass +class Config: + name: str + seed: int + data: Section + render: Section + prompt: Section + model: Section + train: Section + eval: Section + source_path: str = "" + + def dump(self, run_dir: str) -> str: + """Write the fully resolved config beside the run it produced.""" + os.makedirs(run_dir, exist_ok=True) + p = os.path.join(run_dir, "resolved_config.yaml") + with open(p, "w", encoding="utf-8") as f: + f.write(_to_yaml(self.as_dict())) + return p + + def as_dict(self) -> dict: + d = {f.name: getattr(self, f.name) for f in fields(self) if f.name != "source_path"} + return {k: (dict(v) if isinstance(v, Section) else v) for k, v in d.items()} + + def prompt_for(self, a: str, b: str) -> str: + return self.prompt.template.format(a=a, b=b) + + +def _deep_merge(base: dict, over: dict, path: str = "") -> dict: + out = copy.deepcopy(base) + for k, v in (over or {}).items(): + here = f"{path}.{k}" if path else k + if k not in out: + raise KeyError(f"unknown config key {here!r}. " + f"Valid keys here: {sorted(out)}") + if isinstance(v, dict) and isinstance(out[k], dict): + out[k] = _deep_merge(out[k], v, here) + else: + out[k] = v + return out + + +def _to_yaml(d: dict, indent: int = 0) -> str: + pad = " " * indent + lines = [] + for k, v in d.items(): + if isinstance(v, dict): + lines.append(f"{pad}{k}:") + lines.append(_to_yaml(v, indent + 1)) + elif isinstance(v, list): + lines.append(f"{pad}{k}: {json.dumps(v)}") + elif isinstance(v, str): + lines.append(f"{pad}{k}: {json.dumps(v)}") + elif v is None: + lines.append(f"{pad}{k}: null") + else: + lines.append(f"{pad}{k}: {v}") + return "\n".join(l for l in lines if l) + + +def load(path: str) -> Config: + """Read a config, merge it over the defaults, and check it makes sense.""" + try: + import yaml + except ImportError: + with open(path, encoding="utf-8") as source: + raw = json.load(source) + else: + with open(path, encoding="utf-8") as source: + raw = yaml.safe_load(source) or {} + merged = _deep_merge(DEFAULTS, raw) + cfg = Config( + name=merged["name"], seed=merged["seed"], + data=Section(merged["data"]), render=Section(merged["render"]), + prompt=Section(merged["prompt"]), model=Section(merged["model"]), + train=Section(merged["train"]), eval=Section(merged["eval"]), + source_path=os.path.abspath(path)) + validate(cfg) + return cfg + + +def validate(cfg: Config) -> None: + """Fail loudly on the mistakes that would otherwise pass silently.""" + if __package__: + from .representations import RENDERERS + else: + from representations import RENDERERS + + if cfg.render.kind not in RENDERERS: + raise SystemExit(f"render.kind {cfg.render.kind!r} unknown; " + f"have {sorted(RENDERERS)}") + if cfg.data.sampler != "original": + raise SystemExit(f"data.sampler {cfg.data.sampler!r} unknown") + supported_metrics = set(DEFAULTS["eval"]["metrics"]) + if not cfg.eval.metrics or set(cfg.eval.metrics) - supported_metrics: + raise ValueError(f"eval.metrics must select from {sorted(supported_metrics)}") + if not cfg.eval.splits or set(cfg.eval.splits) - {"pairs_val", "classes_val", "test"}: + raise ValueError("eval.splits must select pairs_val, classes_val, or test") + for key in ("rename", "flip"): + if not isinstance(cfg.render[key], bool): + raise ValueError(f"render.{key} must be a boolean") + for section, keys in ((cfg.data, ("per_class_train", "per_class_eval", "pin_size")), + (cfg.train, ("micro_batch", "grad_accum", "max_steps", "eval_every", "eval_rows"))): + for key in keys: + if type(section[key]) is not int or section[key] < 1: + raise ValueError(f"{key} must be a positive integer") + if type(cfg.train.keep_checkpoints) is not int or cfg.train.keep_checkpoints < 0: + raise ValueError("train.keep_checkpoints must be a non-negative integer") + if not (0 < cfg.data.test_frac < 1 and 0 < cfg.data.val_frac < 1 + and cfg.data.test_frac + cfg.data.val_frac < 1): + raise ValueError("test_frac and val_frac must be positive and sum to less than one") + + t = cfg.prompt.template + for tok in ("{a}", "{b}"): + if tok not in t: + raise SystemExit(f"prompt.template must contain {tok}") + if t.endswith(" "): + raise SystemExit( + "prompt.template ends in a space. The label token carries its own " + "leading space (' equivalent'), so a trailing space here shifts " + "the target token and the run trains on the wrong thing.") + try: + t.format(a="A", b="B") + except (KeyError, ValueError, IndexError) as exc: + raise ValueError("prompt.template must format using only {a} and {b}") from exc + if cfg.prompt.labels != DEFAULTS["prompt"]["labels"]: + raise ValueError("prompt.labels must preserve the four trained label identifiers and their order") + + eff = cfg.train.micro_batch * cfg.train.grad_accum + if eff != 128: + print(f"warning: micro_batch x grad_accum = {eff}, not 128. Runs are " + "only comparable to the published ones at an effective batch " + "of 128.") + + +if __name__ == "__main__": + import sys + c = load(sys.argv[1]) + print(_to_yaml(c.as_dict())) + print(f"\nexample prompt:\n{c.prompt_for('x * y = y * x', 'a * b = b * a')}") diff --git a/pipeline/metrics.py b/pipeline/metrics.py new file mode 100644 index 0000000..b0d6959 --- /dev/null +++ b/pipeline/metrics.py @@ -0,0 +1,43 @@ +"""Metrics on oracle-labeled rows, with explicit denominators.""" + +from collections import defaultdict + +LABELS = ("equivalent", "weaker", "stronger", "incomparable") + + +def summarize(rows, metrics): + """Summarize scored JSONL rows. Undefined rates are None, not zero.""" + rows = list(rows) + for row in rows: + if row.get("label") not in LABELS or row.get("pred") not in LABELS: + raise ValueError("evaluation requires resolved gold and predicted relation labels") + n = len(rows) + result = {"n": n} + if "accuracy" in metrics: + result["accuracy"] = sum(r["label"] == r["pred"] for r in rows) / n if n else None + for name, equivalent in (("false_equivalence", False), ("missed_equivalence", True)): + if name in metrics: + eligible = [r for r in rows if (r["label"] == "equivalent") == equivalent] + errors = sum((r["pred"] == "equivalent") != equivalent for r in eligible) + result[name] = {"errors": errors, "eligible": len(eligible), + "rate": errors / len(eligible) if eligible else None} + if "per_label" in metrics: + result["per_label"] = {} + for label in LABELS: + support = sum(r["label"] == label for r in rows) + predicted = sum(r["pred"] == label for r in rows) + correct = sum(r["label"] == r["pred"] == label for r in rows) + result["per_label"][label] = { + "support": support, "predicted": predicted, + "recall": correct / support if support else None, + "precision": correct / predicted if predicted else None, + } + if "per_tier" in metrics: + tiers = defaultdict(list) + for row in rows: + tiers[row.get("tier", "unspecified")].append(row) + result["per_tier"] = { + tier: summarize(group, ["accuracy", "false_equivalence", "missed_equivalence"]) + for tier, group in sorted(tiers.items()) + } + return result diff --git a/pipeline/representations.py b/pipeline/representations.py new file mode 100644 index 0000000..5a97a15 --- /dev/null +++ b/pipeline/representations.py @@ -0,0 +1,211 @@ +#!/usr/bin/env python3 +"""How an equation is written down, as a swappable choice. + +The judge never sees an equation object. It sees a string. Which string is a +research variable, not a detail: the same law can be written as infix magma +notation, as a Lean statement, as LaTeX, or as Python, and whether the judge +understands them equally is an open question the team wants to test. + +Until now the rendering was hard-coded inside `build_pair_dataset.augment()`, +so answering that question meant editing the dataset builder. Each renderer +here takes a parsed `Equation` plus a naming choice and returns a string. + +Every renderer must satisfy two properties, both checked by `selftest()`: + +* **Grouping is explicit.** The magma operation is not associative, so + `x ? y ? z` is ambiguous and must never be emitted. +* **Variables are used exactly as given.** A renderer may rename, but must not + merge two variables into one or invent a new one. Merging is precisely the + drift the judge exists to detect, so a renderer that did it would poison the + labels. + +Adding one: write the function, add it to `RENDERERS`, run `selftest()`. +""" +from __future__ import annotations + +import os +import sys + +sys.path.insert(0, os.environ.get( + "SDC_ORACLE", + os.path.join(os.path.dirname(os.path.dirname(os.path.abspath(__file__))), "oracle"))) + +from normalizer import Equation, Op, Var, parse_equation # noqa: E402 + +# Symbols the oracle's own parser accepts, so anything rendered infix can be +# read back and mapped to a catalogue node. ASCII '.' is NOT among them. +INFIX_SYMBOLS = ["*", "·", "@", "+"] + +VAR_POOLS = [ + list("xyzwuvst"), + list("abcdefgh"), + list("pqrstuvw"), + list("mnkljihg"), +] + + +def _term_infix(t, names, op, top=True): + if isinstance(t, Var): + return names[t.name] + s = f"{_term_infix(t.left, names, op, False)} {op} {_term_infix(t.right, names, op, False)}" + return s if top else f"({s})" + + +def infix(eq: Equation, names: dict, op: str = "*") -> str: + """`x * (y * z) = (x * y) * z` — the original format, and the default.""" + return f"{_term_infix(eq.lhs, names, op)} = {_term_infix(eq.rhs, names, op)}" + + +def _term_prefix(t, names, fn): + if isinstance(t, Var): + return names[t.name] + return f"{fn}({_term_prefix(t.left, names, fn)}, {_term_prefix(t.right, names, fn)})" + + +def function_call(eq: Equation, names: dict, op: str = "op") -> str: + """`op(x, op(y, z)) = op(op(x, y), z)` — grouping carried by brackets alone. + + Useful as a control: it removes infix precedence entirely, so if the judge + behaves differently here the difference is about notation, not structure. + """ + return (f"{_term_prefix(eq.lhs, names, op)} = " + f"{_term_prefix(eq.rhs, names, op)}") + + +def lean(eq: Equation, names: dict, op: str = "∘") -> str: + """A Lean-style theorem statement over a magma. + + Matches how the Equational Theories Project writes its laws, so it is the + representation a formalisation pipeline would actually emit. + """ + vs = sorted(set(names.values())) + binder = " ".join(vs) + body = (f"{_term_infix(eq.lhs, names, op)} = " + f"{_term_infix(eq.rhs, names, op)}") + return f"∀ {binder} : G, {body}" + + +def latex(eq: Equation, names: dict, op: str = r"\circ") -> str: + r"""`x \circ (y \circ z) = (x \circ y) \circ z` inside math delimiters.""" + body = (f"{_term_infix(eq.lhs, names, op)} = " + f"{_term_infix(eq.rhs, names, op)}") + return f"$ {body} $" + + +def python(eq: Equation, names: dict, op: str = "op") -> str: + """`op(x, op(y, z)) == op(op(x, y), z)` — executable-looking, `==` not `=`. + + Deliberately close to `function_call` except for the equality token. The + pair isolates how much the judge leans on the `=` symbol itself, which is + the thing the team suspected was being lost in embedding work. + """ + return (f"{_term_prefix(eq.lhs, names, op)} == " + f"{_term_prefix(eq.rhs, names, op)}") + + +RENDERERS = { + "infix": infix, + "function_call": function_call, + "lean": lean, + "latex": latex, + "python": python, +} + +# Operator symbols that make sense per renderer. Drawing one at random is the +# surface augmentation the original `augment()` did. +DEFAULT_OPS = { + "infix": INFIX_SYMBOLS, + "function_call": ["op", "f", "mul", "star"], + "lean": ["∘", "⋆", "*"], + "latex": [r"\circ", r"\star", r"\cdot"], + "python": ["op", "f", "mul", "star"], +} + + +def render(eq: Equation, rng, kind: str = "infix", op: str | None = None, + rename: bool = True, flip: bool = True) -> str: + """Render one law, drawing the surface choices from `rng`. + + `rename` shuffles the variable names, `flip` may swap the two sides of the + equality. Both are meaning-preserving, and both are what stop the judge + from keying on a particular spelling. Turn them off to render a law the + same way every time. + """ + if kind not in RENDERERS: + raise KeyError(f"unknown representation {kind!r}; " + f"have {sorted(RENDERERS)}") + order: list = [] + eq.lhs.variables(order) + eq.rhs.variables(order) + if rename: + pool = rng.choice(VAR_POOLS)[:] + rng.shuffle(pool) + names = {v: pool[i] for i, v in enumerate(order)} + else: + names = {v: v for v in order} + symbol = op if op is not None else rng.choice(DEFAULT_OPS[kind]) + if flip and rng.random() < 0.5: + eq = Equation(eq.rhs, eq.lhs) + return RENDERERS[kind](eq, names, symbol) + + +def selftest() -> int: + """Check every renderer against the two properties that must never break. + + Both failure modes are silent if unchecked. A renderer that dropped a + bracket would emit an ambiguous term the oracle cannot parse; one that + merged two variables would relabel the pair, and merging variables is + exactly the drift the judge is built to catch. + """ + import random + from collections import Counter + + cases = [ + "x ◇ (y ◇ z) = (x ◇ y) ◇ z", # grouping must survive + "x ◇ x = x ◇ ((y ◇ y) ◇ z)", # a repeated variable must stay repeated + "x = y", # two distinct variables, no operation + ] + bad = 0 + for kind in RENDERERS: + for src in cases: + eq = parse_equation(src) + order: list = [] + eq.lhs.variables(order) + eq.rhs.variables(order) + names = {v: VAR_POOLS[0][i] for i, v in enumerate(order)} + out = RENDERERS[kind](eq, names, DEFAULT_OPS[kind][0]) + + # every variable that went in must come out, the right number of times + want = Counter() + for t in (eq.lhs, eq.rhs): + stack = [t] + while stack: + node = stack.pop() + if isinstance(node, Var): + want[names[node.name]] += 1 + else: + stack += [node.left, node.right] + # Count occurrences in the TERM only. A Lean statement binds its + # variables up front (`forall x y z : G,`), and those mentions are + # correct, not duplicates, so the binder is dropped before counting. + body = out.split(", ", 1)[1] if out.startswith("\u2200") else out + got = Counter() + for tok in body.replace("(", " ").replace(")", " ").replace(",", " ").split(): + if tok in want: + got[tok] += 1 + ok_vars = got == want + + # brackets must balance, and every operation must be bracketed + ok_group = out.count("(") == out.count(")") + + ok = ok_vars and ok_group + bad += not ok + print(f"[{'ok ' if ok else 'FAIL'}] {kind:<14} {out}") + if not ok: + print(f" wanted {dict(want)}, got {dict(got)}") + print("selftest: all passed" if not bad else f"selftest: {bad} FAILURES") + return bad + + +if __name__ == "__main__": + raise SystemExit(1 if selftest() else 0) diff --git a/pipeline/requirements.txt b/pipeline/requirements.txt new file mode 100644 index 0000000..d2d95b1 --- /dev/null +++ b/pipeline/requirements.txt @@ -0,0 +1,4 @@ +PyYAML +numpy + +# Inference and training dependencies are in judge/requirements.txt. diff --git a/pipeline/run.py b/pipeline/run.py new file mode 100644 index 0000000..00d5118 --- /dev/null +++ b/pipeline/run.py @@ -0,0 +1,189 @@ +#!/usr/bin/env python3 +"""Run a stage, or all of them, from one config file. + + python3 pipeline/run.py --config configs/rep-lean.yaml --stage all + python3 pipeline/run.py --config configs/rep-lean.yaml --stage data + python3 pipeline/run.py --config configs/rep-lean.yaml --stage train + python3 pipeline/run.py --config configs/rep-lean.yaml --stage eval + +Stages +------ + data build the four JSONL splits described by cfg.data and cfg.render + train finetune cfg.model on them with cfg.train + eval score configured splits and write the requested metrics + +Each stage writes into the run directory, and `resolved_config.yaml` goes there +first, so the recipe sits beside the artifact rather than in someone's shell +history. + +`--dry-run` prints what each stage would do, builds nothing and trains nothing. +Use it to check a config before spending on a GPU. +""" +from __future__ import annotations + +import argparse +import json +import os +import subprocess +import sys +import time + +HERE = os.path.dirname(os.path.abspath(__file__)) +WORK = os.path.dirname(HERE) +if __package__: + from .config import Config, load, validate + from .metrics import summarize +else: + from config import Config, load, validate + from metrics import summarize + + +def _abs(p: str) -> str: + return p if os.path.isabs(p) else os.path.join(WORK, p) + + +def stage_data(cfg, dry: bool) -> None: + """Build the dataset. Delegates to build_pair_dataset.py rather than + reimplementing it, so the validated sampler stays the only sampler.""" + out = _abs(cfg.data.out_dir) + cmd = [sys.executable, _abs("judge/build_pair_dataset.py"), + "--out", out, "--seed", str(cfg.seed), + "--pin-size", str(cfg.data.pin_size), + "--test-frac", str(cfg.data.test_frac), + "--val-frac", str(cfg.data.val_frac)] + cmd += ["--per-class-train", str(cfg.data.per_class_train), + "--per-class-eval", str(cfg.data.per_class_eval)] + + env = dict(os.environ) + # The renderer is passed through the environment because + # build_pair_dataset.augment() reads it there; anything other than "infix" + # changes every row it writes. + env["SDC_RENDER_KIND"] = cfg.render.kind + env["SDC_RENDER_RENAME"] = str(bool(cfg.render.rename)) + env["SDC_RENDER_FLIP"] = str(bool(cfg.render.flip)) + env.pop("SDC_RENDER_OP", None) + if cfg.render.op: + env["SDC_RENDER_OP"] = str(cfg.render.op) + + print(f"[data] representation={cfg.render.kind} sampler={cfg.data.sampler}") + print("[data] " + " ".join(cmd)) + if dry: + return + recipe = {"data": dict(cfg.data), "render": dict(cfg.render), "seed": cfg.seed} + marker = os.path.join(out, "pipeline_data_config.json") + if os.path.isdir(out) and os.listdir(out): + if os.path.isfile(marker): + with open(marker, encoding="utf-8") as source: + matches = json.load(source) == recipe + if matches and all(os.path.isfile(os.path.join(out, f"{split}.jsonl")) + for split in ("train", "pairs_val", "classes_val", "test")): + print(f"[data] reusing the matching dataset at {out}") + return + raise ValueError(f"{out} contains data with a missing or different recipe; " + "choose a fresh data.out_dir or run only train/eval on existing data") + subprocess.run(cmd, check=True, env=env) + with open(marker, "w", encoding="utf-8") as sink: + json.dump(recipe, sink, indent=2) + + +def stage_train(cfg, dry: bool) -> None: + out = _abs(cfg.train.out_dir) + if not dry and os.path.exists(os.path.join(out, "adapter-final")): + raise SystemExit(f"{out}/adapter-final exists. Refusing to overwrite a " + "trained adapter; move it or change train.out_dir.") + cmd = [sys.executable, _abs("judge/train_judge.py"), + "--data", _abs(cfg.data.out_dir), "--out", out, + "--model", cfg.model.base, "--rank", str(cfg.model.lora_rank), + "--lora-layers", cfg.model.lora_layers, + "--lora-modules", cfg.model.lora_modules, + "--lr", str(cfg.train.lr), + "--micro-batch", str(cfg.train.micro_batch), + "--grad-accum", str(cfg.train.grad_accum), + "--max-steps", str(cfg.train.max_steps), + "--eval-every", str(cfg.train.eval_every), + "--eval-rows", str(cfg.train.eval_rows), + "--keep-checkpoints", str(cfg.train.keep_checkpoints), + "--seed", str(cfg.seed)] + env = dict(os.environ) + env["SDC_PROMPT_TEMPLATE"] = cfg.prompt.template + print(f"[train] {cfg.model.base} rank={cfg.model.lora_rank} " + f"batch={cfg.train.micro_batch}x{cfg.train.grad_accum}") + print("[train] " + " ".join(cmd)) + if dry: + return + subprocess.run(cmd, check=True, env=env) + + +def stage_eval(cfg, dry: bool) -> None: + out = _abs(cfg.train.out_dir) + adapter = os.path.join(out, "adapter-final") + if not dry and not os.path.isdir(adapter): + raise SystemExit(f"no adapter at {adapter}; run --stage train first") + print(f"[eval] splits={cfg.eval.splits} metrics={cfg.eval.metrics}") + if dry: + return + results = {} + for split in cfg.eval.splits: + f = os.path.join(_abs(cfg.data.out_dir), f"{split}.jsonl") + subprocess.run( + [sys.executable, _abs("judge/judge.py"), + "batch", f, "--adapter", adapter, "--batch-size", "64", + "--out", os.path.join(out, f"scored_{split}.jsonl")], + capture_output=True, text=True, check=True) + scored = os.path.join(out, f"scored_{split}.jsonl") + with open(scored, encoding="utf-8") as source: + rows = [json.loads(line) for line in source if line.strip()] + results[split] = summarize(rows, cfg.eval.metrics) + print(f"[eval] {split}: {results[split]}") + with open(os.path.join(out, "eval_summary.json"), "w") as fh: + json.dump(results, fh, indent=2) + + +def run(config: str | os.PathLike | Config, *, stage: str = "all", + dry_run: bool = False) -> dict: + """Run configured stages from Python and return their artifact locations. + + ``config`` is a YAML/JSON path or a Config returned by ``config.load``. + Relative artifact paths resolve against this workspace, as in the CLI. + A dry run validates and prints commands without writing or launching jobs. + Stage failures propagate to the caller; later stages do not run. + """ + if stage not in ("all", "data", "train", "eval"): + raise ValueError("stage must be all, data, train, or eval") + cfg = config if isinstance(config, Config) else load(os.fspath(config)) + validate(cfg) + run_dir = _abs(cfg.train.out_dir) + stages = ["data", "train", "eval"] if stage == "all" else [stage] + if not dry_run and "train" in stages and os.path.exists(os.path.join(run_dir, "adapter-final")): + raise FileExistsError(f"{run_dir}/adapter-final exists; choose a new train.out_dir") + source = cfg.source_path or "in-memory config" + print(f"=== {cfg.name} (from {source})") + if not dry_run: + print("wrote", cfg.dump(run_dir)) + + t0 = time.time() + for name in stages: + print(f"\n--- stage: {name}") + {"data": stage_data, "train": stage_train, "eval": stage_eval}[name](cfg, dry_run) + print(f"\ndone in {(time.time()-t0)/60:.1f} min") + return { + "name": cfg.name, "stages": stages, "dry_run": dry_run, + "run_dir": run_dir, + "data_dir": _abs(cfg.data.out_dir), + "adapter_dir": os.path.join(run_dir, "adapter-final"), + } + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--config", required=True) + ap.add_argument("--stage", default="all", + choices=["all", "data", "train", "eval"]) + ap.add_argument("--dry-run", action="store_true") + args = ap.parse_args() + + run(args.config, stage=args.stage, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/tests/test_callable_interfaces.py b/tests/test_callable_interfaces.py new file mode 100644 index 0000000..59df448 --- /dev/null +++ b/tests/test_callable_interfaces.py @@ -0,0 +1,195 @@ +"""Offline interface tests. Synthetic logits do not validate trained weights.""" + +import contextlib +import io +import json +from pathlib import Path +import subprocess +import tempfile +from types import SimpleNamespace +import unittest +from unittest.mock import patch + +import numpy as np + +from judge import Judge, LABELS +from pipeline.config import load +from pipeline.run import run, stage_eval + + +ROOT = Path(__file__).resolve().parents[1] + + +class Tensor: + """Small NumPy-backed tensor double for exercising the scoring boundary.""" + + def __init__(self, data): + self.data = np.asarray(data) + + def __getitem__(self, key): + if isinstance(key, tuple): + key = tuple(k.data if isinstance(k, Tensor) else k for k in key) + return Tensor(self.data[key]) + + def size(self, axis): + return self.data.shape[axis] + + def float(self): + return self + + def cpu(self): + return self + + def numpy(self): + return self.data + + +def softmax(tensor, axis): + x = tensor.data + exp = np.exp(x - x.max(axis=axis, keepdims=True)) + return Tensor(exp / exp.sum(axis=axis, keepdims=True)) + + +def synthetic_judge(): + """Pad positions favor the wrong label; only true final positions work.""" + judge = Judge.__new__(Judge) + judge.device = "cpu" + judge.template = "FIRST {a}\nSECOND {b}\nDECISION:" + judge.label_ids = [2, 4, 6, 8] + judge.torch = SimpleNamespace( + no_grad=contextlib.nullcontext, + tensor=lambda data, **kwargs: Tensor(data), + arange=lambda n, **kwargs: Tensor(np.arange(n)), + cat=lambda values: Tensor(np.concatenate([v.data for v in values])), + empty=lambda shape, **kwargs: Tensor(np.empty(shape, dtype=np.float32)), + softmax=softmax, float32=np.float32, + ) + + class Tokenizer: + pad_token_id = 0 + + def __call__(self, prompt, **kwargs): + # Encodes the requested winning class and different sequence lengths. + winner = int(prompt.split()[1]) + return SimpleNamespace(input_ids=[winner + 1] * (winner + 2)) + + class Model: + def __call__(self, input_ids, attention_mask): + n, length = input_ids.data.shape + logits = np.zeros((n, length, 10), dtype=float) + logits[:, :, 2] = 100 # Poison all positions except each true last token. + for row in range(n): + last = int(attention_mask.data[row].sum()) - 1 + winner = int(input_ids.data[row, 0]) - 1 + logits[row, last] = 0 + logits[row, last, judge.label_ids[winner]] = 4 + logits[row, last, 9] = 1000 # Ignore non-label vocabulary logits. + return SimpleNamespace(logits=Tensor(logits)) + + judge.tok = Tokenizer() + judge.model = Model() + return judge + + +class JudgeInterfaceTests(unittest.TestCase): + def setUp(self): + self.judge = synthetic_judge() + + def test_batch_uses_last_unpadded_position_and_only_label_tokens(self): + pairs = [(str(i), "x = x") for i in range(4)] + verdicts = self.judge.compare_many(iter(pairs), batch_size=3) + self.assertEqual([v.label for v in verdicts], list(LABELS)) + for (a, b), verdict in zip(pairs, verdicts): + self.assertEqual((verdict.text_a, verdict.text_b), (a, b)) + self.assertAlmostEqual(sum(verdict.probs.values()), 1) + self.assertGreater(verdict.confidence, 0.9) + self.assertEqual(json.loads(json.dumps(verdict.to_dict()))["label"], verdict.label) + + def test_function_call_and_compare_reuse_the_model(self): + model = self.judge.model + self.assertEqual(self.judge("1", "x = x"), self.judge.compare("1", "x = x")) + self.assertIs(self.judge.model, model) + + def test_prompt_preserves_saved_template(self): + self.assertEqual(self.judge.format_prompt("a", "b"), "FIRST a\nSECOND b\nDECISION:") + + def test_empty_batches_and_ranking(self): + self.assertEqual(self.judge.probs([]).shape, (0, 4)) + self.assertEqual(self.judge.compare_many([]), []) + self.assertEqual(self.judge.rank("x = x", []), []) + + def test_invalid_inputs_fail_before_inference(self): + for size in [0, -1, 1.5, True]: + with self.assertRaises(ValueError): + self.judge.probs([], batch_size=size) + for pairs in [[("", "x")], [(1, "x")], ["ab"], [("x",)]]: + with self.assertRaises(ValueError): + self.judge.probs(pairs) + + +class PipelineInterfaceTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.config_path = Path(self.temp.name) / "config.json" + self.out = Path(self.temp.name) / "run" + self.config_path.write_text(json.dumps({ + "name": "interface-test", "train": {"out_dir": str(self.out)}, + "eval": {"splits": ["test"]}, + })) + + def test_dry_run_from_path_writes_and_launches_nothing(self): + with patch("pipeline.run.subprocess.run") as process, contextlib.redirect_stdout(io.StringIO()): + result = run(self.config_path, dry_run=True) + process.assert_not_called() + self.assertFalse(self.out.exists()) + self.assertEqual(result["stages"], ["data", "train", "eval"]) + self.assertTrue(result["dry_run"]) + + def test_config_object_is_revalidated(self): + config = load(self.config_path) + config.render["kind"] = "invalid" + with self.assertRaises(SystemExit): + run(config, dry_run=True) + self.assertFalse(self.out.exists()) + + def test_bad_stage_and_unknown_keys_fail_without_writes(self): + with self.assertRaises(ValueError): + run(self.config_path, stage="typo") + self.config_path.write_text('{"train": {"typo": 5}}') + with self.assertRaises(KeyError): + run(self.config_path, dry_run=True) + self.assertFalse(self.out.exists()) + + def test_existing_adapter_allows_preview_but_prevents_training(self): + (self.out / "adapter-final").mkdir(parents=True) + with patch("pipeline.run.subprocess.run") as process, contextlib.redirect_stdout(io.StringIO()): + run(self.config_path, stage="train", dry_run=True) + with self.assertRaises(FileExistsError): + run(self.config_path, stage="train") + process.assert_not_called() + + def test_stage_failure_stops_later_stages(self): + error = subprocess.CalledProcessError(1, ["dataset-builder"]) + with patch("pipeline.run.stage_data", side_effect=error), \ + patch("pipeline.run.stage_train") as train, \ + patch("pipeline.run.stage_eval") as evaluate, \ + contextlib.redirect_stdout(io.StringIO()): + with self.assertRaises(subprocess.CalledProcessError): + run(self.config_path) + train.assert_not_called() + evaluate.assert_not_called() + + def test_evaluation_failure_is_not_silently_reported_as_success(self): + (self.out / "adapter-final").mkdir(parents=True) + error = subprocess.CalledProcessError(1, ["judge"]) + with patch("pipeline.run.subprocess.run", side_effect=error) as process, \ + contextlib.redirect_stdout(io.StringIO()): + with self.assertRaises(subprocess.CalledProcessError): + stage_eval(load(self.config_path), dry=False) + self.assertTrue(process.call_args.kwargs["check"]) + self.assertFalse((self.out / "eval_summary.json").exists()) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/test_pipeline_contract.py b/tests/test_pipeline_contract.py new file mode 100644 index 0000000..336d74c --- /dev/null +++ b/tests/test_pipeline_contract.py @@ -0,0 +1,107 @@ +"""Dataset reuse, representation controls, metadata, and metric definitions.""" + +import contextlib +import io +import json +import os +from pathlib import Path +import random +import tempfile +import unittest +from unittest.mock import patch + +from judge.build_pair_dataset import augment +from judge.train_judge import save_prompt_metadata +from pipeline.config import load +from pipeline.metrics import summarize +from pipeline.representations import parse_equation, render, selftest +from pipeline.run import run, stage_data + + +class PipelineContractTests(unittest.TestCase): + def setUp(self): + self.temp = tempfile.TemporaryDirectory() + self.addCleanup(self.temp.cleanup) + self.root = Path(self.temp.name) + self.path = self.root / "config.json" + self.path.write_text(json.dumps({"data": {"out_dir": str(self.root / "data")}, + "train": {"out_dir": str(self.root / "run")}})) + self.cfg = load(self.path) + + def test_resolved_configuration_can_be_reloaded(self): + saved = self.cfg.dump(str(self.root / "resolved")) + self.assertEqual(load(saved).as_dict(), self.cfg.as_dict()) + + def test_existing_adapter_preserves_its_configuration(self): + directory = Path(self.cfg.train.out_dir) + (directory / "adapter-final").mkdir(parents=True) + original = directory / "resolved_config.yaml" + original.write_text("original recipe") + with self.assertRaises(FileExistsError): + run(self.cfg, stage="train") + self.assertEqual(original.read_text(), "original recipe") + + def test_changed_or_unknown_data_recipe_is_rejected(self): + directory = Path(self.cfg.data.out_dir) + directory.mkdir() + (directory / "train.jsonl").write_text("{}\n") + with patch("pipeline.run.subprocess.run") as process, contextlib.redirect_stdout(io.StringIO()): + with self.assertRaises(ValueError): + stage_data(self.cfg, False) + process.assert_not_called() + + def test_matching_data_recipe_reuses_complete_splits(self): + directory = Path(self.cfg.data.out_dir) + directory.mkdir() + recipe = {"data": dict(self.cfg.data), "render": dict(self.cfg.render), "seed": self.cfg.seed} + (directory / "pipeline_data_config.json").write_text(json.dumps(recipe)) + for split in ("train", "pairs_val", "classes_val", "test"): + (directory / f"{split}.jsonl").write_text("{}\n") + with patch("pipeline.run.subprocess.run") as process, contextlib.redirect_stdout(io.StringIO()): + stage_data(self.cfg, False) + process.assert_not_called() + + def test_infix_render_controls_reach_builder(self): + eq = parse_equation("a * (b * a) = b") + env = {"SDC_RENDER_KIND": "infix", "SDC_RENDER_RENAME": "False", + "SDC_RENDER_FLIP": "False", "SDC_RENDER_OP": "@"} + with patch.dict(os.environ, env): + actual = augment(eq, random.Random(0)) + self.assertEqual(actual, "a @ (b @ a) = b") + + def test_renderers_preserve_grouping_and_variable_occurrences(self): + with contextlib.redirect_stdout(io.StringIO()): + self.assertEqual(selftest(), 0) + + def test_checkpoint_prompt_metadata(self): + directory = self.root / "checkpoint-10" + template = "Compare {a} with {b}\nRelation:" + save_prompt_metadata(directory, template) + metadata = json.loads((directory / "prompt_template.json").read_text()) + self.assertEqual(metadata["template"], template) + self.assertEqual(metadata["labels"], ["equivalent", "weaker", "stronger", "incomparable"]) + + def test_unknown_metrics_and_changed_label_order_are_rejected(self): + self.cfg.eval["metrics"] = ["made-up"] + with self.assertRaises(ValueError): + run(self.cfg, dry_run=True) + self.cfg.eval["metrics"] = ["accuracy"] + self.cfg.prompt["labels"].reverse() + with self.assertRaises(ValueError): + run(self.cfg, dry_run=True) + + def test_false_equivalence_uses_non_equivalent_denominator(self): + rows = [ + {"label": "weaker", "pred": "equivalent", "tier": "graph"}, + {"label": "stronger", "pred": "stronger", "tier": "graph"}, + {"label": "equivalent", "pred": "incomparable", "tier": "cross-node"}, + ] + result = summarize(rows, ["accuracy", "false_equivalence", "missed_equivalence", "per_tier"]) + self.assertAlmostEqual(result["accuracy"], 1 / 3) + self.assertEqual(result["false_equivalence"], {"errors": 1, "eligible": 2, "rate": 0.5}) + self.assertEqual(result["missed_equivalence"]["rate"], 1) + self.assertEqual(result["per_tier"]["graph"]["n"], 2) + self.assertNotIn("per_label", result) + self.assertIsNone(summarize([], ["accuracy"])["accuracy"]) + with self.assertRaises(ValueError): + summarize([{"label": "unknown", "pred": "equivalent"}], ["accuracy"])