| """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() |
| |
| 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) |
|
|
|
|
| |
| |
| |
| 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)) |
| evf = self.cfg["span"]["evidence"].get("adapter") |
| 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 = {} |
| 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) |
| |
| |
| |
| import os as _os |
| self.entity_route = entity_route or _os.environ.get("VELA3_ENTITY_ROUTE") or self.cfg["span"].get("entity_route", "ner") |
| |
| |
| 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) |
| |
| |
| self.router_task = {} |
| for t, ta in (self.cfg["router"].get("task_adapters") or {}).items(): |
| files = ta.get("files") or [ta["file"]] |
| 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) |
|
|
| |
| def predict(self, state, questions): |
| return self.decision.predict(state, 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=(",", ":")))] |
| 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: |
| 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 |
|
|
| |
| 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"]: |
| 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" |
| nq = _norm(question).lower(); names_only = not any(d for _, d in labels) |
| if lset is None and not evq: |
| 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"]: |
| 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"])]) |
| |
| 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: |
| 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": |
| 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"] |
|
|
| |
| 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}} |
|
|
