solvi 0.1.0__py3-none-any.whl

solvi/__init__.py

@@ -0,0 +1,6 @@
+"""solvi — a decision-system builder: a catalog of functions and checks, questions with typed answers, a strategist that
+assembles the flow, execution with computed_state, a hash chain and independent replay, and an answer head that trains in milliseconds."""
+from .core import Answer, AnswerType, Catalog, Question, Quote
+from .system import Response, System
+
+__all__ = ["Answer", "AnswerType", "Catalog", "Question", "Quote", "Response", "System"]

solvi/core.py

@@ -0,0 +1,100 @@
+"""solvi — a catalog of functions and checks, questions with typed answers.
+
+A part's contract comes from its signature: argument names are the facts it reads, the function name is the fact it sets.
+Part kinds: extract (pull a value from text, with a quote), fn (computation), check (test → bool), rule (answer to a question)."""
+from __future__ import annotations
+
+import inspect
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+
+@dataclass
+class Quote:
+    """A value extracted from text, with its location (doc[start:end] is the supporting quote)."""
+    value: Any
+    start: int
+    end: int
+    source: str = "doc"
+    confidence: float = 1.0
+
+
+@dataclass
+class AnswerType:
+    kind: str                       # yes_no | choice
+    options: list
+
+    def normalize(self, v):
+        if self.kind == "yes_no" and isinstance(v, bool):
+            return "yes" if v else "no"
+        if v not in self.options:
+            raise ValueError(f"answer {v!r} is not one of {self.options}")
+        return v
+
+
+class Answer:
+    @staticmethod
+    def yes_no() -> AnswerType:
+        return AnswerType("yes_no", ["yes", "no"])
+
+    @staticmethod
+    def choice(options) -> AnswerType:
+        return AnswerType("choice", list(options))
+
+
+@dataclass
+class Question:
+    name: str
+    text: str
+    answer: AnswerType
+    checkpoints: list = field(default_factory=list)   # parts required in every flow for this question
+    uses: list | None = None                          # hint to the strategist: which facts matter (when there is no rule or fit)
+
+
+@dataclass
+class Part:
+    kind: str                       # extract | fn | check | rule
+    name: str                       # fact it sets (for rule: "answer:<question>")
+    inputs: list
+    func: Callable
+    doc: str = ""
+    hard: bool = False              # check only: hard check
+    then: dict = field(default_factory=dict)   # check only: {question: answer} when the check is false
+    question: str | None = None     # rule only
+
+
+class Catalog:
+    """Everything the system can do; the strategist decides which parts each question needs."""
+
+    def __init__(self):
+        self.parts: dict[str, Part] = {}
+        self.rules: dict[str, Part] = {}
+
+    def _add(self, kind, f, **kw):
+        sig = inspect.signature(f)
+        p = Part(kind=kind, name=f.__name__, inputs=list(sig.parameters), func=f, doc=(f.__doc__ or "").strip(), **kw)
+        if kind == "rule":
+            p.name = "answer:" + p.question
+            self.rules[p.question] = p
+        else:
+            if p.name in self.parts:
+                raise ValueError(f"part {p.name} is already in the catalog")
+            self.parts[p.name] = p
+        return f
+
+    def extract(self, f):
+        return self._add("extract", f)
+
+    def fn(self, f):
+        return self._add("fn", f)
+
+    def check(self, f=None, *, hard=False, then=None):
+        if f is None:
+            return lambda g: self._add("check", g, hard=hard, then=then or {})
+        return self._add("check", f)
+
+    def rule(self, question):
+        return lambda f: self._add("rule", f, question=question)
+
+    def producer(self, fact):
+        return self.parts.get(fact)

solvi/extract_long.py

