| 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) |
| 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" |
|
|
