File size: 2,143 Bytes
659433c | 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 | """Numpy-only neural inference; no semantic phrase rules. Fixed feature hashing."""
from importlib.resources import files
from core import Diagnostic, assemble, compile_js, frontend
# isort: split
# core must set CPU thread limits before NumPy initializes.
import numpy as np
def default_checkpoint():
return files("tiny_codemode_weights").joinpath("model.npz")
class Model:
def __init__(self, path=None):
with np.load(path if path is not None else default_checkpoint()) as z:
self.w = {k: z[k] for k in z.files}
def logits(self, item):
_, ids, f, _ = frontend(item)
return self.forward(ids, f)
def forward(self, ids, f):
w = self.w
dims = w["convs.0.weight"].shape[1] - w["embed.weight"].shape[1]
h = np.concatenate(
[w["embed.weight"][ids % len(w["embed.weight"])], f[:, :dims]], axis=1
)
for i in range(3):
weight = w[f"convs.{i}.weight"]
bias = w[f"convs.{i}.bias"]
pad = weight.shape[2] // 2
p = np.pad(h, ((pad, pad), (0, 0)))
h = np.maximum(
sum(
p[k : k + len(h)] @ weight[:, :, k].T
for k in range(weight.shape[2])
)
+ bias,
0,
)
pooled = np.concatenate([h.mean(0), h.max(0)])
heads = (
h @ w["roles.weight"].T + w["roles.bias"],
pooled @ w["mode.weight"].T + w["mode.bias"],
pooled @ w["output.weight"].T + w["output.bias"],
)
return (
heads + (pooled @ w["reject.weight"].T + w["reject.bias"],)
if "reject.weight" in w
else heads
)
def predict(self, item):
heads = self.logits(item)
if len(heads) == 4 and int(heads[3].argmax()) == 1:
raise Diagnostic("learned rejection: unsupported/invalid request")
r, m, o = heads[:3]
return assemble(item, r.argmax(1), int(m.argmax()), int(o.argmax()))
def program(self, item):
return compile_js(self.predict(item), item["tools"])
|