@@ -0,0 +1,224 @@
+"""@extract for real-world, general-purpose documents: a field is defined by its DESCRIPTION, the document may be long
+(windows), and the field may be absent ("no answer"). ModernBERT: input "field description [SEP] document window", two pointer
+heads; "no answer" is position 0 (the special token). Prediction: the best span across all windows; an answer exists if its
+score is above the field's threshold (tuned on held-out examples). Also works for fields unseen in training, from the
+description alone ("new field")."""
+from __future__ import annotations
+
+import math
+import random
+import time
+
+import numpy as np
+
+from .core import Quote
+
+
+class LongSpanExtractor:
+    def __init__(self, model_name="answerdotai/ModernBERT-large", max_len=1024, stride=128, max_span=96, device=None):
+        """stride: overlap between adjacent windows in tokens (as in transformers); max_span: maximum answer length in tokens."""
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+        self.torch = torch
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tok = AutoTokenizer.from_pretrained(model_name)
+        self.enc = AutoModel.from_pretrained(model_name).to(self.device)
+        self.head = torch.nn.Linear(self.enc.config.hidden_size, 2).to(self.device)
+        self.max_len, self.stride, self.max_span = max_len, stride, max_span
+        self.thr = {}
+        self.thr_default = 0.0                   # threshold for fields without their own examples (new field by description)
+        self.model_name = model_name
+        self._cache = {}
+        self._tok_key, self._tok_val = None, None
+
+    def _windows(self, desc, text):
+        """Windows "[CLS] description [SEP] text chunk [SEP]" overlapping by stride; split manually (the ModernBERT tokenizer's
+        overflow returns only one extra window). → {input_ids, attention_mask, offset_mapping, ctx}."""
+        key = hash(text)
+        if self._tok_key != key:
+            t = self.tok(text, add_special_tokens=False, return_offsets_mapping=True)
+            self._tok_key, self._tok_val = key, (t["input_ids"], t["offset_mapping"])
+        ids, offs = self._tok_val
+        d = self.tok(desc, add_special_tokens=False)["input_ids"][:64]
+        cls, sep, pad = self.tok.cls_token_id, self.tok.sep_token_id, self.tok.pad_token_id
+        room = self.max_len - len(d) - 3
+        step = max(1, room - self.stride)
+        out = {"input_ids": [], "attention_mask": [], "offset_mapping": [], "ctx": []}
+        for a in range(0, max(1, len(ids)), step):
+            chunk = ids[a:a + room]
+            seq = [cls] + d + [sep] + chunk + [sep]
+            n = len(seq)
+            out["input_ids"].append(seq + [pad] * (self.max_len - n))
+            out["attention_mask"].append([1] * n + [0] * (self.max_len - n))
+            out["offset_mapping"].append([(0, 0)] * (len(d) + 2) + list(offs[a:a + room]) + [(0, 0)] * (self.max_len - n + 1))
+            out["ctx"].append(list(range(len(d) + 2, len(d) + 2 + len(chunk))))
+            if a + room >= len(ids):
+                break
+        return out
+
+    def _ctx(self, enc, i):
+        return enc["ctx"][i]
+
+    def make_examples(self, items, neg_per_item=3, seed=0):
+        """items: [(text, description, (start, end) | None[, neg])] → training windows: every window with the answer + up to
+        neg_per_item (or the item's own neg) windows without it."""
+        rng = random.Random(seed)
+        out = []
+        for it in items:
+            text, desc, span = it[:3]
+            npi = it[3] if len(it) > 3 else neg_per_item
+            enc = self._windows(desc, text)
+            pos, neg = [], []
+            for i in range(len(enc["input_ids"])):
+                offs = enc["offset_mapping"][i]
+                ctx = self._ctx(enc, i)
+                if not ctx:
+                    continue
+                c0, c1 = offs[ctx[0]][0], offs[ctx[-1]][1]
+                if span is not None and c0 <= span[0] < c1:            # answer starts in this window (a long answer is clipped at the window edge)
+                    st = next((j for j in ctx if offs[j][1] > span[0]), ctx[0])
+                    en = next((j for j in reversed(ctx) if offs[j][0] < span[1]), st)
+                    pos.append((enc["input_ids"][i], enc["attention_mask"][i], st, min(max(st, en), st + self.max_span - 1)))
+                else:
+                    neg.append((enc["input_ids"][i], enc["attention_mask"][i], 0, 0))
+            rng.shuffle(neg)
+            out += pos + neg[:npi if pos or span is None else max(1, npi // 3)]
+        return out
+
+    def fit(self, items, epochs=3, lr=3e-5, bs=8, seed=0, neg_per_item=3, log=print):
+        torch = self.torch
+        random.seed(seed)
+        torch.manual_seed(seed)
+        ex = self.make_examples(items, neg_per_item, seed)
+        opt = torch.optim.AdamW([{"params": list(self.enc.parameters()), "lr": lr}, {"params": list(self.head.parameters()), "lr": lr * 10}],
+                                weight_decay=0.01)
+        total = epochs * math.ceil(len(ex) / bs)
+        sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(1, total // 10)) * max(0.0, 1 - s / total))
+        params = list(self.enc.parameters()) + list(self.head.parameters())
+        self.enc.train()
+        t0 = time.time()
+        for ep in range(epochs):
+            random.shuffle(ex)
+            # batches of windows of similar length (sorted within chunks of 64 batches), in random order
+            batches = []
+            for k in range(0, len(ex), 64 * bs):
+                part = sorted(ex[k:k + 64 * bs], key=lambda c: sum(c[1]))
+                batches += [part[b:b + bs] for b in range(0, len(part), bs)]
+            random.shuffle(batches)
+            tot = 0.0
+            for ch in batches:
+                L = max(sum(c[1]) for c in ch)                  # trim the batch to its longest window
+                ids = torch.tensor([c[0][:L] for c in ch], device=self.device)
+                att = torch.tensor([c[1][:L] for c in ch], device=self.device)
+                with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
+                    h = self.enc(input_ids=ids, attention_mask=att).last_hidden_state
+                    lg = self.head(h).float()
+                m = att.bool()
+                ls, le = lg[..., 0].masked_fill(~m, -1e4), lg[..., 1].masked_fill(~m, -1e4)
+                loss = (torch.nn.functional.cross_entropy(ls, torch.tensor([c[2] for c in ch], device=self.device)) +
+                        torch.nn.functional.cross_entropy(le, torch.tensor([c[3] for c in ch], device=self.device))) / 2
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(params, 1.0)
+                opt.step()
+                sched.step()
+                tot += loss.item()
+            log(f"[extractL] epoch {ep + 1}/{epochs}: loss {tot / max(1, math.ceil(len(ex) / bs)):.4f}, windows {len(ex)}, {time.time() - t0:.0f} s")
+        self.enc.eval()
+        self._cache = {}
+        return self
+
+    def predict(self, text, desc, bs=8):
+        """→ (start, end, span score, "no answer" score) — the best span across all windows."""
+        key = (hash(text), desc)
+        if key in self._cache:
+            return self._cache[key]
+        torch = self.torch
+        max_span = self.max_span
+        enc = self._windows(desc, text)
+        best, null = (0, 0, -1.0), 0.0
+        with torch.no_grad():
+            n = len(enc["input_ids"])
+            for b in range(0, n, bs):
+                L = max(sum(a) for a in enc["attention_mask"][b:b + bs])
+                ids = torch.tensor([x[:L] for x in enc["input_ids"][b:b + bs]], device=self.device)
+                att = torch.tensor([x[:L] for x in enc["attention_mask"][b:b + bs]], device=self.device)
+                with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
+                    lg = self.head(self.enc(input_ids=ids, attention_mask=att).last_hidden_state).float()
+                for r in range(lg.shape[0]):
+                    i = b + r
+                    ctx = self._ctx(enc, i)
+                    if not ctx:
+                        continue
+                    offs = enc["offset_mapping"][i]
+                    ps = torch.softmax(lg[r, :, 0].masked_fill(~att[r].bool(), -1e4), -1).cpu().numpy()
+                    pe = torch.softmax(lg[r, :, 1].masked_fill(~att[r].bool(), -1e4), -1).cpu().numpy()
+                    null = max(null, float(ps[0] * pe[0]))
+                    c0, c1 = ctx[0], ctx[-1] + 1
+                    sc = np.triu(np.outer(ps[c0:c1], pe[c0:c1])) - np.triu(np.outer(ps[c0:c1], pe[c0:c1]), max_span)
+                    s, e = np.unravel_index(int(sc.argmax()), sc.shape)
+                    if sc[s, e] > best[2]:
+                        best = (offs[c0 + s][0], offs[c0 + e][1], float(sc[s, e]))
+        st, en = int(best[0]), int(best[1])
+        while st < en and text[st].isspace():              # BPE offsets include the leading space
+            st += 1
+        out = (st, en, best[2], null)
+        self._cache[key] = out
+        return out
+
+    def tune_threshold(self, name, items, grid=(0.0003, 0.001, 0.003, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8)):
+        """Per-field "answer present" threshold from held-out examples: maximizes present/absent decision accuracy."""
+        best = None
+        for t in grid:
+            acc = np.mean([(self.predict(x, d)[2] >= t) == (sp is not None) for x, d, sp in items])
+            if best is None or acc > best[0]:
+                best = (acc, t)
+        self.thr[name] = best[1]
+        return best
+
+    def tune_default_threshold(self, items, grid=(0.0003, 0.001, 0.003, 0.01, 0.02, 0.05, 0.1, 0.2, 0.3, 0.4, 0.5)):
+        """One threshold from held-out examples of many fields — used for fields that have no examples of their own."""
+        best = None
+        for t in grid:
+            acc = np.mean([(self.predict(x, d)[2] >= t) == (sp is not None) for x, d, sp in items])
+            if best is None or acc > best[0]:
+                best = (acc, t)
+        self.thr_default = best[1]
+        return best
+
+    def save(self, path):
+        """Encoder weights (bf16 safetensors), tokenizer, span head and settings into a directory."""
+        import json
+        from pathlib import Path
+        Path(path).mkdir(parents=True, exist_ok=True)
+        self.enc.to(self.torch.bfloat16).save_pretrained(path)
+        self.enc.to(self.torch.float32)
+        self.tok.save_pretrained(path)
+        self.torch.save(self.head.state_dict(), f"{path}/span_head.pt")
+        json.dump({"max_len": self.max_len, "stride": self.stride, "max_span": self.max_span, "thr_default": self.thr_default,
+                   "thr": self.thr, "base_model": self.model_name}, open(f"{path}/solvi_extract.json", "w"), indent=1)
+
+    @classmethod
+    def load(cls, path, device=None):
+        """From a directory written by save(), or a Hugging Face model id (downloaded once and cached)."""
+        import json
+        import os
+        if not os.path.isdir(path):
+            from huggingface_hub import snapshot_download
+            path = snapshot_download(path)
+        cfg = json.load(open(f"{path}/solvi_extract.json"))
+        ex = cls(path, max_len=cfg["max_len"], stride=cfg["stride"], max_span=cfg["max_span"], device=device)
+        ex.enc.to(ex.torch.float32)
+        ex.head.load_state_dict(ex.torch.load(f"{path}/span_head.pt", map_location=ex.device))
+        ex.thr_default, ex.thr = cfg["thr_default"], cfg.get("thr", {})
+        return ex
+
+    def field(self, name, desc):
+        def f(doc):
+            s, e, sc, _ = self.predict(doc, desc)
+            if sc < self.thr.get(name, self.thr_default):
+                return Quote("", 0, 0, confidence=1 - sc)
+            return Quote(doc[s:e], s, e, confidence=sc)
+        f.__name__ = name
+        f.__doc__ = desc
+        return f

solvi/extract_model.py

@@ -0,0 +1,105 @@
+"""Model-backed @extract: ModernBERT + two pointer heads (span start and end) that "cut a field's value out of the text given
+its description" (extractive QA). The answer is always a span of the text itself, so it comes with a quote and offsets by
+construction. The field description is the @extract function's docstring (a field can be added without retraining if its
+description resembles the trained ones)."""
+from __future__ import annotations
+
+import math
+import random
+import time
+
+import numpy as np
+
+
+class SpanExtractor:
+    def __init__(self, model_name="answerdotai/ModernBERT-large", max_len=1024, device=None):
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+        self.torch = torch
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tok = AutoTokenizer.from_pretrained(model_name)
+        self.enc = AutoModel.from_pretrained(model_name).to(self.device)
+        self.head = torch.nn.Linear(self.enc.config.hidden_size, 2).to(self.device)
+        self.max_len = max_len
+
+    def _batch(self, items):
+        """items: [(text, description, (start, end) | None)] → tensors and target positions in tokens."""
+        enc = self.tok([d for _, d, _ in items], [t for t, _, _ in items], truncation="only_second", max_length=self.max_len,
+                       padding=True, return_offsets_mapping=True, return_tensors="pt")
+        starts, ends, ctx_masks = [], [], []
+        for i, (_, _, span) in enumerate(items):
+            seq = enc.sequence_ids(i)
+            offs = enc["offset_mapping"][i].tolist()
+            ctx = [j for j, s in enumerate(seq) if s == 1]
+            m = [s == 1 for s in seq]
+            ctx_masks.append(m)
+            if span is None:
+                starts.append(0)
+                ends.append(0)
+                continue
+            s_c, e_c = span
+            st = next((j for j in ctx if offs[j][0] <= s_c < offs[j][1] or offs[j][0] >= s_c), ctx[0])
+            en = next((j for j in reversed(ctx) if offs[j][0] < e_c), st)
+            starts.append(st)
+            ends.append(max(en, st))
+        return enc, starts, ends, ctx_masks
+
+    def _logits(self, enc):
+        out = self.enc(input_ids=enc["input_ids"].to(self.device), attention_mask=enc["attention_mask"].to(self.device)).last_hidden_state
+        lg = self.head(out)
+        return lg[..., 0], lg[..., 1]
+
+    def fit(self, examples, epochs=4, lr=3e-5, bs=8, seed=0, log=print):
+        torch = self.torch
+        random.seed(seed)
+        torch.manual_seed(seed)
+        ex = [e for e in examples if e[2] is not None]
+        params = list(self.enc.parameters()) + list(self.head.parameters())
+        opt = torch.optim.AdamW([{"params": list(self.enc.parameters()), "lr": lr}, {"params": list(self.head.parameters()), "lr": lr * 10}],
+                                weight_decay=0.01)
+        total = epochs * math.ceil(len(ex) / bs)
+        sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(1, total // 10)) * max(0.0, 1 - s / total))
+        self.enc.train()
+        t0 = time.time()
+        for ep in range(epochs):
+            random.shuffle(ex)
+            tot = 0.0
+            for b in range(0, len(ex), bs):
+                enc, st, en, m = self._batch(ex[b:b + bs])
+                with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
+                    ls, le = self._logits(enc)
+                mask = torch.tensor(m, device=self.device)
+                ls = ls.float().masked_fill(~mask, -1e4)
+                le = le.float().masked_fill(~mask, -1e4)
+                loss = (torch.nn.functional.cross_entropy(ls, torch.tensor(st, device=self.device)) +
+                        torch.nn.functional.cross_entropy(le, torch.tensor(en, device=self.device))) / 2
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(params, 1.0)
+                opt.step()
+                sched.step()
+                tot += loss.item()
+            log(f"[extract] epoch {ep + 1}/{epochs}: loss {tot / max(1, math.ceil(len(ex) / bs)):.4f}, {time.time() - t0:.0f} s")
+        self.enc.eval()
+        return self
+
+    def predict(self, items, bs=16, max_span=64):
+        """items: [(text, description)] → [(start, end, confidence)] in text characters."""
+        torch = self.torch
+        out = []
+        with torch.no_grad():
+            for b in range(0, len(items), bs):
+                ch = [(t, d, None) for t, d in items[b:b + bs]]
+                enc, _, _, m = self._batch(ch)
+                with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
+                    ls, le = self._logits(enc)
+                mask = torch.tensor(m, device=self.device)
+                ps = torch.softmax(ls.float().masked_fill(~mask, -1e4), -1).cpu().numpy()
+                pe = torch.softmax(le.float().masked_fill(~mask, -1e4), -1).cpu().numpy()
+                for i in range(len(ch)):
+                    offs = enc["offset_mapping"][i].tolist()
+                    L = len(offs)
+                    sc = np.triu(np.outer(ps[i, :L], pe[i, :L])) - np.triu(np.outer(ps[i, :L], pe[i, :L]), max_span)
+                    s, e = np.unravel_index(int(sc.argmax()), sc.shape)
+                    out.append((int(offs[s][0]), int(offs[e][1]), float(sc[s, e])))
+        return out

solvi/extract_multi.py

@@ -0,0 +1,155 @@
+"""Single-pass @extract: ModernBERT reads the document once, with a pair of pointer heads (start / end) per field.
+For solvi: extractor.field(name) returns a function doc → Quote; all fields of a document come from one pass (cached by text)."""
+from __future__ import annotations
+
+import hashlib
+import math
+import random
+import time
+
+import numpy as np
+
+from .core import Quote
+
+
+class MultiSpanExtractor:
+    def __init__(self, fields, model_name="answerdotai/ModernBERT-large", max_len=1024, device=None):
+        import torch
+        from transformers import AutoModel, AutoTokenizer
+        self.torch = torch
+        self.fields = list(fields)
+        self.device = device or ("cuda" if torch.cuda.is_available() else "cpu")
+        self.tok = AutoTokenizer.from_pretrained(model_name)
+        self.enc = AutoModel.from_pretrained(model_name).to(self.device)
+        self.head = torch.nn.Linear(self.enc.config.hidden_size, 2 * len(self.fields)).to(self.device)
+        self.max_len = max_len
+        self._cache = {}
+        self.temp = {f: 1.0 for f in self.fields}
+
+    def _enc(self, texts):
+        return self.tok(texts, truncation=True, max_length=self.max_len, padding=True, return_offsets_mapping=True, return_tensors="pt")
+
+    def _logits(self, enc):
+        h = self.enc(input_ids=enc["input_ids"].to(self.device), attention_mask=enc["attention_mask"].to(self.device)).last_hidden_state
+        return self.head(h)                                          # [B, L, 2F]
+
+    @staticmethod
+    def _tok_span(offs, span):
+        s_c, e_c = span
+        idx = [j for j, (a, b) in enumerate(offs) if b > a]
+        st = next((j for j in idx if offs[j][1] > s_c), idx[0])
+        en = next((j for j in reversed(idx) if offs[j][0] < e_c), st)
+        return st, max(st, en)
+
+    def fit(self, docs, spans, epochs=4, lr=3e-5, bs=8, seed=0, log=print):
+        """docs: [text]; spans: [{field: (start, end) | None}]."""
+        torch = self.torch
+        random.seed(seed)
+        torch.manual_seed(seed)
+        idx = list(range(len(docs)))
+        opt = torch.optim.AdamW([{"params": list(self.enc.parameters()), "lr": lr}, {"params": list(self.head.parameters()), "lr": lr * 10}],
+                                weight_decay=0.01)
+        total = epochs * math.ceil(len(idx) / bs)
+        sched = torch.optim.lr_scheduler.LambdaLR(opt, lambda s: min(1.0, (s + 1) / max(1, total // 10)) * max(0.0, 1 - s / total))
+        params = list(self.enc.parameters()) + list(self.head.parameters())
+        self.enc.train()
+        t0 = time.time()
+        F = len(self.fields)
+        for ep in range(epochs):
+            random.shuffle(idx)
+            tot = 0.0
+            for b in range(0, len(idx), bs):
+                ch = idx[b:b + bs]
+                enc = self._enc([docs[i] for i in ch])
+                with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
+                    lg = self._logits(enc).float()
+                mask = enc["attention_mask"].to(self.device).bool()
+                loss, n = 0.0, 0
+                for fi, f in enumerate(self.fields):
+                    tgt_s, tgt_e, rows = [], [], []
+                    for r, i in enumerate(ch):
+                        sp = spans[i].get(f)
+                        if sp is None:
+                            continue
+                        s, e = self._tok_span(enc["offset_mapping"][r].tolist(), sp)
+                        tgt_s.append(s)
+                        tgt_e.append(e)
+                        rows.append(r)
+                    if not rows:
+                        continue
+                    ls = lg[rows, :, 2 * fi].masked_fill(~mask[rows], -1e4)
+                    le = lg[rows, :, 2 * fi + 1].masked_fill(~mask[rows], -1e4)
+                    loss = loss + torch.nn.functional.cross_entropy(ls, torch.tensor(tgt_s, device=self.device)) + \
+                        torch.nn.functional.cross_entropy(le, torch.tensor(tgt_e, device=self.device))
+                    n += 2
+                if n == 0:
+                    continue
+                loss = loss / n
+                opt.zero_grad(set_to_none=True)
+                loss.backward()
+                torch.nn.utils.clip_grad_norm_(params, 1.0)
+                opt.step()
+                sched.step()
+                tot += loss.item()
+            log(f"[extract1] epoch {ep + 1}/{epochs}: loss {tot / max(1, math.ceil(len(idx) / bs)):.4f}, {time.time() - t0:.0f} s")
+        self.enc.eval()
+        self._cache = {}
+        return self
+
+    def predict_doc(self, text, max_span=64):
+        """→ {field: (start, end, confidence)} in a single pass."""
+        key = hashlib.sha1(text.encode()).hexdigest()
+        if key in self._cache:
+            return self._cache[key]
+        torch = self.torch
+        with torch.no_grad():
+            enc = self._enc([text])
+            with torch.autocast(self.device, dtype=torch.bfloat16, enabled=self.device == "cuda"):
+                lg = self._logits(enc).float()[0]
+        offs = enc["offset_mapping"][0].tolist()
+        valid = np.array([b > a for a, b in offs])
+        out = {}
+        for fi, f in enumerate(self.fields):
+            zs = lg[:, 2 * fi].cpu().numpy()
+            ze = lg[:, 2 * fi + 1].cpu().numpy()
+            zs[~valid], ze[~valid] = -1e4, -1e4
+            T = self.temp[f]
+            ps = np.exp((zs - zs.max()) / T)
+            ps /= ps.sum()
+            pe = np.exp((ze - ze.max()) / T)
+            pe /= pe.sum()
+            L = len(offs)
+            sc = np.triu(np.outer(ps, pe)) - np.triu(np.outer(ps, pe), max_span)
+            s, e = np.unravel_index(int(sc.argmax()), sc.shape)
+            out[f] = (int(offs[s][0]), int(offs[e][1]), float(sc[s, e]))
+        if len(self._cache) > 5000:
+            self._cache = {}
+        self._cache[key] = out
+        return out
+
+    def field(self, name):
+        def f(doc):
+            s, e, c = self.predict_doc(doc)[name]
+            return Quote(doc[s:e], s, e, confidence=c)
+        f.__name__ = name
+        return f
+
+    def fit_temperature(self, docs, gold_ok, grid=(0.5, 0.75, 1.0, 1.5, 2.0, 3.0, 5.0)):
+        """Per-field temperature fitted on held-out documents: minimizes log loss of confidence vs. whether the span matched gold.
+        gold_ok(field, doc_index, (start, end)) → bool."""
+        for f in self.fields:
+            best = None
+            for T in grid:
+                self.temp[f] = T
+                self._cache = {}
+                nll = 0.0
+                for i, d in enumerate(docs):
+                    s, e, c = self.predict_doc(d)[f]
+                    y = gold_ok(f, i, (s, e))
+                    c = min(max(c, 1e-6), 1 - 1e-6)
+                    nll -= math.log(c if y else 1 - c)
+                if best is None or nll < best[0]:
+                    best = (nll, T)
+            self.temp[f] = best[1]
+        self._cache = {}
+        return dict(self.temp)