| """vela3_coreml.py: vllm-sr/vela3-fast on Core ML (macOS 15+), no PyTorch weights needed. |
| |
| The encoder runs as a fixed-length fp16 Core ML function (the smallest bucket that fits); the small heads run in fp32 on the host. |
| Input rendering, span decoding and thresholds are the release's own code (upstream/, Apache-2.0, unmodified). |
| |
| from vela3_coreml import Vela3CoreML |
| m = Vela3CoreML(".") # this folder; compute_units="CPU_AND_GPU" by default |
| m.decide("My card was charged twice.", {"intent": {"type": "choice", "instructions": "What does the customer want?", |
| "criteria": {"refund": "Money back", "cancel": "Cancel the order"}}, |
| "urgent": {"type": "noul", "instructions": "The customer needs a fast response."}}) |
| m.pii("Call Dana Whitfield at 312-555-0110.") # -> [{"label", "start", "end", "text", "score"}] |
| m.unsupported(context="...", answer="...") # words of the answer the context does not support |
| |
| Scope: decisions (choice / noul) on the decision backbone, and the two label sets of the span_sets adapter (17 PII types, |
| unsupported claims). Inputs are limited to 1,024 tokens per encoder pass: decision inputs longer than that raise Unsupported; |
| span inputs longer than that are windowed at 1,024 tokens (the release windows at 8,192), so long-document spans can differ.""" |
| import json, sys, time |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| HERE = Path(__file__).resolve().parent |
| sys.path.insert(0, str(HERE / "upstream")) |
| import vela3_schema as VS |
| from vela3_encoder import DecisionEncoder, Vela3Encoder, Unsupported |
|
|
| PII = None |
|
|
|
|
| class _Tok: |
| """The slice of the HF tokenizer API the release's DecisionEncoder uses (ids identical to AutoTokenizer on this tokenizer.json).""" |
|
|
| def __init__(self, path, cfg): |
| from tokenizers import Tokenizer |
| self.t = Tokenizer.from_file(str(path / "tokenizer.json")); self.t.no_padding(); self.t.no_truncation() |
| self.bos_token_id, self.eos_token_id, self.pad_token_id = cfg["bos_token_id"], cfg["eos_token_id"], cfg["pad_token_id"] |
|
|
| def __call__(self, s, add_special_tokens=False): |
| return {"input_ids": self.t.encode(s, add_special_tokens=add_special_tokens).ids} |
|
|
|
|
| class _Encoder: |
| """One multifunction Core ML package: L<n> functions, input_ids / attention_mask int32 [1, n] -> hidden fp16 [1, n, 1024].""" |
|
|
| def __init__(self, package, lengths, compute_units): |
| import coremltools as ct |
| cu = getattr(ct.ComputeUnit, compute_units) |
| self.lengths = tuple(lengths) |
| self.fn = {L: ct.models.MLModel(str(package), function_name=f"L{L}", compute_units=cu) for L in self.lengths} |
| self.last_ms = [] |
|
|
| def __call__(self, ids): |
| n = len(ids) |
| L = next((L for L in self.lengths if n <= L), None) |
| if L is None: |
| raise Unsupported(f"input of {n} tokens is longer than the largest Core ML function (L{self.lengths[-1]})") |
| x = np.zeros((1, L), np.int32); x[0, :n] = ids; m = np.zeros((1, L), np.int32); m[0, :n] = 1 |
| t = time.perf_counter(); h = self.fn[L].predict({"input_ids": x, "attention_mask": m})["hidden"] |
| self.last_ms.append((L, 1e3 * (time.perf_counter() - t))) |
| return torch.from_numpy(h[0, :n].astype(np.float32)) |
|
|
|
|
| class _SpanRun: |
| """Drop-in for vela3_schema._GatedBackend: __call__(x, m, qi, oi, ui, ei) -> (choice logits, span logits).""" |
|
|
| def __init__(self, net, enc): |
| self.net, self.enc = net, enc |
|
|
| def __call__(self, x, m, qi, oi, ui, ei): |
| B, W = x.shape |
| H = torch.zeros(B, W, 1024) |
| for b in range(B): |
| n = int(m[b].sum()); H[b, :n] = self.enc(x[b, :n].tolist()) |
| with torch.no_grad(): |
| lg, sp = self.net.heads(H, *(torch.from_numpy(a) for a in (qi, oi, ui, ei))) |
| return lg.float().numpy(), sp.float().numpy() |
|
|
|
|
| class Vela3CoreML: |
| def __init__(self, path=HERE, compute_units="CPU_AND_GPU"): |
| global PII |
| path = Path(path) |
| cfg = json.loads((path / "config.json").read_text()); self.cfg = json.loads((path / "vela3_config.json").read_text()) |
| PII = self.cfg["span"]["router_label_sets"]["pii"] |
| from safetensors.torch import load_file |
| |
| d = DecisionEncoder.__new__(DecisionEncoder); nn.Module.__init__(d) |
| h, p = cfg["hidden_size"], cfg["decision"]["proj_dim"] |
| d.norm = nn.LayerNorm(h); d.w_o, d.w_q, d.w_p = nn.Linear(h, p), nn.Linear(h, p), nn.Linear(h, p) |
| d.log_tau = nn.Parameter(torch.zeros(())); d.cls_mlp = nn.Sequential(nn.Linear(h, 2 * h), nn.ReLU(), nn.Dropout(0.1), nn.Linear(2 * h, 1)) |
| d.load_state_dict(load_file(str(path / "heads_decision.safetensors")), strict=True); d.eval() |
| d.tok = _Tok(path, cfg); d.mk = {t: int(i) for t, i in cfg["decision"]["marker_ids"].items()} |
| d.device, d.max_len, d.budget = "cpu", 1024, 65536 |
| self.dec_enc = _Encoder(path / "Vela3Decision.mlpackage", (128, 256, 512, 1024), compute_units) |
| d.forward_recs = self._decision_forward |
| self.decision = d |
| |
| net = VS.SchemaNet(nn.Identity(), h, p, span=True) |
| net.load_state_dict(load_file(str(path / "heads_span_sets.safetensors")), strict=True); net.eval() |
| self.span_enc = _Encoder(path / "Vela3SpanSets.mlpackage", (512, 1024), compute_units) |
| eng = VS.SchemaEngine(path, net, "cpu", max_len=1024) |
| eng.run = _SpanRun(net, self.span_enc) |
| v = Vela3Encoder.__new__(Vela3Encoder) |
| v.cfg, v._VS, v.span, v.span_ev, v.span_set, v.task, v.entity_route = self.cfg, VS, eng, eng, eng, {}, "ner" |
| self.v = v |
|
|
| @torch.no_grad() |
| def _decision_forward(self, recs): |
| out = [] |
| for r in recs: |
| H = self.dec_enc(r["ids"])[None] |
| z = self.decision._logits(H, [r])[0].double().numpy(); e = np.exp(z - z.max()); out.append(e / e.sum()) |
| return out |
|
|
| def decide(self, state, questions): |
| """questions: {key: {type: choice | noul, instructions, criteria}} -> {key: answer} (the release's predict()).""" |
| return self.decision.predict(state, questions)[0] |
|
|
| def pii(self, text, threshold=None): |
| q = {"type": "span", "criteria": PII} |
| if threshold is not None: |
| q["threshold"] = threshold |
| return self.v.span_answer(text, q)[0]["spans"] |
|
|
| def unsupported(self, context, answer, threshold=None): |
| q = {"type": "span", "criteria": ["unsupported"]} |
| if threshold is not None: |
| q["threshold"] = threshold |
| return self.v.span_answer({"context": context, "answer": answer}, q)[0]["spans"] |
|
|
|
|
| if __name__ == "__main__": |
| m = Vela3CoreML() |
| print(m.decide("My card was charged twice and I need this fixed today.", |
| {"intent": {"type": "choice", "instructions": "What does the customer want?", |
| "criteria": {"refund": "Money back for the duplicate charge", "cancel": "Cancel the order"}}, |
| "urgent": {"type": "noul", "instructions": "The customer needs a fast response."}})) |
| print(m.pii("Hi, I'm Dana Whitfield (dana@example.com, 312-555-0110), 4410 Elm Ridge Drive, Dayton, OH 45429.")) |
| print(m.unsupported(context="Adults may take up to 4 g of paracetamol in 24 hours.", |
| answer="Adults may take up to 6 g of paracetamol in 24 hours.")) |
|
|
