"""Depth-batched best-first proof search (Phase 3). Design (chosen for this machine, after the resident-REPL route proved unreliable): * one `lean --json` call verifies ALL candidates of one search round (~5 s fixed Mathlib import + ~free per candidate), so rounds are the unit of cost; * candidates come from the 40M tactic model, greedy-decoded under an FSM mask whose first token must be one of the N most frequent first tokens in the training split (data-driven tactic whitelist); * frontier = best-first by cumulative log-prob, pruned to `beam`. Reported metric: SPS = verified candidates / wall-clock second, plus solved count. """ import os import re import sys import time from dataclasses import dataclass, field import numpy as np import torch import torch.nn.functional as F HOME = os.path.expanduser('~') sys.path.insert(0, f'{HOME}/leanprover/src') from leanoar.data import load_specials # noqa: E402 from leanoar.lean_batch import (LeanBatchVerifier, split_decl, narrowed_header, # noqa: E402 MISSING_IMPORT_RE) from leanoar.model import ModelConfig, SchemeC # noqa: E402 from tokenizers import Tokenizer # noqa: E402 @dataclass class Node: state: str tactics: list = field(default_factory=list) score: float = 0.0 depth: int = 0 class ProofSearch: def __init__(self, ckpt=f'{HOME}/leanprover/checkpoints/stage1b/best', tok_file=f'{HOME}/leanprover/data/tok/tokenizer_v1.json', steps_root=f'{HOME}/leanprover/data/steps', k=8, beam=6, max_depth=6, ctx=768, max_new_tokens=24, whitelist_n=64, time_budget=300.0, lean_timeout=180, verify_chunk=24, round_budget=90.0, force_tactics=None, force_depth=99, lean_parallel=1, narrow_imports=False, model_proposals=True, verifier=None, verbose=True): self.k, self.beam, self.max_depth = k, beam, max_depth self.ctx, self.max_new_tokens = ctx, max_new_tokens self.time_budget, self.verbose = time_budget, verbose self.verify_chunk, self.round_budget = verify_chunk, round_budget self.lean_parallel = max(1, int(lean_parallel)) self.narrow_imports = bool(narrow_imports) self.model_proposals = bool(model_proposals) self._hdr_cache = {} self.force_tactics = list(force_tactics or []) self.force_depth = force_depth S = load_specials(f'{HOME}/leanprover/data/tok') self.S = S self.pad, self.fwd = S['<|pad|>'], S['<|FORWARD|>'] self.state_id, self.tac_id, self.eos_id = S['<|state|>'], S['<|tactic|>'], S['<|eos|>'] self.tok = Tokenizer.from_file(tok_file) sd = torch.load(f'{ckpt}/lit_model.pth', map_location='cpu') from dataclasses import fields known = {f.name for f in fields(ModelConfig)} cfg = ModelConfig(**{x: y for x, y in sd['config'].items() if x in known}) self.model = SchemeC(cfg).cuda().eval() self.model.load_state_dict(sd['model']) self.step = sd.get('step') # data-driven FSM whitelist: most frequent first tokens in the train split ids = np.load(f'{steps_root}/train_ids.npy', mmap_mode='r') ms = np.load(f'{steps_root}/train_mask_start.npy') firsts = np.array([ids[i, int(ms[i])] for i in range(len(ms))], dtype=np.int64) vals, cnts = np.unique(firsts, return_counts=True) order = np.argsort(-cnts)[:whitelist_n] self.whitelist = vals[order] self.whitelist_counts = cnts[order] self.verifier = verifier or LeanBatchVerifier(verbose=verbose, timeout=lean_timeout) self.stats = dict(rounds=0, candidates=0, lean_calls=0, dead_ends=0) # ---------- model side ---------- def encode_prompt(self, state): sids = self.tok.encode(state).ids # 4 fixed tokens surround the state: [bos, <|FORWARD|>, <|state|>] + sids + [<|tactic|>], # so the budget must subtract 4 (not 3) or a full-length generation on a long state # overflows ctx by exactly one token (observed: seq 769 > block 768). budget = self.ctx - 4 - self.max_new_tokens if len(sids) > budget: # keep head + tail (the ⊢ goal line is last) head = sids[:int(budget * 0.6)] tail = sids[-(budget - len(head)):] sids = head + tail p = [1, self.fwd, self.state_id] + list(sids) + [self.tac_id] return torch.tensor([p], device='cuda', dtype=torch.long) @torch.no_grad() def propose(self, state, k=None, return_ids=False): """Return [(tactic_text, total_logprob)] for the top-k FSM-legal candidates. return_ids=True also yields the generated token ids (prompt tail + candidate + eos), which the Value head needs to score the candidate (E9 rerank experiment). """ k = k or self.k ids = self.encode_prompt(state) out = self.model(ids)['logits'][0, -1] allow = torch.full_like(out, float('-inf')) allow[torch.tensor(self.whitelist.tolist(), device='cuda')] = 0.0 lg = out + allow lps = F.log_softmax(lg.float(), -1) top = torch.topk(lg, k).indices firsts = top.tolist() vals = lps[top].tolist() # batched greedy continuation: one forward per token, k rows seqs = torch.cat([ids.repeat(len(firsts), 1), torch.tensor(firsts, device='cuda')[:, None]], dim=1) scores = list(vals) done = [False] * len(firsts) for _ in range(self.max_new_tokens): if seqs.shape[1] > self.ctx: # hard stop: never feed the model more than ctx break lg2 = self.model(seqs)['logits'][:, -1] nxt = lg2.argmax(-1) lp = F.log_softmax(lg2.float(), -1)[torch.arange(len(firsts), device='cuda'), nxt] for j in range(len(firsts)): if not done[j]: scores[j] += float(lp[j]) if int(nxt[j]) == self.eos_id: done[j] = True seqs = torch.cat([seqs, nxt[:, None]], dim=1) if all(done): break cands = [] seen = set() DANGLING = {'by', 'at', 'with', 'using', 'from', 'in', ',', ';', ':', '=>', '|', '·', '(', '['} for j, f in enumerate(firsts): gen = seqs[j, ids.shape[1]:].tolist() if gen and gen[-1] == self.eos_id: gen = gen[:-1] text = self.tok.decode(gen).strip() # FSM 1: a candidate is ONE tactic — the model writes `A B` (double space) # when it wants two steps, so cut at the first such boundary. text = re.split(r'\s{2,}', text)[0].strip() # FSM 2: drop dangling connectives left at the end of a truncated tactic while True: w = text.split() if w and w[-1] in DANGLING: text = ' '.join(w[:-1]) else: break if not text or '<|' in text or text in seen: continue if re.search(r'\b(sorry|admit|Admitted)\b', text): continue # FSM 4: never let `sorry` "close" a goal if text.count('[') != text.count(']') or text.count('(') != text.count(')'): continue # FSM 3: brackets must balance seen.add(text) cands.append((text, scores[j], list(gen) + [self.eos_id])) if return_ids: return cands return [(t, s) for t, s, _g in cands] # ---------- search ---------- def eff_header(self, header, statement): """Header actually used for verification. With `narrow_imports` on, `import Mathlib` is replaced by the minimal module set that this statement needs (Mathlib's `#min_imports`) plus `Mathlib.Tactic` for the injected automation table and for model tactics. Measured: 3.41 s -> 1.74 s per `lean --json` call, i.e. ~1.9x, because the fixed import cost dominates every call. Candidates that then fail on an apparently missing name are re-checked with the full header, so the narrowed header can never lose a proof. Cached per (header, statement). """ if not self.narrow_imports: return header key = (header, statement) if key in self._hdr_cache: return self._hdr_cache[key] eff = narrowed_header(self.verifier, header, statement) self.stats['minimp_calls'] = self.stats.get('minimp_calls', 0) + 1 if eff: self.stats['narrowed'] = self.stats.get('narrowed', 0) + 1 else: eff = header self._hdr_cache[key] = eff return eff def solve(self, header, statement, name=None, max_depth=None, time_budget=None): t0 = time.time() budget = time_budget or self.time_budget depth_max = max_depth or self.max_depth name, decl_rest = split_decl(statement) hdr = self.eff_header(header, statement) goal, raw = self.verifier.initial_state(hdr, name, decl_rest) self.stats['lean_calls'] += 1 if goal is None and hdr != header: # narrowing must never break a problem goal, raw = self.verifier.initial_state(header, name, decl_rest) self.stats['lean_calls'] += 1 hdr = header if goal is None: return {'solved': False, 'reason': 'initial-state-failed', 'proof': None, 'raw': raw['messages'][:2], 'seconds': time.time() - t0, 'stats': self.stats} frontier = [Node(state=goal, tactics=[], score=0.0, depth=0)] verified = 0 for depth in range(1, depth_max + 1): if not frontier: self.stats['dead_ends'] += 1 break items, meta = [], [] # Force a short list of known goal-closing automation tactics into every round up # to `force_depth` (default: all rounds). Measured motivation: every proof our # search ever closed ended in `ring`/`rfl`-style automation, while the model's own # pool rarely contains them; at depth>1 the automation has to close a *subgoal*, # which is exactly where these tactics shine. if self.force_tactics and depth <= self.force_depth: for root in frontier: for t in self.force_tactics: items.append({'name': name, 'decl_rest': decl_rest, 'tactics': root.tactics + [t]}) meta.append((root, t, 0.0)) for node in frontier: props = self.propose(node.state) if self.model_proposals else [] if self.verbose and props: print(f' d{depth} cands: ' + ' | '.join(t for t, _ in props[:6])[:200], flush=True) for text, sc in props: items.append({'name': name, 'decl_rest': decl_rest, 'tactics': node.tactics + [text]}) meta.append((node, text, sc)) if not items: break # Verify in chunks under a per-round wall-clock cap: cost per candidate is # ~0.4 s (measured) and a single heavy tactic can cost minutes, so one round # must never be allowed to eat a whole problem's budget. res = [] full_ids = set() # items whose verdict came from the full `import Mathlib` t_round = time.time() wave = self.verify_chunk * self.lean_parallel for i in range(0, len(items), wave): if res and (time.time() - t_round) > self.round_budget: self.stats['skipped'] = self.stats.get('skipped', 0) + (len(items) - len(res)) break part = items[i:i + wave] rs, raws = self.verifier.verify_many(hdr, part, chunk=self.verify_chunk, parallel=self.lean_parallel, name=f'{name}_d{depth}_{i}') self.stats['lean_calls'] += len(raws) if hdr != header: # A failure that smells like a missing import is re-checked against the full # `import Mathlib`: the narrowed header must never lose a proof. Only those # candidates pay the full import cost. retry = [j for j, r in enumerate(rs) if r['status'] == 'invalid' and MISSING_IMPORT_RE.search(r.get('error') or '')] if retry: rs2, raws2 = self.verifier.verify_many( header, [part[j] for j in retry], chunk=self.verify_chunk, parallel=self.lean_parallel, name=f'{name}_d{depth}_{i}_full') self.stats['lean_calls'] += len(raws2) self.stats['recheck_full'] = self.stats.get('recheck_full', 0) + len(retry) for j, r2 in zip(retry, rs2): rs[j] = r2 full_ids.add(i + j) res.extend(rs) meta = meta[:len(res)] self.stats['rounds'] += 1 self.stats['candidates'] += len(res) verified += len(res) nxt = [] for k, (r, (node, text, sc)) in enumerate(zip(res, meta)): if r['status'] == 'maybe-closed': # "no errors in the span" is NOT proof: re-compile it standalone. h_used = header if k in full_ids else hdr conf = self.verifier.confirm(h_used, name, decl_rest, node.tactics + [text]) if not conf['ok'] and h_used != header: # never lose a proof to a narrowed import set: re-confirm with everything conf = self.verifier.confirm(header, name, decl_rest, node.tactics + [text]) self.stats['confirm_calls'] = self.stats.get('confirm_calls', 0) + 1 if conf['ok']: proof = node.tactics + [text] self.stats['confirmed'] = self.stats.get('confirmed', 0) + 1 return {'solved': True, 'proof': proof, 'depth': depth, 'seconds': time.time() - t0, 'verified': verified, 'stats': self.stats, 'sps': verified / max(1e-9, time.time() - t0), 'frontier': len(frontier), 'confirm': conf} self.stats['false_positives'] = self.stats.get('false_positives', 0) + 1 continue if r['status'] == 'open': nxt.append(Node(state=r['goals'], tactics=node.tactics + [text], score=node.score + sc, depth=depth)) nxt.sort(key=lambda n: n.score, reverse=True) frontier = nxt[:self.beam] self.stats['depth'] = depth if self.verbose: print(f' depth {depth}: {len(items)} candidates ({len(res)} verified), ' f'{len(nxt)} open, round {time.time()-t_round:.0f}s (total {time.time()-t0:.0f}s)', flush=True) if time.time() - t0 > budget: break return {'solved': False, 'proof': None, 'seconds': time.time() - t0, 'verified': verified, 'stats': self.stats, 'sps': verified / max(1e-9, time.time() - t0)}