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 bc4f72a..89fa77f 100644 --- a/.gitignore +++ b/.gitignore @@ -2,3 +2,9 @@ __pycache__/ *.pyc .venv/ .env + +# judge/: generated dataset and trained weights +judge/data/ +judge/data-*/ +judge/runs/ +*.tar.gz diff --git a/README.md b/README.md index 50febe1..ebc9901 100644 --- a/README.md +++ b/README.md @@ -35,6 +35,15 @@ experiments of your own as a [design](designs/E0-TEMPLATE.md) sheet and bring it 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. + + ## The design-sheet workflow 1. Copy `designs/E0-TEMPLATE.md` to `designs/E-short-keyname.md` and fill it in 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 new file mode 100644 index 0000000..b1e3bf7 --- /dev/null +++ b/judge/README.md @@ -0,0 +1,96 @@ +# ETP relation judge: Python API and CLI + +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. + +## Python interface + +Run from the repository root: + +```python +from judge import Judge + +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()) + +results = judge.compare_many([ + ("x * y = y * x", "a * b = b * a"), + ("x * y = y * x", "x = x"), +], batch_size=2) +``` + +`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. + +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. + +## Installation and artifacts + +```bash +python -m pip install -r judge/requirements.txt -r pipeline/requirements.txt +``` + +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. + +## Command line + +```bash +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 +``` + +The original `python judge/judge.py ...` invocation remains supported. + +## Training and evaluation + +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. + +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. + +## Interpretability and limitations + +`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. + +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. + +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. + +## Verification + +```bash +python -m unittest discover -s tests -v +python -m pipeline.run --config configs/judge-v0-baseline.yaml --stage all --dry-run +``` + +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 new file mode 100644 index 0000000..f505f3f --- /dev/null +++ b/judge/build_pair_dataset.py @@ -0,0 +1,260 @@ +#!/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.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")) + +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 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) + 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..8fee92f --- /dev/null +++ b/judge/judge.py @@ -0,0 +1,288 @@ +#!/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 label. Scores are normalized over four labels, not calibrated guarantees. + +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 + +CPU, CUDA, and Apple-Silicon MPS are supported. Runtime depends on the hardware. +""" +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" +DEFAULT_TEMPLATE = "A: {a}\nB: {b}\nRelation:" + + +@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})" + + 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.""" + + 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() + + # 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: + """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.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 + 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 torch.empty((0, len(LABELS)), dtype=torch.float32).numpy()) + + def compare(self, text_a: str, text_b: str) -> Verdict: + 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. + + 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/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/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..407344c --- /dev/null +++ b/judge/train_judge.py @@ -0,0 +1,268 @@ +#!/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 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") + 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, TrainerCallback, 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.") + + # 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 = 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"]]] + 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, + ) + 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() + + 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) + 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) + 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() 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"])