bev-decisions 0.2.0__py3-none-any.whl

bev/__init__.py

@@ -0,0 +1,24 @@
+"""Bev: open-weight models that answer typed questions about a document with calibrated probabilities.
+
+    import bev
+    model = bev.load()                              # Bev-Large (laptop or GPU); bev.load(bev.PRO_MODEL) for Bev-Pro (GPU)
+    answers = model.decide(
+        "Invoice INV-2291. Total: $9,412.50. Approval limit: $8,900.",
+        {"over_limit": bev.presets.over_limit("$8,900")},
+    )
+    sure, review = bev.gate(answers, threshold=0.8)
+
+Question types: "noul" (yes/no), "choice" (criteria: {key: description}), "score" (criteria: levels, lowest first).
+Up to 16 questions per call, up to 300 options per question, 4,096 input tokens.
+"""
+
+from . import presets
+from .evaluate import evaluate, evaluate_file
+from .gating import confidence, gate, value
+from .model import DEFAULT_MODEL, PRO_MODEL, Bev, load
+from .router import Router
+from .schema import decide_as, questions_for
+
+__all__ = ["DEFAULT_MODEL", "PRO_MODEL", "Bev", "Router", "confidence", "decide_as", "evaluate", "evaluate_file", "gate", "load",
+           "presets", "questions_for", "value"]
+__version__ = "0.2.0"

bev/_core/features.py

