upstream/vela3_encoder.py

"""vela3_encoder.py: inference for vela3-fast (UltraBERT-380M encoder).

Three networks on one frozen backbone:
  DecisionEncoder  the decision backbone (UBD3 stage 1b, soup a+c) and its heads: every choice / noul decision that is not a router
                   question. This class and model.safetensors are unchanged from the decision-only preview (8667a7bb), so these
                   answers are the preview's, bit for bit.
  span adapter     LoRA r64 on all 24 layers + the top 4 layers + span head (span_adapter.safetensors): every span question.
  router adapter   LoRA r64 on all 24 layers + choice / set head (router_adapter.safetensors): choice / set questions whose
                   instructions are one of the router texts in router_questions.json.

Decision rendering (one encoder input per question, exactly as in training, ub_d3_common.to_row):
    [BOS] [Q] instructions [O] key_1: description_1 [O] key_2 ... [SEP_SCHEMA] [SEG_user] state [EOS]
state / instructions: strings as is, JSON objects as compact JSON; noul = the 2-option choice {no, yes}, answer = p(yes).
Readout (both, as trained): cosine(w_o(norm(H[O_i])) + w_q(norm(H[Q])), w_p(norm(mean H over [SEG_user] .. state))) / tau
+ MLP(H[O_i]); softmax over the question's options. Nothing is truncated: an input over max_len raises Unsupported.
Span and router passes: vela3_schema.py (the Vela 2.0 layout, 8,192-token windows). Evidence questions (EV-FIX, 18 Oct 2026,
DIAG_04B_EV_HALU.md): one window up to span.evidence.max_len (32,768) on the same span adapter, and a threshold by target length
(span.evidence.buckets: <= 2,048 tokens the short-dev pick, longer the QASPER-validation pick).
Task adapters (round 1, 19 Oct 2026; vela3_config.json span.task_adapters): relation questions ("Which spans are the <relation> of <head>?",
"Which spans are related to <head> by these relations?", ...) -> the relation adapter, per-form thresholds (pair = one relation, multi =
several); open-label entity questions (names only) stay on the span adapter.

    from vela3_encoder import Vela3Encoder
    m = Vela3Encoder.load("path/to/this/folder")              # cuda if available
    m.system_one(state, {"q": {"type": "choice", "instructions": "...", "criteria": {"a": "...", "b": None}},
                         "ents": {"type": "span", "instructions": "Which spans are named entities of these types?",
                                  "criteria": {"person": None, "company": None}}})
    m.predict(state, questions)    # decisions only, the preview's API and answers (the Decision Index engine uses it)
"""
import json, math
from pathlib import Path

import numpy as np
import torch
import torch.nn as nn

from ub_d3_common import to_row

MARKERS = ["[Q]", "[O]", "[ABS]", "[SEP_SCHEMA]", "[SEG_user]", "[SEG_answer]", "[E]", "[SEG_context]"]


class Unsupported(ValueError):
    pass


def opt_text(name, desc):
    return name if not desc else f"{name}: {desc}"


