File size: 4,121 Bytes
c735e00 | 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 | """Run a chunked D1A Core ML trunk end to end: host embeddings -> chunk 0 .. chunk N (h, per-layer inputs and the shared
keys/values passed along) -> hidden states at the picks -> pointer head on the host.
Usage as a library: Pipeline(export_dir, chunk_dir, L, units).probs(enc) -> [per-question probabilities]."""
import glob
import os
import re
import sys
import time
import coremltools as ct
import numpy as np
import torch
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__))) # d1a_coreml.py sits next to this file
from d1a_coreml import Export, HostEmbed, masks_for # noqa: E402
P_PICKS = 128
class Pipeline:
def __init__(self, export_dir, chunk_dir, L, units="CPU_AND_NE"):
self.ex = Export(export_dir)
self.host = HostEmbed(self.ex)
self.hq, self.hk = self.ex.head()
self.T = self.ex.d1a["temperature"]
self.L = L
cu = getattr(ct.ComputeUnit, units)
self.chunks = []
if chunk_dir.endswith(".mlpackage"): # one multifunction package: functions c00_07, c07_15, ...
names = ["c00_07", "c07_15", "c15_22", "c22_29", "c29_35"]
for name in names:
a, b = map(int, name[1:].split("_"))
m = ct.models.MLModel(chunk_dir, compute_units=cu, function_name=name)
fn = next(f for f in m.get_spec().description.functions if f.name == name)
self.chunks.append((a, b, m, [i.name for i in fn.input], [o.name for o in fn.output]))
return
files = sorted(glob.glob(f"{chunk_dir}/L{L}_*.mlpackage"))
for f in files:
a, b = map(int, re.search(r"_(\d+)-(\d+)\.mlpackage$", f).groups())
m = ct.models.MLModel(f, compute_units=cu)
spec = m.get_spec()
self.chunks.append((a, b, m, [i.name for i in spec.description.input], [o.name for o in spec.description.output]))
def hidden(self, enc):
L, n = self.L, len(enc["ids"])
# right-pad with the pad token's own embeddings (as d1a's torch backend pads ids), never zero rows: an all-zero row
# makes every RMSNorm compute 0 / sqrt(eps), and eps (1e-6) is subnormal in fp16, which the Neural Engine flushes
ids = list(enc["ids"]) + [self.ex.d1a.get("pad_id", 0)] * (L - n)
emb, ple = self.host(ids)
e = emb.numpy().astype(np.float16)
p = ple.numpy().astype(np.float16)
mf, ms, pos = masks_for(enc, self.ex.tc.sliding_window, L=L)
picks = list(enc["decide_idx"]) + [o for oo in enc["opt_idx"] for o in oo]
pk = np.zeros(P_PICKS, np.int32); pk[: len(picks)] = picks
common = {"pos": pos.numpy().astype(np.int32), "mask_full": mf.clamp(min=-1e4).numpy().astype(np.float16),
"mask_slide": ms.clamp(min=-1e4).numpy().astype(np.float16)}
state = {"emb": e, "ple_tok": p, "picks": pk}
for a, b, m, ins, outs in self.chunks:
feed = {k: np.ascontiguousarray(common[k] if k in common else state[k]) for k in ins}
feed = {k: (v.astype(np.float16) if v.dtype == np.float32 else v) for k, v in feed.items()} # chunk outputs come back fp32
out = {k: np.array(v, copy=True) for k, v in m.predict(feed).items()} # Core ML may reuse its output buffers on the next call
if "hidden" in out:
return torch.tensor(out["hidden"][: len(picks)], dtype=torch.float32)
state["h_in"] = out["h"]
if "ple_out" in out:
state["ple"] = out["ple_out"]
for k, v in out.items():
if k.endswith("_out") and k != "ple_out":
state[k[: -len("_out")]] = v
raise RuntimeError("no chunk returned hidden")
def probs(self, enc):
h = self.hidden(enc)
Q = len(enc["decide_idx"])
out, j = [], Q
with torch.no_grad():
for qi, oo in enumerate(enc["opt_idx"]):
z = (self.hk(h[j: j + len(oo)]) @ self.hq(h[qi])) / (self.hq.out_features ** 0.5)
out.append(torch.softmax(z / self.T, -1).numpy()); j += len(oo)
return out
|