@@ -0,0 +1,313 @@
+"""Typed values and facts read from the document and the questions: amounts, dates, times, numbers, and the comparisons,
+sums, counts and time windows they imply. The network uses them alongside the text."""
+
+from __future__ import annotations
+
+import datetime as dt
+import math
+import re
+from dataclasses import dataclass, field
+
+MONTHS = {m: i + 1 for i, m in enumerate(["jan", "feb", "mar", "apr", "may", "jun", "jul", "aug", "sep", "oct", "nov", "dec"])}
+STOP = set(
+    "a an the of to in on at by for from with and or is are was were be been did does do how many much what which who whose whether any at least more than less each this that these those it its as per within single there their ask asks".split()
+)
+_money = re.compile("(?:USD\\s*|\\$)\\s?(\\d[\\d,]*(?:\\.\\d+)?)(k)?|(\\d[\\d,]*(?:\\.\\d+)?)\\s*USD", re.I)
+_iso_ts = re.compile("\\b(\\d{4})-(\\d{2})-(\\d{2})(?:[ T](\\d{2}):(\\d{2})(?::(\\d{2}))?)?\\b")
+_dmy = re.compile("\\b(\\d{1,2})\\s+([A-Za-z]{3,9})\\.?\\s+(\\d{4})\\b")
+_mdy = re.compile("\\b([A-Za-z]{3,9})\\.?\\s+(\\d{1,2})(?:st|nd|rd|th)?,?\\s+(\\d{4})\\b")
+_clock = re.compile("(?<![\\d-])\\b(\\d{1,2}):(\\d{2})\\b(?!:\\d{2}\\s*\\d)")
+_num = re.compile("(?<![\\w.])(\\d+(?:\\.\\d+)?)(?=(?:ms|%|\\s|$|[^\\w.]))")
+_ip = re.compile("\\b(?:\\d{1,3}\\.){3}\\d{1,3}\\b")
+_word = re.compile("[A-Za-z][A-Za-z0-9_\\-]*")
+
+
+@dataclass
+class Value:
+    kind: str
+    value: float
+    start: int
+    end: int
+    line: int
+
+
+@dataclass
+class Fact:
+    kind: str
+    feats: list[float]
+    spans: list[tuple[int, int]] = field(default_factory=list)
+    where: str = "state"
+    result: float | None = None
+
+
+FACT_KINDS = ["cmp", "pair", "sum", "count", "window", "run"]
+FACT_DIM = 8
+VAL_DIM = 4 + 2 + 16
+
+
+def _month(s: str) -> int | None:
+    return MONTHS.get(s[:3].lower())
+
+
+def _minutes(d: dt.datetime) -> float:
+    return (d - dt.datetime(2000, 1, 1)).total_seconds() / 60
+
+
+def _line_of(text: str) -> list[int]:
+    starts = [0] + [m.end() for m in re.finditer("\n", text)]
+    return starts
+
+
+def extract(text: str) -> list[Value]:
+    out: list[Value] = []
+    taken: list[tuple[int, int]] = []
+    starts = _line_of(text)
+
+    def line(pos: int) -> int:
+        lo, hi = (0, len(starts) - 1)
+        while lo < hi:
+            mid = (lo + hi + 1) // 2
+            if starts[mid] <= pos:
+                lo = mid
+            else:
+                hi = mid - 1
+        return lo
+
+    def free(a: int, b: int) -> bool:
+        return all((b <= x or a >= y for x, y in taken))
+
+    def add(kind: str, v: float, a: int, b: int) -> None:
+        if free(a, b):
+            out.append(Value(kind, v, a, b, line(a)))
+            taken.append((a, b))
+
+    for m in _money.finditer(text):
+        raw = m.group(1) or m.group(3)
+        v = float(raw.replace(",", ""))
+        if m.group(2):
+            v *= 1000
+        add("money", v, m.start(), m.end())
+    for m in _iso_ts.finditer(text):
+        y, mo, d, hh, mm, ss = m.groups()
+        try:
+            t = dt.datetime(int(y), int(mo), int(d), int(hh or 0), int(mm or 0), int(ss or 0))
+        except ValueError:
+            continue
+        add("date", _minutes(t), m.start(), m.end())
+    for rx, order in ((_dmy, "dmy"), (_mdy, "mdy")):
+        for m in rx.finditer(text):
+            g = m.groups()
+            d, mo, y = (g[0], g[1], g[2]) if order == "dmy" else (g[1], g[0], g[2])
+            mon = _month(mo)
+            if mon is None:
+                continue
+            try:
+                add("date", _minutes(dt.datetime(int(y), mon, int(d))), m.start(), m.end())
+            except ValueError:
+                continue
+    for m in _clock.finditer(text):
+        add("time", int(m.group(1)) * 60 + int(m.group(2)), m.start(), m.end())
+    for m in _num.finditer(text):
+        add("number", float(m.group(1)), m.start(), m.end())
+    out.sort(key=lambda v: v.start)
+    return out
+
+
+def key_terms(question: str, state_vocab: set[str]) -> list[str]:
+    terms = []
+    for w in _word.findall(question.lower()):
+        if w in STOP or len(w) < 3:
+            continue
+        stem = w[:-1] if w.endswith("s") and len(w) > 4 else w
+        if any(v.startswith(stem) for v in state_vocab):
+            terms.append(stem)
+    terms += _ip.findall(question)
+    return sorted(set(terms))
+
+
+def _lines_with(lines: list[str], terms: list[str]) -> list[int]:
+    low = [x.lower() for x in lines]
+    return [i for i, x in enumerate(low) if all(t.lower() in x for t in terms)]
+
+
+def _signed_log(x: float) -> float:
+    return math.copysign(math.log1p(abs(x)), x)
+
+
+def facts(state: str, question: str) -> list[Fact]:
+    sv, qv = (extract(state), extract(question))
+    lines = state.split("\n")
+    vocab = {w for w in _word.findall(state.lower())}
+    terms = key_terms(question, vocab)
+    return _question_facts(sv, qv, lines, vocab, terms, question)
+
+
+def state_facts(state: str) -> list[Fact]:
+    sv = extract(state)
+    fs: list[Fact] = []
+    for kind in ("date", "money"):
+        vals = [s for s in sv if s.kind == kind][:12]
+        for i in range(len(vals)):
+            for j in range(len(vals)):
+                if i != j:
+                    diff = vals[j].value - vals[i].value
+                    unit = diff / 1440 if kind == "date" else diff
+                    fs.append(
+                        Fact(
+                            "pair",
+                            [float(diff > 0), float(diff == 0), _signed_log(unit), float(kind == "date"), 0, 0, 0, 0],
+                            [(vals[i].start, vals[i].end), (vals[j].start, vals[j].end)],
+                        )
+                    )
+    return fs
+
+
+def _question_facts(sv: list[Value], qv: list[Value], lines: list[str], vocab: set[str], terms: list[str], question: str) -> list[Fact]:
+    fs: list[Fact] = []
+    for q in qv:
+        same = [s for s in sv if s.kind == q.kind]
+        for s in same[:32]:
+            diff = s.value - q.value
+            fs.append(
+                Fact(
+                    "cmp",
+                    [float(diff > 0), float(diff == 0), _signed_log(diff), _signed_log(s.value), 0, 0, 0, 0],
+                    [(s.start, s.end), (q.start, q.end)],
+                    "both",
+                )
+            )
+    money = [s for s in sv if s.kind == "money"]
+    if money:
+        tot = sum(s.value for s in money)
+        for q in [q for q in qv if q.kind == "money"] or [None]:
+            ref = q.value if q else 0.0
+            fs.append(
+                Fact(
+                    "sum",
+                    [
+                        float(tot > ref) if q else 0.0,
+                        _signed_log(tot - ref),
+                        _signed_log(tot),
+                        float(len(money)),
+                        float(q is not None),
+                        0,
+                        0,
+                        0,
+                    ],
+                    [(q.start, q.end)] if q else [],
+                    "question" if q else "state",
+                )
+            )
+    quoted = [m.lower() for m in re.findall("'([^']+)'|\\\"([^\\\"]+)\\\"", question) for m in m if m]
+    if terms or quoted:
+        k = len(_lines_with(lines, sorted(set(terms + quoted))))
+        fs.append(Fact("count", [float(k), math.log1p(k), float(len(terms)), 0, 0, 0, 0, 0], result=float(k)))
+        allw = sorted(
+            {w[:-1] if w.endswith("s") and len(w) > 4 else w for w in _word.findall(question.lower()) if w not in STOP and len(w) > 2}
+            | set(quoted)
+            | set(_ip.findall(question))
+        )
+        k_all = len(_lines_with(lines, allw))
+        fs.append(Fact("count", [float(k_all), math.log1p(k_all), float(len(allw)), 0, 0, 1.0, 0, 0], result=float(k_all)))
+        for drop in allw:
+            sub = [w for w in allw if w != drop]
+            if sub:
+                kd = len(_lines_with(lines, sub))
+                fs.append(Fact("count", [float(kd), math.log1p(kd), float(len(sub)), 0, 0, 0, 1.0, 0], result=float(kd)))
+        for t in terms:
+            kt = len(_lines_with(lines, [t]))
+            fs.append(Fact("count", [float(kt), math.log1p(kt), 1.0, 1.0, 0, 0, 0, 0], result=float(kt)))
+    for q in qv:
+        same = [s for s in sv if s.kind == q.kind]
+        if same:
+            above = sum(s.value > q.value for s in same)
+            fs.append(
+                Fact(
+                    "count",
+                    [float(above), math.log1p(above), 0, 0, 1.0, float(len(same)), 0, 0],
+                    [(q.start, q.end)],
+                    "question",
+                    result=float(above),
+                )
+            )
+    wq = re.search("(\\d+)\\s*-?\\s*minute", question)
+    if wq and terms:
+        w = float(wq.group(1))
+        stamps: dict[str, list[float]] = {}
+        content_terms = [t for t in terms if not _ip.fullmatch(t)]
+        for i in _lines_with(lines, content_terms):
+            ts = [v.value for v in sv if v.line == i and v.kind == "date"]
+            ip = _ip.findall(lines[i])
+            if ts:
+                stamps.setdefault(ip[0] if ip else "*", []).append(ts[0])
+        best = 0
+        for ts in stamps.values():
+            ts.sort()
+            j = 0
+            for i in range(len(ts)):
+                while ts[i] - ts[j] > w:
+                    j += 1
+                best = max(best, i - j + 1)
+        fs.append(Fact("window", [float(best), math.log1p(best), w, 0, 0, 0, 0, 0], result=float(best)))
+    ents = [t for t in terms if t in vocab or "-" in t]
+    for q in [q for q in qv if q.kind == "number"]:
+        for ent in ents[:3]:
+            idx = _lines_with(lines, [ent])
+            if len(idx) < 2:
+                continue
+            per_line = [[v.value for v in sv if v.line == i and v.kind == "number"] for i in idx]
+            width = min((len(x) for x in per_line), default=0)
+            for col in range(min(width, 4)):
+                run = best = 0
+                for nums in per_line:
+                    run = run + 1 if nums[col] > q.value else 0
+                    best = max(best, run)
+                fs.append(
+                    Fact(
+                        "run",
+                        [float(best), math.log1p(best), float(col), float(len(idx)), 0, 0, 0, 0],
+                        [(q.start, q.end)],
+                        "question",
+                        result=float(best),
+                    )
+                )
+    return fs[:128]
+
+
+def option_feats(state: str, options: list[str]) -> list[list[float]]:
+    low = state.lower()
+    lines = low.split("\n")
+    vocab = set(_word.findall(low))
+    out = []
+    for o in options:
+        ol = o.lower()
+        toks = [t for t in _word.findall(ol) if t not in STOP]
+        overlap = sum(t in vocab for t in toks) / max(len(toks), 1)
+        try:
+            num = float(o)
+        except ValueError:
+            num = -1.0
+        first = 1.0
+        idx = [i for i, x in enumerate(lines) if ol in x] if len(ol) > 2 else []
+        if len(idx) >= 4:
+            cols = [[float(n) for n in _num.findall(lines[i])] for i in idx]
+            width = min(len(c) for c in cols)
+            for col in range(min(width, 4)):
+                series = [c[col] for c in cols]
+                head = sorted(series[: max(3, len(series) // 4)])
+                base = head[len(head) // 2]
+                hit = next((k for k, v in enumerate(series) if base > 0 and v > 2.5 * base and (v - base > 0.5)), None)
+                if hit is not None:
+                    first = min(first, idx[hit] / max(len(lines), 1))
+        out.append([float(ol in low), overlap, num, first, float(num >= 0)])
+    return out
+
+
+def value_features(v: Value) -> list[float]:
+    """A typed value as 22 numbers: its kind, signed log magnitude, a day ordinal for dates, and 16 Fourier features."""
+    kinds = ["money", "date", "time", "number"]
+    raw = math.log1p(abs(v.value)) if math.isfinite(v.value) else 709.78 if not math.isnan(v.value) else 0.0
+    lv = math.copysign(min(raw, 709.78), v.value) if not math.isnan(v.value) else 0.0
+    four = [f(lv * 2**k / 4) for k in range(8) for f in (math.sin, math.cos)]
+    ordinal = (v.value / 1440 - 9500) / 1000 if v.kind == "date" else 0.0
+    return [float(v.kind == k) for k in kinds] + [lv / 10, ordinal] + four

bev/_core/network.py

@@ -0,0 +1,342 @@
+"""The Bev-Large network (inference only).
+
+A ModernBERT encoder reads the questions, their options and the document in one sequence; typed values (amounts, dates,
+times, numbers) are added to the token embeddings. Each question becomes a vector that gathers evidence from the
+document, is refined over three recursive steps together with the facts read from the text, and scores its options.
+The three steps are combined with learned halting weights, and the per-type temperature turns scores into calibrated
+probabilities.
+"""
+
+from __future__ import annotations
+
+import bisect
+import math
+import re
+
+import torch
+from torch import nn
+
+from .features import FACT_DIM, FACT_KINDS, VAL_DIM, extract, facts, option_feats, state_facts, value_features
+
+D = 256  # width of the decision layers
+OPT_FEATS = 5
+STEPS = 3
+QTYPES = {"noul": 0, "choice": 1, "score": 2}
+MIN_STATE_TOKENS = 160  # long option lists are shortened so that at least this much of the document fits
+ALL_QUESTIONS = 15  # fact tag for facts about the document as a whole
+
+
+class Prompt:
+    """One request as a single token sequence: [Q1] type: instructions [O] option ... [STATE] document."""
+
+    def __init__(self, tok, item: dict, max_len: int) -> None:
+        self.names = list(item["questions"])
+        qs = item["questions"]
+        self.late: set[str] = set()  # questions whose options are encoded separately (too many to fit)
+        self.compact: set[str] = set()  # questions whose option descriptions are shortened
+        while True:
+            self._build(item)
+            n_head = len(tok(self.head, add_special_tokens=False)["input_ids"])
+            joint = [n for n in self.names if qs[n]["type"] != "noul" and n not in self.late]
+            if n_head <= max_len - MIN_STATE_TOKENS or not joint:
+                break
+            big = max(joint, key=lambda n: len(qs[n]["criteria"]))
+            (self.late if big in self.compact else self.compact).add(big)
+        self.state_off = len(self.head)
+        enc = tok(self.head + item["state"], truncation=True, max_length=max_len, return_offsets_mapping=True)
+        self.ids, self.offsets = enc["input_ids"], enc["offset_mapping"]
+        self.ends, m = [], 0
+        for _, e in self.offsets:
+            m = max(m, e)
+            self.ends.append(m)
+        self.values = [(v, self.state_off) for v in extract(item["state"])]
+        for base, txt in self.value_texts:
+            self.values += [(v, base) for v in extract(txt)]
+
+    def _build(self, item: dict) -> None:
+        parts: list[str] = []
+        pos = 0
+        self.qspans, self.instr_off, self.ospans, self.value_texts = [], [], {}, []
+
+        def add(s: str) -> int:
+            nonlocal pos
+            start = pos
+            parts.append(s)
+            pos += len(s)
+            return start
+
+        for i, n in enumerate(self.names):
+            q = item["questions"][n]
+            add(f"[Q{i + 1}] {q['type']}: ")
+            s = add(q["instructions"])
+            self.instr_off.append(s)
+            self.value_texts.append((s, q["instructions"]))
+            self.qspans.append((s, pos))
+            crit = q.get("criteria")
+            if q["type"] == "noul":
+                c = crit if isinstance(crit, dict) else {}
+                spans = []
+                for txt in (f"yes: {c.get('true', 'true')}", f"no: {c.get('false', 'false')}"):
+                    add(" [O] ")
+                    o = add(txt)
+                    spans.append((o, o + len(txt)))
+                self.ospans[n] = spans
+            elif n not in self.late:
+                spans = []
+                for k, desc in crit.items():
+                    add(" [O] ")
+                    txt = f"{k}: {desc}"
+                    if n in self.compact:
+                        same = re.sub(r"[\W_]+", " ", str(k)).strip().lower() == re.sub(r"[\W_]+", " ", str(desc)).strip().lower()
+                        txt = str(k) if same else f"{k}: {' '.join(str(desc).split()[:6])}"
+                    o = add(txt)
+                    spans.append((o, o + len(txt)))
+                    self.value_texts.append((o, txt))
+                self.ospans[n] = spans
+            add(" ")
+        add("[STATE] ")
+        self.head = "".join(parts)
+
+    def tokens_in(self, a: int, b: int) -> list[int]:
+        """Token indices covering characters a..b of the prompt."""
+        out = []
+        k = bisect.bisect_right(self.ends, a)
+        while k < len(self.offsets) and self.offsets[k][0] < b:
+            s, e = self.offsets[k]
+            if e > s:
+                out.append(k)
+            k += 1
+        return out
+
+
+class BevNetwork(nn.Module):
+    """Bev-Large: ModernBERT-large encoder plus the decision layers. Build it from the encoder configuration that ships
+    with the weights (no download), then load model.safetensors."""
+
+    def __init__(self, encoder_dir: str) -> None:
+        super().__init__()
+        from transformers import AutoConfig, AutoModel
+
+        self.enc = AutoModel.from_config(AutoConfig.from_pretrained(encoder_dir), attn_implementation="sdpa")
+        H = self.enc.config.hidden_size
+        self.H, self.pad_id = H, self.enc.config.pad_token_id
+        self.val_proj = nn.Linear(VAL_DIM, H)
+        self.q_proj, self.o_proj, self.o_late = nn.Linear(H, D), nn.Linear(H, D), nn.Linear(H, D)
+        self.qtype, self.qid = nn.Embedding(3, D), nn.Embedding(16, D)
+        self.ev_k, self.ev_v, self.ev_out = nn.Linear(H, D), nn.Linear(H, D), nn.Linear(2 * D, D)
+        self.fact_in = nn.Sequential(nn.Linear(FACT_DIM + len(FACT_KINDS) + 2 * H, D), nn.GELU(), nn.Linear(D, D))
+        self.core = nn.TransformerEncoderLayer(D, 8, 4 * D, dropout=0.0, batch_first=True, norm_first=True)
+        self.core_gate = nn.Parameter(torch.zeros(1))
+        self.halt = nn.Linear(D, 1)
+        self.opt_mlp = nn.Sequential(nn.Linear(3 * D, D), nn.GELU(), nn.Linear(D, 1))
+        self.op_gate, self.opt_lex = nn.Linear(D, 1), nn.Linear(OPT_FEATS, 1)
+        self.fact_att = nn.Linear(D, D, bias=False)
+        self.num_w = nn.Parameter(torch.tensor(1.0))
+        self.fact_q, self.fact_mix = nn.Linear(D, D, bias=False), nn.Linear(D, D)
+        self.dec_sel = nn.Sequential(nn.Linear(3 * D, D), nn.GELU(), nn.Linear(D, 1))
+        self.dec_none, self.dec_w = nn.Linear(D, 1), nn.Linear(D, 1)
+        self.temp = nn.ParameterDict({t: nn.Parameter(torch.tensor(1.0), requires_grad=False) for t in QTYPES})
+
+    # yes/no questions whose answer follows from an exact comparison of numbers or dates found in the text
+    def _exact_yes_no(self, base: torch.Tensor, zq: torch.Tensor, Fc: torch.Tensor, bits: torch.Tensor) -> torch.Tensor:
+        zz = zq.unsqueeze(0).expand(Fc.shape[0], -1)
+        a = torch.softmax(torch.cat([self.dec_none(zq), self.dec_sel(torch.cat([zz, Fc, zz * Fc], -1)).squeeze(-1)]), 0)
+        comp = torch.sigmoid(torch.cat([base, base + self.dec_w(zq) * (2 * bits - 1)]))
+        p = (a * comp).sum().clamp(1e-06, 1 - 1e-06)
+        return torch.log(p / (1 - p)).reshape(1)
+
+    def _facts(self, it: dict, e: Prompt, L: int):
+        nq = len(e.names)
+        fl: list[tuple[int, object]] = [(ALL_QUESTIONS, f) for f in state_facts(it["state"])]
+        for i, n in enumerate(e.names):
+            fl += [(i, f) for f in facts(it["state"], it["questions"][n]["instructions"])]
+        if not fl:
+            return None
+        feats = torch.zeros(len(fl), FACT_DIM + len(FACT_KINDS))
+        smat = torch.zeros(2 * len(fl), L)
+        tags, res_v = [], []
+        for r_, (i, f) in enumerate(fl):
+            feats[r_] = torch.tensor(f.feats + [float(f.kind == k) for k in FACT_KINDS])
+            wheres = {"state": ("state", "state"), "both": ("state", "question"), "question": ("question", "question")}[f.where]
+            for sidx, ((s0, s1), where) in enumerate(zip(f.spans[:2], wheres, strict=False)):
+                base = e.state_off if where == "state" else e.instr_off[i if i < nq else 0]
+                tk = e.tokens_in(base + s0, base + s1)
+                if tk:
+                    smat[2 * r_ + sidx, tk] = 1.0 / len(tk)
+            tags.append(i)
+            res_v.append(f.result if f.result is not None else float("nan"))
+        return feats, smat, tags, res_v
+
+    def _encode(self, prompts: list[Prompt], dev: torch.device, L: int) -> torch.Tensor:
+        ids = torch.full((len(prompts), L), self.pad_id, dtype=torch.long)
+        mask = torch.zeros((len(prompts), L), dtype=torch.long)
+        vfeat = torch.zeros((len(prompts), L, VAL_DIM))
+        for b, e in enumerate(prompts):
+            ids[b, : len(e.ids)] = torch.tensor(e.ids)
+            mask[b, : len(e.ids)] = 1
+            for v, off in e.values:
+                for k in e.tokens_in(off + v.start, off + v.end):
+                    vfeat[b, k] = torch.tensor(value_features(v))
+        ids, mask, vfeat = ids.to(dev), mask.to(dev), vfeat.to(dev)
+        emb = self.enc.get_input_embeddings()(ids) + self.val_proj(vfeat)
+        return self.enc(inputs_embeds=emb, attention_mask=mask).last_hidden_state
+
+    def _encode_options(self, texts: list[str], tok, dev: torch.device) -> torch.Tensor:
+        """Mean-pooled encoder vectors of option texts encoded on their own (questions with very many options)."""
+        uniq = sorted(set(texts))
+        out = {}
+        for i in range(0, len(uniq), 256):
+            chunk = uniq[i : i + 256]
+            n_real = len(chunk)
+            chunk = chunk + [""] * (next(bk for bk in (32, 64, 128, 256) if bk >= n_real) - n_real)
+            t = tok(chunk, padding="max_length", truncation=True, max_length=32, return_tensors="pt").to(dev)
+            h = self.enc(input_ids=t["input_ids"], attention_mask=t["attention_mask"]).last_hidden_state
+            pooled = (h * t["attention_mask"].unsqueeze(-1)).sum(1) / t["attention_mask"].sum(1, keepdim=True)
+            for s, v in zip(chunk[:n_real], pooled[:n_real], strict=True):
+                out[s] = v
+        return torch.stack([out[s] for s in texts])
+
+    def forward(self, items: list[dict], tok, max_len: int = 4096) -> list[dict[str, torch.Tensor]]:
+        """Per request: {question: logits per step, "__halt__": halting weights per step and question}."""
+        dev = next(self.parameters()).device
+        prompts = [Prompt(tok, it, max_len) for it in items]
+        L = max(256, -(-max(len(e.ids) for e in prompts) // 256) * 256)
+        fdata = [self._facts(it, e, L) for it, e in zip(items, prompts, strict=True)]
+        hid = self._encode(prompts, dev, L)
+        late_texts, late_index = [], []
+        for b, (it, e) in enumerate(zip(items, prompts, strict=True)):
+            for n in e.late:
+                for k, desc in it["questions"][n]["criteria"].items():
+                    late_index.append((b, n))
+                    late_texts.append(f"{k}: {desc}")
+        late_emb = self.o_late(self._encode_options(late_texts, tok, dev)) if late_texts else None
+        late_map: dict[tuple[int, str], list[int]] = {}
+        for j, key in enumerate(late_index):
+            late_map.setdefault(key, []).append(j)
+        outs = []
+        for b, it in enumerate(items):
+            e, h = prompts[b], hid[b]
+            nq = len(e.names)
+            rows = [e.tokens_in(*sp) or [0] for sp in e.qspans]
+            orows, owner = [], []
+            for i, n in enumerate(e.names):
+                for sp in e.ospans.get(n, []):
+                    orows.append(e.tokens_in(*sp) or [0])
+                    owner.append(i)
+            pm = torch.zeros(nq + len(orows), L)
+            for r_, toks in enumerate(rows + orows):
+                pm[r_, toks] = 1.0 / len(toks)
+            pooled = pm.to(dev) @ h
+            qtypes = torch.tensor([QTYPES[it["questions"][n]["type"]] for n in e.names], device=dev)
+            q = self.q_proj(pooled[:nq]) + self.qid.weight[:nq] + self.qtype(qtypes)
+            o_joint = self.o_proj(pooled[nq:]) if orows else None
+            valid = torch.zeros(L, dtype=torch.bool, device=dev)
+            valid[: len(e.ids)] = True
+            att = torch.softmax((q @ self.ev_k(h).T / math.sqrt(D)).masked_fill(~valid.unsqueeze(0), float("-inf")), -1)
+            q = self.ev_out(torch.cat([q, att @ self.ev_v(h)], -1))  # question + evidence from the document
+            F = fq_t = fres_t = None
+            if fdata[b] is not None:
+                feats, smat, tags, res_v = fdata[b]
+                span_h = (smat.to(dev) @ h).view(len(tags), 2 * self.H)
+                F = self.fact_in(torch.cat([feats.to(dev), span_h], 1)) + self.qid.weight[torch.tensor(tags, device=dev)]
+                fq_t = torch.tensor(tags, device=dev)
+                fres_t = torch.tensor(res_v, device=dev, dtype=torch.float32)
+            x = (torch.cat([q, F]) if F is not None else q).unsqueeze(0)
+            zs, lams = [], []
+            for _ in range(STEPS):  # recursive refinement with learned halting
+                x = self.core(x)
+                zs.append(q + self.core_gate * x[0, :nq])
+                lams.append(torch.sigmoid(self.halt(zs[-1])).squeeze(-1))
+            remain, ps = torch.ones(nq, device=dev), []
+            for k, lam in enumerate(lams):
+                ps.append(remain * (lam if k < len(lams) - 1 else torch.ones_like(lam)))
+                remain = remain * (1 - lam)
+            p_halt = torch.stack(ps)
+            dec_ok = dec_bits = None
+            if F is not None:  # each question reads the facts that belong to it
+                own = (fq_t.unsqueeze(0) == torch.arange(nq, device=dev).unsqueeze(1)) | (fq_t == ALL_QUESTIONS).unsqueeze(0)
+                has = own.any(1, keepdim=True)
+                zs = [
+                    z
+                    + has.float()
+                    * self.fact_mix(
+                        torch.softmax((self.fact_q(z) @ F.T / math.sqrt(D)).masked_fill(~own, float("-inf")), -1).nan_to_num(0.0) @ F
+                    )
+                    for z in zs
+                ]
+                ft = fdata[b][0].to(dev)
+                kind = ft[:, FACT_DIM : FACT_DIM + len(FACT_KINDS)]
+                dec_ok = (
+                    (kind[:, FACT_KINDS.index("cmp")] > 0)
+                    | (kind[:, FACT_KINDS.index("pair")] > 0)
+                    | (kind[:, FACT_KINDS.index("sum")] > 0) & (ft[:, 4] > 0)
+                )
+                dec_bits = ft[:, 0]
+            res: dict[str, torch.Tensor] = {}
+            for i, n in enumerate(e.names):
+                qd = it["questions"][n]
+                if qd["type"] == "noul":
+                    o = o_joint[torch.tensor([j for j, w in enumerate(owner) if w == i], device=dev)]
+                else:
+                    o = (
+                        late_emb[torch.tensor(late_map[b, n], device=dev)]
+                        if n in e.late
+                        else o_joint[torch.tensor([j for j, w in enumerate(owner) if w == i], device=dev)]
+                    )
+                    ofe = torch.tensor(option_feats(it["state"], list(qd["criteria"])), dtype=torch.float32, device=dev)
+                step_logits = []
+                for z in zs:
+                    zq = z[i]
+                    zz = zq.unsqueeze(0).expand(o.shape[0], -1)
+                    lg = self.opt_mlp(torch.cat([zz, o, zz * o], -1)).squeeze(-1)
+                    if qd["type"] == "noul":
+                        base = (lg[0] - lg[1]).reshape(1)
+                        if dec_ok is not None:
+                            cand = dec_ok & ((fq_t == i) | (fq_t == ALL_QUESTIONS))
+                            if bool(cand.any()):
+                                base = self._exact_yes_no(base, zq, F[cand], dec_bits[cand])
+                        step_logits.append(base)
+                        continue
+                    lex = ofe.clone()
+                    lex[:, 2] = 0
+                    op = self.opt_lex(lex).squeeze(-1)
+                    numeric = ofe[:, 4] > 0
+                    if F is not None and bool(numeric.any()):  # numeric options matched against computed results
+                        sel = (fq_t == i) & ~torch.isnan(fres_t)
+                        if bool(sel.any()):
+                            a_ = torch.softmax(F[sel] @ self.fact_att(zq) / math.sqrt(D), 0)
+                            match = torch.exp(-(ofe[:, 2:3] - fres_t[sel].unsqueeze(0)).abs())
+                            op = op + self.num_w * numeric.float() * (match @ a_) * 10
+                    step_logits.append(lg + torch.sigmoid(self.op_gate(zq)) * op)
+                res[n] = torch.stack(step_logits)
+            res["__halt__"] = p_halt
+            outs.append(res)
+        return outs
+
+
+def probabilities(model: BevNetwork, out: dict[str, torch.Tensor], qtype: str, name: str) -> torch.Tensor:
+    """Calibrated probabilities of one question: per step, temperature-scaled; steps combined by their halting weights."""
+    i = [k for k in out if k != "__halt__"].index(name)
+    w = out["__halt__"][:, i : i + 1]
+    t = model.temp[qtype]
+    if qtype == "noul":
+        return (w.squeeze(-1) * torch.sigmoid(out[name].squeeze(-1) / t)).sum().reshape(1)
+    return (w * torch.softmax(out[name] / t, -1)).sum(0)
+
+
+def predict(model: BevNetwork, tok, items: list[dict], batch: int, max_len: int) -> list[dict]:
+    """[{"id", "answers": {question: {"pred", "probs"}}}] for items {"id", "state", "questions"}."""
+    rows = []
+    for k in range(0, len(items), batch):
+        chunk = items[k : k + batch]
+        for it, out in zip(chunk, model(chunk, tok, max_len), strict=True):
+            ans = {}
+            for n, q in it["questions"].items():
+                p = probabilities(model, out, q["type"], n).float().cpu().numpy()
+                if q["type"] == "noul":
+                    ans[n] = {"pred": int(p[0] >= 0.5), "probs": {"1": float(p[0]), "0": float(1 - p[0])}}
+                else:
+                    opts = list(q["criteria"])
+                    ans[n] = {"pred": opts[int(p.argmax())], "probs": {o: float(v) for o, v in zip(opts, p, strict=True)}}
+            rows.append({"id": it["id"], "answers": ans})
+    return rows