File size: 3,829 Bytes
ba1c728 | 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 | import json, os, pickle, re, numpy as np, onnxruntime as ort
from tokenizers import Tokenizer
_TP = re.compile(r"(?u)\b\w\w+\b")
class Strat:
"""strat (fast kNN): tier-1 tf-idf, escalating to tier-2 int8 student + class-restricted kNN memory. CPU, 1 thread."""
def __init__(s, path, threads=1):
c = json.load(open(os.path.join(path, "config.json"))); s.c = c; s.labels = c["labels"]; C = len(s.labels)
s.wv, s.wi, s.cv, s.ci, s.cT, s.b = pickle.load(open(os.path.join(path, "fast.pkl"), "rb")); s.off = len(s.wv)
so = ort.SessionOptions(); so.intra_op_num_threads = threads; so.inter_op_num_threads = 1
s.sess = ort.InferenceSession(os.path.join(path, "student_int8.onnx"), so, providers=["CPUExecutionProvider"])
s.bt = Tokenizer.from_file(os.path.join(path, "tokenizer.json")); s.bt.enable_truncation(c["maxlen"]); s.bt.no_padding()
k = np.load(os.path.join(path, "knn.npz")); E, y = k["E"], k["y"]
o = np.argsort(y, kind="stable"); s.Ei = np.ascontiguousarray(E[o]); s.yi = y[o].astype(np.int64) # sorted by class
s.coff = np.concatenate([[0], np.cumsum(np.bincount(s.yi, minlength=C))]).astype(np.int64)
s.M = int(c.get("knn_classes", 20)); s.nl = C
def _tf(s, terms, vocab, idf, off):
cnt = {}
for t in terms:
j = vocab.get(t)
if j is not None: cnt[j] = cnt.get(j, 0) + 1
if not cnt: return None
idx = np.fromiter(cnt.keys(), dtype=np.int64, count=len(cnt)); tf = np.fromiter(cnt.values(), dtype=np.float32, count=len(cnt))
v = (1.0 + np.log(tf)) * idf[idx]; n = np.sqrt((v * v).sum())
return idx + off, (v / n if n > 0 else v)
def feat(s, text):
t = text.lower(); w = _TP.findall(t)
a = s._tf(w + [w[i] + " " + w[i + 1] for i in range(len(w) - 1)], s.wv, s.wi, 0)
cg = []
for wd in t.split():
wd = " " + wd + " "; L = len(wd)
for n in range(2, 6):
o = 0; cg.append(wd[o:o + n])
while o + n < L:
o += 1; cg.append(wd[o:o + n])
if o == 0: break
b = s._tf(cg, s.cv, s.ci, s.off)
ps = [x for x in (a, b) if x is not None]
if not ps: return np.zeros(0, np.int64), np.zeros(0, np.float32)
return np.concatenate([p[0] for p in ps]), np.concatenate([p[1] for p in ps]).astype(np.float32)
def fast_p(s, text):
idx, val = s.feat(text)
z = val @ s.cT[idx] + s.b
z = z - z.max(); p = np.exp(z); return p / p.sum()
def _knn(s, e, p, p_tf):
if s.M <= 0 or s.M >= s.nl:
sim = s.Ei @ e; lab = s.yi
else:
cand = np.argpartition(-p, s.M)[:s.M]
if p_tf is not None: cand = np.union1d(cand, np.argpartition(-p_tf, 4)[:4])
co = s.coff; sims = [s.Ei[co[c]:co[c + 1]] @ e for c in cand]
sim = np.concatenate(sims); lab = np.concatenate([s.yi[co[c]:co[c + 1]] for c in cand])
top = np.argpartition(-sim, 10)[:10]; w = np.exp((sim[top] - sim[top].max()) / 0.05)
pk = np.bincount(lab[top], weights=w, minlength=s.nl); return pk / pk.sum()
def student_p(s, text, p_tf=None):
ids = np.array([s.bt.encode(text).ids], dtype=np.int64)
lg, e = s.sess.run(None, {"input_ids": ids, "attention_mask": np.ones_like(ids)})
z = lg[0] - lg[0].max(); p = np.exp(z); p /= p.sum()
lam = s.c["lam"]
if lam > 0: p = (1 - lam) * p + lam * s._knn(e[0], p, p_tf)
mu = s.c.get("mu", 0.0)
if mu > 0 and p_tf is not None: p = (1 - mu) * p + mu * p_tf
return p
def predict_one(s, text):
p = s.fast_p(text)
if p.max() >= s.c["tau"]: return s.labels[int(p.argmax())], "fast"
return s.labels[int(s.student_p(text, p).argmax())], "student"
|