code/leanoar/search.py

"""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)}