class DecisionEncoder(nn.Module):
    def __init__(self, cfg, proj_dim=256):
        super().__init__()
        from transformers import AutoModel
        self.encoder = AutoModel.from_config(cfg, trust_remote_code=True, attn_implementation="sdpa", dtype=torch.float32)
        h = cfg.hidden_size
        self.norm = nn.LayerNorm(h)
        self.w_o, self.w_q, self.w_p = nn.Linear(h, proj_dim), nn.Linear(h, proj_dim), nn.Linear(h, proj_dim)
        self.log_tau = nn.Parameter(torch.tensor(math.log(0.05)))
        self.cls_mlp = nn.Sequential(nn.Linear(h, 2 * h), nn.ReLU(), nn.Dropout(0.1), nn.Linear(2 * h, 1))

    @classmethod
    def load(cls, path, device=None, max_len=32768, token_budget=65536):
        from safetensors.torch import load_file
        from transformers import AutoConfig, AutoTokenizer
        path = Path(path); device = device or ("cuda" if torch.cuda.is_available() else "cpu")
        cfg = AutoConfig.from_pretrained(str(path), trust_remote_code=True)
        m = cls(cfg, cfg.decision["proj_dim"])
        m.load_state_dict(load_file(str(path / "model.safetensors")), strict=True)
        m = m.to(device).eval()
        # text is tokenised with the base UltraBERT tokenizer (as in training); the 8 marker tokens are the rows appended after it
        m.tok = AutoTokenizer.from_pretrained(str(path), trust_remote_code=True)
        m.mk = {t: int(cfg.decision["marker_ids"][t]) for t in MARKERS}
        assert len(m.tok) == cfg.decision["n_vocab_base"] == min(m.mk.values()), "tokenizer / marker mismatch"
        m.device, m.max_len, m.budget = device, int(max_len), int(token_budget)
        return m

    def _ids(self, s):
        return self.tok(s, add_special_tokens=False)["input_ids"] if s else []

    def encode(self, state, q, state_ids=None):
        """-> dict(ids, q_pos, opt_pos, names, pool) for one question, or raise Unsupported."""
        qq = to_row("q", state, q)["questions"][0]; st = to_row("q", state, q)["parts"][0]["text"]
        ids = [self.tok.bos_token_id, self.mk["[Q]"]] + self._ids(qq["text"]); opt_pos, names = [], []
        for n, d in qq["options"].items():
            opt_pos.append(len(ids)); names.append(n); ids += [self.mk["[O]"]] + self._ids(opt_text(n, d))
        ids.append(self.mk["[SEP_SCHEMA]"]); ps = len(ids)
        ids += [self.mk["[SEG_user]"]] + (state_ids if state_ids is not None else self._ids(st)); pe = len(ids); ids.append(self.tok.eos_token_id)
        if len(ids) > self.max_len:
            raise Unsupported(f"input longer than the {self.max_len}-token context")
        return dict(ids=ids, q_pos=1, opt_pos=opt_pos, names=names, pool=(ps, pe))

    def _logits(self, H, recs):
        H = H.float(); dev = H.device
        ob = torch.tensor([b for b, r in enumerate(recs) for _ in r["opt_pos"]], device=dev)
        op = torch.tensor([p for r in recs for p in r["opt_pos"]], device=dev)
        own = torch.tensor([k for k, r in enumerate(recs) for _ in r["opt_pos"]], device=dev)
        b = torch.arange(len(recs), device=dev)
        ps = torch.tensor([r["pool"][0] for r in recs], device=dev); pe = torch.tensor([r["pool"][1] for r in recs], device=dev)
        cs = nn.functional.pad(H.cumsum(1), (0, 0, 1, 0))
        pooled = (cs[b, pe] - cs[b, ps]) / (pe - ps).clamp(min=1).unsqueeze(-1).float()
        v = nn.functional.normalize(self.w_p(self.norm(pooled)), dim=-1)
        uq = self.w_q(self.norm(H[b, torch.tensor([r["q_pos"] for r in recs], device=dev)]))
        u = nn.functional.normalize(self.w_o(self.norm(H[ob, op])) + uq[own], dim=-1)
        lg = (u * v[own]).sum(-1) / self.log_tau.exp() + self.cls_mlp(H[ob, op]).squeeze(-1)
        return list(torch.split(lg, [len(r["opt_pos"]) for r in recs]))

    @torch.no_grad()
    def forward_recs(self, recs):
        """Batched by length under a token budget (as the evaluated engine); bf16 autocast on CUDA. -> list of prob arrays."""
        out = [None] * len(recs); order = sorted(range(len(recs)), key=lambda i: len(recs[i]["ids"])); i = 0
        amp = torch.autocast("cuda", dtype=torch.bfloat16, enabled=self.device == "cuda")
        while i < len(order):
            j = i + 1
            while j < len(order) and len(recs[order[j]]["ids"]) * (j - i + 1) <= self.budget:
                j += 1
            grp = [recs[k] for k in order[i:j]]; w = max(len(r["ids"]) for r in grp)
            x = torch.full((len(grp), w), self.tok.pad_token_id, dtype=torch.long); msk = torch.zeros_like(x)
            for r_, r in enumerate(grp):
                x[r_, :len(r["ids"])] = torch.tensor(r["ids"]); msk[r_, :len(r["ids"])] = 1
            with amp:
                H = self.encoder(input_ids=x.to(self.device), attention_mask=msk.to(self.device)).last_hidden_state
            for k, lg in zip(order[i:j], self._logits(H, grp)):
                z = lg.double().cpu().numpy(); p = np.exp(z - z.max()); out[k] = p / p.sum()
            i = j
        return out

    def predict(self, state, questions):
        """questions: {key: {type: choice | noul, instructions, criteria}} -> {key: answer} (choice: choice + probabilities; noul: p(yes))."""
        recs, keys, st = [], [], None
        for k, q in questions.items():
            if q["type"] not in ("choice", "noul"):
                raise Unsupported("Unsupported question type " + str(q["type"]))
            if st is None:
                st = self._ids(to_row("q", state, q)["parts"][0]["text"])
            recs.append(self.encode(state, q, st)); keys.append(k)
        ans = {}
        for k, r, p in zip(keys, recs, self.forward_recs(recs)):
            if questions[k]["type"] == "noul":
                ans[k] = {"type": "noul", "noul": float(p[r["names"].index("yes")])}
            else:
                pr = {n: float(v) for n, v in zip(r["names"], p)}
                ans[k] = {"type": "choice", "choice": max(pr, key=pr.get), "probabilities": pr}
        return ans, sum(len(r["ids"]) for r in recs)


