| |
| """In-band triage classifier (plan v3 Phase 0a item 4): predict, from statement text alone (CPU baseline) or from |
| 8B hidden states (GPU hook, computed later), whether the base model's pass count on a statement is IN-BAND |
| (1 <= k <= n-1 accepted samples out of n), i.e. neither never-solved nor always-solved. |
| |
| numpy only: hashed char n-gram features + a handful of dense statement features -> L2-regularised logistic |
| regression fit by full-batch Adam. Evaluation: stratified k-fold CV AUC (rank-based, ties averaged) and a |
| held-out split chosen by a hash of problem_id (not by row order), with problem-level bootstrap intervals, and the |
| AUC of trivial baselines (statement length) so "the classifier is better than length" is a number, not a feeling. |
| |
| Cautions recorded with the result (they are properties of the data, not of the code): |
| * labels are noisy: k out of 8 samples, so a problem with true p=0.9 is labelled in-band only ~57% of the time; |
| * the pool is one source (Goedel-Pset) filtered by the SAME model that labels it, and near-templated statements |
| (lean_workbook / orca_math) can sit on both sides of a random split. CV AUC is therefore optimistic for |
| OTHER sources: the real test is the transfer to the profiled sample of FormalMATH / Numina / Nemotron. |
| |
| .venv/bin/python scripts/p2/supply/triage.py train --problems .tmp/p2/problems/train_pool.jsonl \ |
| --verdicts .tmp/p2/r1/pool/ver_r0.jsonl --n 8 --out .tmp/p2/supply/triage |
| .venv/bin/python scripts/p2/supply/triage.py score --model .tmp/p2/supply/triage/model.npz \ |
| --problems <set>.jsonl --out <set>.triage.jsonl |
| """ |
|
|
| from __future__ import annotations |
|
|
| import argparse |
| import hashlib |
| import importlib.util |
| import json |
| import re |
| import sys |
| import zlib |
| from collections import defaultdict |
| from pathlib import Path |
| from typing import Iterable, Optional, Sequence |
|
|
| import numpy as np |
|
|
| REPO = Path(__file__).resolve().parents[3] |
| sys.path.insert(0, str(REPO)) |
|
|
| N_FEATURES = 1 << 14 |
| DENSE_NAMES = ("log_chars", "log_tokens", "n_hyp", "n_digits", "n_quant", "n_real", "n_nat_int", "has_informal", "log_informal") |
|
|
|
|
| def _load_common(): |
| if "p2_common" in sys.modules: |
| return sys.modules["p2_common"] |
| spec = importlib.util.spec_from_file_location("p2_common", REPO / "scripts" / "p2" / "p2_common.py") |
| mod = importlib.util.module_from_spec(spec) |
| sys.modules["p2_common"] = mod |
| spec.loader.exec_module(mod) |
| return mod |
|
|
|
|
| |
|
|
|
|
| def read_pass_counts(verdicts_path: str | Path, n: int = 8) -> dict[str, tuple[int, int]]: |
| """problem_id -> (accepted, total) from a verdict JSONL, keeping ONLY problems with exactly `n` distinct |
| samples (a problem still being generated has a partial count that would mislabel it). Tolerates a torn last |
| line (the file may be appended to by a live chain). Duplicate (problem, sample_idx) keeps the last verdict.""" |
| seen: dict[str, dict[int, bool]] = defaultdict(dict) |
| with open(verdicts_path, encoding="utf-8") as f: |
| for line in f: |
| try: |
| r = json.loads(line) |
| except json.JSONDecodeError: |
| continue |
| seen[r["problem_id"]][int(r["sample_idx"])] = bool(r["accepted"]) |
| return {p: (sum(v.values()), len(v)) for p, v in seen.items() if len(v) == n} |
|
|
|
|
| def in_band(k: int, n: int) -> int: |
| return int(1 <= k <= n - 1) |
|
|
|
|
| |
|
|
|
|
| def statement_body(statement: str) -> str: |
| s = re.sub(r"/-.*?-/", " ", statement, flags=re.S) |
| s = re.sub(r"--[^\n]*", "", s) |
| s = re.sub(r"^\s*(import|open|set_option)\b[^\n]*$", "", s, flags=re.M) |
| s = re.sub(r"\btheorem\s+\S+", "theorem", s, count=1) |
| return " ".join(s.replace(":= by sorry", "").split()) |
|
|
|
|
| def _hash_ngrams(text: str, ns: Sequence[int], salt: str, vec: dict[int, float]) -> None: |
| for n in ns: |
| for i in range(len(text) - n + 1): |
| h = zlib.crc32(f"{salt}{text[i : i + n]}".encode()) |
| idx, sign = h % N_FEATURES, 1.0 if (h >> 31) & 1 else -1.0 |
| vec[idx] = vec.get(idx, 0.0) + sign |
|
|
|
|
| def featurize(statement: str, informal: Optional[str] = None) -> tuple[np.ndarray, np.ndarray, np.ndarray]: |
| """-> (indices, values, dense). Signed hashing of char 2..4-grams of the statement body plus word unigrams and |
| bigrams of the informal text; values are sublinear (sign*log1p|c|) and the sparse part is L2-normalised.""" |
| body = statement_body(statement) |
| vec: dict[int, float] = {} |
| _hash_ngrams(body, (2, 3, 4), "s", vec) |
| iw = re.findall(r"[a-z0-9]+", (informal or "").lower()) |
| for i, w in enumerate(iw): |
| for g in (w, w + " " + iw[i + 1] if i + 1 < len(iw) else None): |
| if g: |
| h = zlib.crc32(("w" + g).encode()) |
| idx = h % N_FEATURES |
| vec[idx] = vec.get(idx, 0.0) + (1.0 if (h >> 31) & 1 else -1.0) |
| idx = np.fromiter(vec.keys(), dtype=np.int64, count=len(vec)) |
| val = np.array([np.sign(c) * np.log1p(abs(c)) for c in vec.values()], dtype=np.float64) |
| nrm = np.linalg.norm(val) |
| if nrm > 0: |
| val /= nrm |
| toks = body.split() |
| dense = np.array([ |
| np.log1p(len(body)), np.log1p(len(toks)), len(re.findall(r"\(h[\w₀-₉']*\s*:", body)), |
| len(re.findall(r"\d", body)) / max(1, len(body)) * 10, len(re.findall(r"[∀∃]", body)), |
| body.count("ℝ") + body.count("Real"), body.count("ℕ") + body.count("ℤ"), |
| float(bool(informal)), np.log1p(len(iw)), |
| ], dtype=np.float64) |
| return idx, val, dense |
|
|
|
|
| def build_matrix(rows: Sequence[dict]) -> np.ndarray: |
| """Dense float32 [n, N_FEATURES + len(DENSE_NAMES)]. 850 rows is ~56 MB; for large sets use `batches`.""" |
| X = np.zeros((len(rows), N_FEATURES + len(DENSE_NAMES)), dtype=np.float32) |
| for r, row in enumerate(rows): |
| idx, val, dense = featurize(row["formal_statement"], row.get("informal_prefix")) |
| X[r, idx] = val |
| X[r, N_FEATURES:] = dense |
| return X |
|
|
|
|
| def batches(rows: Iterable[dict], size: int = 2048): |
| buf: list[dict] = [] |
| for r in rows: |
| buf.append(r) |
| if len(buf) >= size: |
| yield buf, build_matrix(buf) |
| buf = [] |
| if buf: |
| yield buf, build_matrix(buf) |
|
|
|
|
| |
|
|
|
|
| class HiddenStateFeatures: |
| """Hook for 8B hidden-state features, computed LATER on GPU (not here: no GPU in this lane). |
| |
| Contract: a `.npz` with `ids` (str array of problem_id) and `X` (float32 [n, d]) (or a directory holding |
| `ids.json` + `X.npy`, same content) -- e.g. the mean-pooled |
| last-layer (or a middle-layer) hidden state of the Goedel-Prover-V2-8B prompt tokens of each statement, one |
| forward pass, no sampling. `train --hidden-npz` aligns it to the labelled problems by id and fits the same |
| logistic regression on it (alone, or concatenated with the text features via --combine).""" |
|
|
| def __init__(self, path: str | Path): |
| p = Path(path) |
| if p.is_dir(): |
| ids, X = json.loads((p / "ids.json").read_text()), np.load(p / "X.npy", allow_pickle=False) |
| else: |
| z = np.load(p, allow_pickle=False) |
| ids, X = z["ids"], z["X"] |
| if len(ids) != len(X): |
| raise ValueError(f"{len(ids)} ids but {len(X)} feature rows") |
| self.index = {str(i): j for j, i in enumerate(ids)} |
| self.X = X.astype(np.float32) |
|
|
| def matrix(self, ids: Sequence[str]) -> np.ndarray: |
| miss = [i for i in ids if i not in self.index] |
| if miss: |
| raise KeyError(f"{len(miss)} problems have no hidden-state row, e.g. {miss[:3]}") |
| return self.X[[self.index[i] for i in ids]] |
|
|
| @staticmethod |
| def compute_on_gpu(*_a, **_k): |
| raise NotImplementedError("run on GPU in a heavy-lock window: forward the prompt through the 8B, pool, np.savez(ids=..., X=...)") |
|
|
|
|
| |
|
|
|
|
| def auc(y: np.ndarray, s: np.ndarray) -> float: |
| """Rank AUC with average ranks for ties (== Mann-Whitney U / (n_pos n_neg)).""" |
| y = np.asarray(y).astype(bool) |
| npos, nneg = int(y.sum()), int((~y).sum()) |
| if npos == 0 or nneg == 0: |
| return float("nan") |
| order = np.argsort(s, kind="mergesort") |
| ss = np.asarray(s)[order] |
| ranks = np.empty(len(s)) |
| i = 0 |
| while i < len(ss): |
| j = i |
| while j + 1 < len(ss) and ss[j + 1] == ss[i]: |
| j += 1 |
| ranks[order[i : j + 1]] = (i + j) / 2 + 1 |
| i = j + 1 |
| return float((ranks[y].sum() - npos * (npos + 1) / 2) / (npos * nneg)) |
|
|
|
|
| class Logistic: |
| """L2 logistic regression, full-batch Adam. Features are standardised using training statistics (dense part |
| only; the sparse part is already unit-norm). Intercept unpenalised.""" |
|
|
| def __init__(self, l2: float = 1e-2, iters: int = 400, lr: float = 0.05, seed: int = 0): |
| self.l2, self.iters, self.lr, self.seed = l2, iters, lr, seed |
| self.w: Optional[np.ndarray] = None |
| self.b = 0.0 |
| self.mu: Optional[np.ndarray] = None |
| self.sd: Optional[np.ndarray] = None |
|
|
| def _prep(self, X: np.ndarray) -> np.ndarray: |
| return (X - self.mu) / self.sd |
|
|
| def fit(self, X: np.ndarray, y: np.ndarray, sample_weight: Optional[np.ndarray] = None) -> "Logistic": |
| X = X.astype(np.float64) |
| self.mu = X.mean(0) |
| self.sd = X.std(0) |
| self.sd[self.sd < 1e-8] = 1.0 |
| |
| self.mu[:N_FEATURES] = 0.0 |
| self.sd[:N_FEATURES] = 1.0 |
| if X.shape[1] <= N_FEATURES: |
| self.mu, self.sd = X.mean(0), np.maximum(X.std(0), 1e-8) |
| Z = self._prep(X) |
| n, d = Z.shape |
| sw = np.ones(n) if sample_weight is None else sample_weight / sample_weight.mean() |
| w, b = np.zeros(d), float(np.log((y.mean() + 1e-3) / (1 - y.mean() + 1e-3))) |
| m, v, mb, vb = np.zeros(d), np.zeros(d), 0.0, 0.0 |
| for t in range(1, self.iters + 1): |
| p = 1 / (1 + np.exp(-(Z @ w + b))) |
| r = (p - y) * sw |
| g = Z.T @ r / n + self.l2 * w |
| gb = r.mean() |
| m, v = 0.9 * m + 0.1 * g, 0.999 * v + 0.001 * g * g |
| mb, vb = 0.9 * mb + 0.1 * gb, 0.999 * vb + 0.001 * gb * gb |
| mh, vh = m / (1 - 0.9**t), v / (1 - 0.999**t) |
| w -= self.lr * mh / (np.sqrt(vh) + 1e-8) |
| b -= self.lr * (mb / (1 - 0.9**t)) / (np.sqrt(vb / (1 - 0.999**t)) + 1e-8) |
| self.w, self.b = w, b |
| return self |
|
|
| def decision(self, X: np.ndarray) -> np.ndarray: |
| return self._prep(X.astype(np.float64)) @ self.w + self.b |
|
|
| def proba(self, X: np.ndarray) -> np.ndarray: |
| return 1 / (1 + np.exp(-self.decision(X))) |
|
|
|
|
| def stratified_folds(y: np.ndarray, k: int, seed: int) -> np.ndarray: |
| rs = np.random.RandomState(seed) |
| fold = np.zeros(len(y), dtype=int) |
| for c in (0, 1): |
| idx = np.where(y == c)[0] |
| rs.shuffle(idx) |
| fold[idx] = np.arange(len(idx)) % k |
| return fold |
|
|
|
|
| def cross_val_auc(X: np.ndarray, y: np.ndarray, k: int = 5, seed: int = 0, **kw) -> dict: |
| fold = stratified_folds(y, k, seed) |
| oof = np.zeros(len(y)) |
| per = [] |
| for f in range(k): |
| tr, te = fold != f, fold == f |
| m = Logistic(**kw).fit(X[tr], y[tr]) |
| oof[te] = m.decision(X[te]) |
| per.append(auc(y[te], oof[te])) |
| return {"fold_auc": [round(a, 4) for a in per], "mean": float(np.mean(per)), "std": float(np.std(per)), "oof_auc": auc(y, oof)} |
|
|
|
|
| def bootstrap_auc_ci(y: np.ndarray, s: np.ndarray, B: int = 1000, seed: int = 0) -> tuple[float, float]: |
| rs = np.random.RandomState(seed) |
| vals = [] |
| for _ in range(B): |
| i = rs.randint(0, len(y), len(y)) |
| a = auc(y[i], s[i]) |
| if not np.isnan(a): |
| vals.append(a) |
| return float(np.percentile(vals, 2.5)), float(np.percentile(vals, 97.5)) |
|
|
|
|
| def holdout_mask(ids: Sequence[str], frac: float = 0.25, salt: str = "triage-v1") -> np.ndarray: |
| """Deterministic by problem_id hash, independent of row order.""" |
| return np.array([int(hashlib.sha256((salt + i).encode()).hexdigest()[:8], 16) / 0xFFFFFFFF < frac for i in ids]) |
|
|
|
|
| |
|
|
|
|
| def cmd_train(a) -> int: |
| C = _load_common() |
| probs = {r["problem_id"]: r for r in C.iter_jsonl(a.problems)} |
| counts = read_pass_counts(a.verdicts, a.n) |
| ids = sorted(i for i in counts if i in probs) |
| rows = [probs[i] for i in ids] |
| k = np.array([counts[i][0] for i in ids]) |
| y = np.array([in_band(int(x), a.n) for x in k]) |
| print(f"labelled problems: {len(ids)} (of {len(counts)} with n={a.n} verdicts, {len(probs)} in pool); " |
| f"in-band {int(y.sum())} ({y.mean():.1%}); k histogram {np.bincount(k, minlength=a.n + 1).tolist()}") |
| Xt = build_matrix(rows) |
| hidden = HiddenStateFeatures(a.hidden_npz).matrix(ids) if a.hidden_npz else None |
| X = Xt if hidden is None else (np.hstack([Xt, hidden]) if a.combine else hidden) |
| held = holdout_mask(ids) |
| res: dict = {"n_labelled": len(ids), "in_band": int(y.sum()), "k_hist": np.bincount(k, minlength=a.n + 1).tolist(), |
| "features": "text_hash_char_ngram" + ("+hidden" if hidden is not None and a.combine else "_hidden" if hidden is not None else ""), |
| "n_train": int((~held).sum()), "n_heldout": int(held.sum()), "heldout_in_band": int(y[held].sum())} |
| best = None |
| for l2 in (1e-3, 1e-2, 1e-1, 1.0): |
| cv = cross_val_auc(X[~held], y[~held], k=5, seed=0, l2=l2) |
| res.setdefault("cv_by_l2", {})[str(l2)] = round(cv["oof_auc"], 4) |
| if best is None or cv["oof_auc"] > best[1]["oof_auc"]: |
| best = (l2, cv) |
| l2, cv = best |
| res["chosen_l2"], res["train_cv"] = l2, cv |
| m = Logistic(l2=l2).fit(X[~held], y[~held]) |
| s = m.decision(X[held]) |
| lo, hi = bootstrap_auc_ci(y[held], s) |
| res["heldout_auc"] = round(auc(y[held], s), 4) |
| res["heldout_auc_ci95"] = [round(lo, 4), round(hi, 4)] |
| |
| length = np.array([len(statement_body(r["formal_statement"])) for r in rows], dtype=float) |
| res["baseline_heldout_auc"] = {"statement_length": round(auc(y[held], length[held]), 4), |
| "neg_statement_length": round(auc(y[held], -length[held]), 4)} |
| |
| for name, tgt in (("solved_at_least_once", (k >= 1).astype(int)), ("always_solved", (k == a.n).astype(int))): |
| mm = Logistic(l2=l2).fit(X[~held], tgt[~held]) |
| res.setdefault("secondary_heldout_auc", {})[name] = round(auc(tgt[held], mm.decision(X[held])), 4) |
| final = Logistic(l2=l2).fit(X, y) |
| out = Path(a.out) |
| out.mkdir(parents=True, exist_ok=True) |
| np.savez(out / "model.npz", w=final.w, b=final.b, mu=final.mu, sd=final.sd, n_features=N_FEATURES) |
| (out / "triage_report.json").write_text(json.dumps(res, indent=1)) |
| print(json.dumps(res, indent=1)) |
| return 0 |
|
|
|
|
| def cmd_score(a) -> int: |
| C = _load_common() |
| z = np.load(a.model) |
| m = Logistic() |
| m.w, m.b, m.mu, m.sd = z["w"], float(z["b"]), z["mu"], z["sd"] |
| n = 0 |
| with open(a.out, "w", encoding="utf-8") as f: |
| for rows, X in batches(C.iter_jsonl(a.problems)): |
| for r, p in zip(rows, m.proba(X)): |
| f.write(json.dumps({"problem_id": r["problem_id"], "p_in_band": round(float(p), 5)}) + "\n") |
| n += len(rows) |
| print(f"scored {n} -> {a.out}") |
| return 0 |
|
|
|
|
| def main(argv=None) -> int: |
| ap = argparse.ArgumentParser(description=__doc__.split("\n")[0]) |
| sub = ap.add_subparsers(dest="cmd", required=True) |
| t = sub.add_parser("train") |
| t.add_argument("--problems", required=True) |
| t.add_argument("--verdicts", required=True) |
| t.add_argument("--n", type=int, default=8) |
| t.add_argument("--out", required=True) |
| t.add_argument("--hidden-npz", help="GPU features (ids, X); see HiddenStateFeatures") |
| t.add_argument("--combine", action="store_true", help="concatenate hidden-state and text features") |
| t.set_defaults(fn=cmd_train) |
| s = sub.add_parser("score") |
| s.add_argument("--model", required=True) |
| s.add_argument("--problems", required=True) |
| s.add_argument("--out", required=True) |
| s.set_defaults(fn=cmd_score) |
| a = ap.parse_args(argv) |
| return a.fn(a) |
|
|
|
|
| if __name__ == "__main__": |
| raise SystemExit(main()) |
|
|
