| """AksharaLP v3 tiny sentence LM: Devanagari subword (unigram-LM pieces, Viterbi segmentation reimplemented here so Python and |
| JS tokenize identically) + 2-layer LSTM with projection (LSTMP) + tied input/output embedding, trained with a self-normalisation |
| penalty (alpha * logZ^2) so the runtime may skip the softmax normaliser (norm='self') or compute it exactly (norm='exact'). |
| File format (same low-bit scheme as bhc/xlit_q.py): u32 header length | JSON header {cfg, bits, G, tensors, vocab:[[piece, score]...]} | data. |
| Sentence = '</s>' w1 w2 ... wn '</s>' ; digits runs are the word '<n>'.""" |
| import io, json, math |
| import numpy as np |
|
|
| SP = "\u2581" |
|
|
|
|
| class Tok: |
| """Viterbi over piece log-probs (identical to sentencepiece unigram encode for single whitespace-free words, no sampling).""" |
| def __init__(self, vocab): |
| self.pieces = [p for p, _ in vocab]; self.score = {p: s for p, s in vocab}; self.id = {p: i for i, p in enumerate(self.pieces)} |
| self.maxlen = max(len(p) for p in self.pieces); self.unk = self.id["<unk>"]; self.eos = self.id["</s>"] |
| self.unk_score = min(s for _, s in vocab) - 10.0; self.cache = {} |
|
|
| def encode_word(self, w): |
| r = self.cache.get(w) |
| if r is not None: return r |
| if w in self.id and w.startswith("<"): r = [self.id[w]]; self.cache[w] = r; return r |
| s = SP + w; n = len(s); best = [-1e30] * (n + 1); back = [None] * (n + 1); best[0] = 0.0 |
| for i in range(n): |
| if best[i] <= -1e29: continue |
| hit = False |
| for j in range(i + 1, min(n, i + self.maxlen) + 1): |
| p = s[i:j]; sc = self.score.get(p) |
| if sc is not None and p not in ("<unk>", "<s>", "</s>", "<n>"): |
| hit = True |
| if best[i] + sc > best[j]: best[j] = best[i] + sc; back[j] = (i, self.id[p]) |
| if not hit or (back[i + 1] is None): |
| if best[i] + self.unk_score > best[i + 1]: best[i + 1] = best[i] + self.unk_score; back[i + 1] = (i, self.unk) |
| out = []; j = n |
| while j > 0: i, pid = back[j]; out.append(pid); j = i |
| r = out[::-1] |
| if len(self.cache) < 200000: self.cache[w] = r |
| return r |
|
|
|
|
| def _dq(data, off, shape, kind, G): |
| r, c = shape; ng = (c + G - 1) // G |
| if kind == "q8": |
| q = np.frombuffer(data, np.int8, r * c, off).astype(np.float32).reshape(r, c); off += r * c |
| else: |
| nb = (r * c + 1) // 2; u = np.frombuffer(data, np.uint8, nb, off); off += nb |
| q = np.empty(nb * 2, np.int8); q[0::2] = (u >> 4).astype(np.int8) - 8; q[1::2] = (u & 15).astype(np.int8) - 8 |
| q = q[:r * c].astype(np.float32).reshape(r, c) |
| off += (-off) % 4 |
| sc = np.frombuffer(data, np.float16, r * ng, off).astype(np.float32).reshape(r, ng) |
| return q * np.repeat(sc, G, 1)[:, :c] |
|
|
|
|
| class NLM: |
| """numpy float32 inference. state = (h[L,E], c[L,H]) per hypothesis, stacked into arrays for batching.""" |
| def __init__(self, cfg, W, vocab, norm="exact"): |
| self.cfg = cfg; self.W = W; self.tok = Tok(vocab); self.L = cfg["layers"]; self.norm = norm |
| self.emb = W["emb.weight"]; self.ob = W["out_bias"] |
| self._init = None; L = self.L |
| self._wihT = [np.ascontiguousarray(W[f"lstm.weight_ih_l{l}"].T) for l in range(L)]; self._wih0T = self._wihT[0] |
| self._whhT = [np.ascontiguousarray(W[f"lstm.weight_hh_l{l}"].T) for l in range(L)] |
| self._whrT = [np.ascontiguousarray(W[f"lstm.weight_hr_l{l}"].T) for l in range(L)] |
| self._b = [W[f"lstm.bias_ih_l{l}"] + W[f"lstm.bias_hh_l{l}"] for l in range(L)]; self._embT = np.ascontiguousarray(self.emb.T) |
|
|
| @staticmethod |
| def load(path, norm=None): |
| b = open(path, "rb").read(); n = int.from_bytes(b[:4], "little"); meta = json.loads(b[4:4 + n]); data = b[4 + n:] |
| W = {} |
| for k, shape, kind, off, ln in meta["tensors"]: |
| if kind == "f16": |
| cnt = int(np.prod(shape)); W[k] = np.frombuffer(data, np.float16, cnt, off).astype(np.float32).reshape(shape) |
| elif kind == "f32": |
| cnt = int(np.prod(shape)); W[k] = np.frombuffer(data, np.float32, cnt, off).reshape(shape).copy() |
| else: W[k] = _dq(data, off, shape, kind, meta["G"]) |
| return NLM(meta["cfg"], W, meta["vocab"], norm or meta["cfg"].get("norm", "exact")) |
|
|
| def _step(self, h, c, ids, g0=None): |
| """h: [L,B,E], c: [L,B,H], ids: [B] -> new h, c. g0: optional precomputed layer-0 gate pre-activations [B,4H].""" |
| hs = []; cs = [] |
| if g0 is None: |
| u, inv = np.unique(ids, return_inverse=True) |
| g0 = (self.emb[u] @ self._wih0T + self._b[0])[inv] + h[0] @ self._whhT[0] |
| for l in range(self.L): |
| g = g0 if l == 0 else x @ self._wihT[l] + h[l] @ self._whhT[l] + self._b[l] |
| H = g.shape[1] // 4 |
| i_ = 1 / (1 + np.exp(-g[:, :H])); f_ = 1 / (1 + np.exp(-g[:, H:2 * H])); g_ = np.tanh(g[:, 2 * H:3 * H]); o_ = 1 / (1 + np.exp(-g[:, 3 * H:])) |
| cn = f_ * c[l] + i_ * g_; hn = (o_ * np.tanh(cn)) @ self._whrT[l] |
| hs.append(hn); cs.append(cn); x = hn |
| return np.stack(hs), np.stack(cs) |
|
|
| def _logp(self, h, ids): |
| """log P(ids[b] | state b) using the top-layer output h[-1]: [B]""" |
| top = h[-1]; z = (top * self.emb[ids]).sum(1) + self.ob[ids] |
| if self.norm == "exact": |
| lg = top @ self._embT + self.ob; m = lg.max(1); z = z - (m + np.log(np.exp(lg - m[:, None]).sum(1))) |
| return z |
|
|
| def init_state(self): |
| if self._init is None: |
| E = self.emb.shape[1]; H = self.W["lstm.weight_hh_l0"].shape[1] * 0 + self.W["lstm.weight_ih_l0"].shape[0] // 4 |
| h = np.zeros((self.L, 1, E), np.float32); c = np.zeros((self.L, 1, H), np.float32) |
| h, c = self._step(h, c, np.array([self.tok.eos])); self._init = (h[:, 0], c[:, 0]) |
| return self._init |
|
|
| def score_words(self, states, words): |
| """states: list of B states; words: list of K Devanagari words -> (lp[B][K], new_states[B][K]). |
| Exact chain rule over each word's pieces from each state, batched over all (b,k) pairs. First piece: one softmax |
| normaliser per beam state and the layer-0 recurrent term once per beam (broadcast over the K candidates).""" |
| B, K = len(states), len(words); seqs = [self.tok.encode_word(w) for w in words] |
| h = np.stack([s[0] for s in states], 1); c = np.stack([s[1] for s in states], 1) |
| id0 = np.array([q[0] for q in seqs]); top = h[-1] |
| lp = top @ self.emb[id0].T + self.ob[id0][None, :] |
| if self.norm == "exact": |
| lg = top @ self._embT + self.ob; m = lg.max(1); lp = lp - (m + np.log(np.exp(lg - m[:, None]).sum(1)))[:, None] |
| g0 = ((self.emb[id0] @ self._wih0T + self._b[0])[None, :, :] + (h[0] @ self._whhT[0])[:, None, :]).reshape(B * K, -1) |
| hb, cb = self._step(np.repeat(h, K, 1), np.repeat(c, K, 1), None, g0=g0) |
| flat = lp.reshape(-1).astype(np.float32); maxp = max(len(q) for q in seqs) |
| for t in range(1, maxp): |
| ks = [k for k in range(K) if len(seqs[k]) > t] |
| idx = (np.arange(B)[:, None] * K + np.array(ks)[None, :]).reshape(-1) |
| ids = np.tile(np.array([seqs[k][t] for k in ks]), B) |
| flat[idx] += self._logp(hb[:, idx], ids) |
| hn, cn = self._step(hb[:, idx], cb[:, idx], ids); hb[:, idx] = hn; cb[:, idx] = cn |
| ns = [[(hb[:, b * K + k], cb[:, b * K + k]) for k in range(K)] for b in range(B)] |
| return flat.reshape(B, K).tolist(), ns |
|
|
| def advance_many(self, states, words): |
| lp, ns = self.score_words(states, words); return ns |
|
|
| def eos_logp(self, state): |
| h = state[0][:, None, :] |
| return float(self._logp(h, np.array([self.tok.eos]))[0]) |
|
|
| def sent_logp(self, words): |
| st = self.init_state(); tot = 0.0 |
| for w in words: |
| lp, ns = self.score_words([st], [w]); tot += lp[0][0]; st = ns[0][0] |
| return tot + self.eos_logp(st) |
|
|
|
|
| |
| def torch_model(cfg): |
| import torch, torch.nn as nn |
|
|
| class LSTMP_LM(nn.Module): |
| def __init__(s, cfg): |
| super().__init__(); s.cfg = cfg; V, E, H, L = cfg["vocab"], cfg["emb"], cfg["hid"], cfg["layers"] |
| s.emb = nn.Embedding(V, E); nn.init.normal_(s.emb.weight, 0, E ** -0.5) |
| s.lstm = nn.LSTM(E, H, num_layers=L, proj_size=E, batch_first=True, dropout=cfg.get("dropout", 0.0)) |
| s.out_bias = nn.Parameter(torch.zeros(V)); s.drop = nn.Dropout(cfg.get("dropout", 0.0)) |
|
|
| def forward(s, x, state=None): |
| h, st = s.lstm(s.drop(s.emb(x)), state) |
| return s.drop(h) @ s.emb.weight.T + s.out_bias, st |
| return LSTMP_LM(cfg) |
|
|
|
|
| def _q(v, bits, G): |
| """symmetric b-bit ints, one fp16 scale per G consecutive inputs per row, MSE-optimal clipping for 4-bit (same as bhc/xlit_q.py).""" |
| r, c = v.shape; pad = (-c) % G; vp = np.pad(v, ((0, 0), (0, pad))).reshape(r, -1, G) |
| L = 2 ** (bits - 1) - 1; base = np.abs(vp).max(2, keepdims=True) / L; base[base == 0] = 1 |
| best_err = None |
| for al in ((1.0,) if bits >= 8 else (1.0, 0.95, 0.9, 0.85, 0.8, 0.75, 0.7, 0.65)): |
| s_ = (base * al).astype(np.float16).astype(np.float32); q_ = np.clip(np.round(vp / s_), -L, L) |
| e = ((q_ * s_ - vp) ** 2).sum(2, keepdims=True) |
| if best_err is None: best_err, sc, q = e, s_, q_ |
| else: m = e < best_err; best_err = np.where(m, e, best_err); sc = np.where(m, s_, sc); q = np.where(m, q_, q) |
| return q.astype(np.int8).reshape(r, -1)[:, :c], sc[..., 0].astype(np.float16) |
|
|
|
|
| def export(model, vocab, path, bits=4, G=32, emb_bits=8, norm="self"): |
| """bits/emb_bits in {4, 8} or 0 (= float32, exact reference).""" |
| sd = model.state_dict(); cfg = dict(model.cfg); cfg["norm"] = norm |
| meta = {"cfg": cfg, "bits": bits, "G": G, "tensors": [], "vocab": vocab}; buf = io.BytesIO() |
| for k, v in sd.items(): |
| v = v.float().cpu().numpy(); buf.write(b"\0" * ((-buf.tell()) % 4)) |
| b_ = emb_bits if k.startswith("emb") else bits |
| if v.ndim == 2 and b_: |
| q, sc = _q(v, b_, G); flat = q.reshape(-1) |
| if b_ == 4: |
| u = (flat + 8).astype(np.uint8); u = np.r_[u, np.zeros(len(u) % 2, np.uint8)]; data = (u[0::2] << 4 | u[1::2]).astype(np.uint8).tobytes() |
| else: data = flat.astype(np.int8).tobytes() |
| meta["tensors"].append([k, list(v.shape), "q%d" % b_, buf.tell(), len(data)]); buf.write(data); buf.write(b"\0" * ((-buf.tell()) % 4)); buf.write(sc.tobytes()) |
| else: |
| meta["tensors"].append([k, list(v.shape), "f32", buf.tell(), 0]); buf.write(v.astype(np.float32).tobytes()) |
| hdr = json.dumps(meta, separators=(",", ":"), ensure_ascii=False).encode(); hdr += b" " * ((-len(hdr) - 4) % 4) |
| with open(path, "wb") as f: f.write(len(hdr).to_bytes(4, "little")); f.write(hdr); f.write(buf.getvalue()) |
|
|