# =====================================================================================================================
# vela3-fast: decisions (the preview's DecisionEncoder, unchanged) + gated span / router adapters
# =====================================================================================================================
import re as _re

ROLES = ("user", "context", "answer")
_REL_Q = _re.compile(r"\bspans? (?:are|is) the .+ (?:of|for)\b|\brelat(?:ed|ion|ions|es|ing)?\b|\b(?:head|tail)\b|\bstands? in\b"
                     r"|\bwith respect to\b|\bcounterpart\b|\bsuch that\b|\brole of\b")
_EV_Q = _re.compile(r"\b(?:evidence|support(?:s|ing|ed)?|cite[sd]?|citations?)\b")


def _norm(s):
    return " ".join(str(s or "").split())


def _parts(state):
    """System One state -> [(role, text)]: a string is the user part; an object of text parts uses the roles user / context / answer."""
    if isinstance(state, str):
        return [("user", state)]
    if isinstance(state, dict) and state and all(isinstance(v, str) for v in state.values()):
        bad = [k for k in state if k not in ROLES]
        if bad:
            raise ValueError(f"span / router state parts must be named {ROLES}, got {bad}")
        return list(state.items())
    raise ValueError("span and router questions need a text state or an object of text parts {user | context | answer: text}")


class Vela3Encoder:
    """vela3-fast. system_one(state, questions): choice / noul (decision backbone), router choice / set (router adapter), span
    (span adapter). predict(state, questions): the preview's decision API (decision backbone only)."""

    def __init__(self, path, device=None, max_len=32768, token_budget=65536, adapter_bf16=False, entity_route=None):
        """adapter_bf16: run the span / router passes under bf16 autocast (default fp32, as evaluated). Decisions: unchanged (bf16 autocast on GPU)."""
        import vela3_schema as VS
        path = Path(path)
        self.decision = DecisionEncoder.load(path, device=device, max_len=max_len, token_budget=token_budget)
        dev = self.decision.device
        self.cfg = json.loads((path / "vela3_config.json").read_text())
        rq = json.loads((path / "router_questions.json").read_text())["questions"]
        self.router_q = {_norm(t): s for t, s in rq.items()}
        span_net = VS.build_adapter(self.decision, path / "span_adapter.safetensors", dev)
        self.span = VS.SchemaEngine(path, span_net, dev, amp_bf16=adapter_bf16)
        ev_len = int(self.cfg["span"]["evidence"].get("max_len", self.span.max_len))  # EV-FIX: evidence engine, same net (no extra weights)
        evf = self.cfg["span"]["evidence"].get("adapter")  # EVFAST: optional dedicated evidence adapter (evidence questions only; missing file = off)
        ev_net = VS.build_adapter(self.decision, path / evf, dev) if evf and (path / evf).exists() else span_net
        self.span_ev = self.span if (ev_len == self.span.max_len and ev_net is span_net) else VS.SchemaEngine(path, ev_net, dev, tokenizer=self.span.tok, max_len=ev_len,
                                                                                      token_budget=max(ev_len, 32768), amp_bf16=adapter_bf16)
        self.task = {}  # per-task span adapters (span.task_adapters: {task: file}); a missing file = route off
        for t, f in (self.cfg["span"].get("task_adapters") or {}).items():
            if (path / f).exists():
                self.task[t] = VS.SchemaEngine(path, VS.build_adapter(self.decision, path / f, dev), dev, tokenizer=self.span.tok, amp_bf16=adapter_bf16)
        # entity route for names-only open-label entity questions: "ner" (the NER adapter, default) or "span" (the span adapter). The
        # end-to-end relation pipeline (entity step, then relation questions) runs with "span" (span.entity_route_relation_pipeline):
        # set VELA3_ENTITY_ROUTE=span, or Vela3Encoder.load(..., entity_route="span"), for that pipeline.
        import os as _os
        self.entity_route = entity_route or _os.environ.get("VELA3_ENTITY_ROUTE") or self.cfg["span"].get("entity_route", "ner")
        # R2c: a second span adapter for the router label sets (span.set_adapter: {file, label_sets, router_set_temperature,
        # thresholds_by_set, pii_buckets}): halu and PII questions are answered by it with its own dev-picked temperature / thresholds
        ss = self.cfg["span"].get("set_adapter"); self.span_set = None
        if ss and (path / ss["file"]).exists():
            self.span_set = VS.SchemaEngine(path, VS.build_adapter(self.decision, path / ss["file"], dev), dev, tokenizer=self.span.tok, amp_bf16=adapter_bf16)
        self.router = VS.SchemaEngine(path, VS.build_adapter(self.decision, path / "router_adapter.safetensors", dev), dev, tokenizer=self.span.tok, amp_bf16=adapter_bf16)
        # DOM04B: optional per-task router adapters (vela3_config.json router.task_adapters: {task: {file, ids}}): a router question whose
        # trained schema id is in ids is answered by that adapter; a missing file = route off (the shared router adapter answers)
        self.router_task = {}
        for t, ta in (self.cfg["router"].get("task_adapters") or {}).items():
            files = ta.get("files") or [ta["file"]]  # several files = logit average of the adapters (an ensemble entry)
            if all((path / f).exists() for f in files):
                engs = [VS.SchemaEngine(path, VS.build_adapter(self.decision, path / f, dev), dev, tokenizer=self.span.tok, amp_bf16=adapter_bf16) for f in files]
                for i in ta["ids"]:
                    self.router_task[i] = (t, engs)
        self._VS = VS
        self.name = self.cfg["name"]

    @classmethod
    def load(cls, path, **kw):
        return cls(path, **kw)

    # ------------------------------------------------------------------ decisions (the preview, unchanged)
    def predict(self, state, questions):
        return self.decision.predict(state, questions)

    # ------------------------------------------------------------------ router questions
    def _router_answers(self, state, qs):
        textual = isinstance(state, str) or (isinstance(state, dict) and state and all(isinstance(v, str) for v in state.values()) and set(state) <= set(ROLES))
        parts = _parts(state) if textual else [("user", json.dumps(state, ensure_ascii=False, separators=(",", ":")))]  # JSON state: compact JSON
        roles = [r for r, _ in parts]; rows = {}
        for k, q in qs.items():
            s = self.router_q[_norm(q.get("instructions"))]
            crit = q.get("criteria") or {}
            if isinstance(crit, (list, tuple)):
                crit = {str(c): None for c in crit}
            opts = {str(n): (str(d) if d else (s["options"] or {}).get(str(n), "")) for n, d in crit.items()} or dict(s["options"])
            over = q.get("over") or (s["over"] if s["over"] in roles else roles[-1])
            rows.setdefault((over, s["id"] if s["id"] in self.router_task else None), []).append(dict(id=k, type=s["type"], over=over, text=q["instructions"], options=opts, abstain=bool(s["abstain"])))
        out, n_tok = {}, 0
        for (over, task), gq in rows.items():
            engs = self.router_task[task][1] if task else [self.router]; head = f"router:{self.router_task[task][0]}" if task else "router"
            raws = [e.predict_raw([dict(parts=[dict(role=r, text=t) for r, t in parts], questions=gq)])[0] for e in engs]; raw = raws[0]
            if len(raws) > 1:  # logit average over the ensemble's adapters (same questions, same option order)
                for qi, q in enumerate(raw["qs"]):
                    q["logits"] = list(np.mean([np.asarray(r_["qs"][qi]["logits"], np.float64) for r_ in raws], axis=0))
            n_tok += raw["n_input"]
            for q in raw["qs"]:
                lg = np.asarray(q["logits"], np.float64); n = len(q["names"])
                if q["qtype"] == "set":
                    p = 1 / (1 + np.exp(-lg[:n])); thr = float(self.cfg["router"]["set_threshold"])
                    pr = {nm: float(v) for nm, v in zip(q["names"], p)}
                    out[q["qid"]] = {"type": "set", "selected": [nm for nm, v in pr.items() if v > thr], "probabilities": pr, "threshold": thr, "head": head}
                else:
                    z = lg[:n]; p = np.exp(z - z.max()); p /= p.sum(); pr = {nm: float(v) for nm, v in zip(q["names"], p)}
                    out[q["qid"]] = {"type": "choice", "choice": max(pr, key=pr.get), "probabilities": pr, "head": head}
        return out, n_tok

    # ------------------------------------------------------------------ span questions
    def _label_set(self, names):
        sp = self.cfg["span"]
        for k, ls in sp["router_label_sets"].items():
            if names and set(names) <= set(ls):
                return k
        if names and {n.strip().lower() for n in names} <= {x.lower() for x in sp["evidence"]["labels"]}:
            return "evidence"
        return None

    def span_answer(self, state, q):
        VS = self._VS; sp = self.cfg["span"]
        parts = _parts(state); roles = [r for r, _ in parts]
        target = q.get("target") or q.get("over") or ("answer" if "answer" in roles else roles[-1])
        if target not in roles:
            raise ValueError(f"target {target!r} is not a part of the state")
        crit = q.get("criteria") or q.get("labels")
        if isinstance(crit, dict):
            labels = [(str(k), "" if v is None else str(v)) for k, v in crit.items()]
        elif isinstance(crit, (list, tuple)):
            labels = [(str(k), "") for k in crit]
        else:
            raise ValueError("span questions need criteria: {label: description or null} or a list of labels")
        names = [n for n, _ in labels]
        if not names or len(set(names)) != len(names):
            raise ValueError("span labels must be distinct and at least one")
        lset = self._label_set(names); question = q.get("instructions")
        if lset in sp["schemas"]:  # a trained router label set: its question text and label descriptions fill what the request leaves empty
            sch = sp["schemas"][lset]; labels = [(n, d or sch["labels"].get(n, "")) for n, d in labels]; question = question or sch["question"]
        if lset == "evidence":
            question = question or sp["evidence"]["question"]
        question = question or sp["default_question"]
        row = dict(parts=[dict(role=r, text=t) for r, t in parts], questions=[dict(id="q", type="span", over=target, text=question, labels=dict(labels))])
        evq = lset == "evidence" or (lset not in ("pii", "halu", "toxic") and bool(_EV_Q.search(_norm(question).lower())))
        eng, head = (self.span_ev if evq else self.span), "span"  # EV-FIX: evidence questions in one long window
        nq = _norm(question).lower(); names_only = not any(d for _, d in labels)
        if lset is None and not evq:  # task routes: relation questions -> the relation adapter; names-only entity questions -> ner (when shipped)
            if _REL_Q.search(nq):
                if "relation" in self.task:
                    eng, head = self.task["relation"], "relation"
            elif names_only and "ner" in self.task and self.entity_route == "ner":
                eng, head = self.task["ner"], "ner"
        ss = sp.get("set_adapter")
        if ss and self.span_set is not None and lset in ss["label_sets"]:  # R2c: router label sets on the set adapter, with its thresholds
            eng, head = self.span_set, "span_set"
            sp = dict(sp, router_set_temperature=ss["router_set_temperature"], thresholds_by_set=dict(sp["thresholds_by_set"], **ss["thresholds_by_set"]), pii_buckets=ss["pii_buckets"])
        raw = eng.predict_raw([row])[0]
        n_tok = raw["n_tokens"][target]
        self.span.last = eng.last = dict(target_tokens=int(n_tok), windows=max(1, raw["windows"]), prompt_tokens=[int(raw["n_input"])])
        # threshold and temperature: the dev picks of this checkpoint (vela3_config.json span.threshold_sources)
        if lset == "pii":
            b = next(b for b in sp["pii_buckets"] if n_tok <= b["max_tokens"]); thr0, T = b["threshold"], sp["router_set_temperature"]
            why = f"pii label set, {b['name']} bucket ({n_tok} target tokens)"
        elif lset in ("halu", "toxic"):
            thr0, T, why = sp["thresholds_by_set"][lset], sp["router_set_temperature"], f"{lset} label set"
        elif evq:
            bs = sp["evidence"].get("buckets")
            if bs:  # EV-FIX: threshold by target length
                b = next(b for b in bs if n_tok <= b["max_tokens"]); thr0, T, why = b["threshold"], 1.0, f"evidence question, {b['name']} bucket ({n_tok} target tokens)"
            else:
                thr0, T, why = sp["evidence"]["threshold"], 1.0, "evidence question"
        elif head == "relation":  # per-form dev picks: one relation (pair) / several (multi)
            thr0 = sp["relation"]["threshold_pair" if len(labels) == 1 else "threshold_multi"]; T, why = 1.0, f"relation question ({'pair' if len(labels) == 1 else 'multi'})"
        elif _REL_Q.search(nq):
            thr0, T, why = sp["relation_threshold"], 1.0, "relation question"
        elif head == "ner":
            thr0, T, why = sp["ner"]["threshold"], 1.0, "open labels, names only"
        else:
            thr0, T, why = sp["threshold"], 1.0, "open labels"
        thr = float(q.get("threshold", thr0)); text = dict(parts)[target]
        S = raw["span"]; spans = []
        if S is not None and len(S["offs"]):
            p = VS._sigmoid(S["logits"] / T); wo = S["offs"]
            spans = VS.decode_spans(p, wo, S["names"], text, thr)
            for s in spans:
                k = S["names"].index(s["label"]); m = (wo[:, 0] < s["end"]) & (wo[:, 1] > s["start"])
                s["text"] = text[s["start"]:s["end"]]; s["score"] = float(p[m, k].mean()) if m.any() else 0.0
        return {"type": "span", "target": target, "spans": spans, "head": head, "head_reason": f"{head} adapter; {why}", "threshold": thr}, raw["n_input"]

    # ------------------------------------------------------------------ System One
    def system_one(self, state, questions, model=None):
        """questions {qid: question} -> {"model", "answers": {qid: answer}, "usage"}. Choice / noul: the decision backbone (the
        preview's answers); choice / set whose instructions are a router text: the router adapter; span: the span adapter.
        A question that cannot be answered gets {"type", "error", "message"}; the others are still answered."""
        dec, rout, span, answers, n = {}, {}, {}, {}, 0
        for k, q in questions.items():
            t = q.get("type") if isinstance(q, dict) else None
            if t == "span":
                span[k] = q
            elif t in ("choice", "set") and _norm(q.get("instructions")) in self.router_q:
                if self.router_q[_norm(q.get("instructions"))]["type"] != t:
                    answers[k] = {"type": t, "error": "invalid_question", "message": "router question asked with another type than trained"}
                else:
                    rout[k] = q
            elif t in ("choice", "noul"):
                dec[k] = q
            else:
                answers[k] = {"type": t, "error": "invalid_question", "message": f"unsupported question type {t!r} (choice, noul, span; set for router questions)"}
        if dec:
            try:
                a, m = self.decision.predict(state, dec); answers.update(a); n += m
            except Unsupported as e:
                answers.update({k: {"type": q["type"], "error": "unsupported", "message": str(e)} for k, q in dec.items()})
        if rout:
            try:
                a, m = self._router_answers(state, rout); answers.update(a); n += m
            except ValueError as e:
                answers.update({k: {"type": q["type"], "error": "invalid_question", "message": str(e)} for k, q in rout.items()})
        for k, q in span.items():
            try:
                a, m = self.span_answer(state, q); answers[k] = a; n += m
            except ValueError as e:
                answers[k] = {"type": "span", "error": "invalid_question", "message": str(e)}
        return {"model": "vllm-sr/" + self.name, "answers": {k: answers[k] for k in questions}, "usage": {"input_tokens": int(n), "output_tokens": 0}}