File size: 7,789 Bytes
deaf348 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 | """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 # noqa: E402
from vela3_encoder import DecisionEncoder, Vela3Encoder, Unsupported # noqa: E402
PII = None # filled from vela3_config.json
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
# decisions: the release DecisionEncoder without its torch encoder (heads + rendering only), encoder on Core ML
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
# spans: the release SchemaNet heads of the span_sets adapter + its SchemaEngine, encoder on Core ML
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."))
|