| """Compact Devanagari lexicon + pruned backoff word-bigram LM (BhashaCore Track B/C assets), format v2. |
| |
| Word ids: words sorted by (unigram bucket = round(-2*ln count) i.e. frequency class, then alphabetical) so |
| ids are roughly frequency-ordered (small ids for frequent words -> small LM deltas), while the list is still |
| front-coded within each bucket. Unigram ln-prob is per bucket (implicit; no per-word bytes). |
| lex.bin : header | per-bucket (n_words varint, ln-prob f16) | front-coded list (1 byte/char = codepoint-0x0900, |
| last suffix byte has bit 7 set; 0x7F = empty suffix) |
| lm.bin : header | successor count per history (varint; history N = sentence start) | delta-coded successor |
| ids (varint) | 4-bit codebook ln p(w|v), 16 centroids f32 | 4-bit codebook ln backoff for histories |
| with successors (16 centroids) |
| """ |
| import gzip, math, struct |
| import numpy as np |
|
|
| HK = 1 << 21 |
|
|
|
|
| def varints(vals): |
| out = bytearray() |
| for v in vals: |
| v = int(v) |
| while v >= 0x80: out.append((v & 0x7F) | 0x80); v >>= 7 |
| out.append(v) |
| return bytes(out) |
|
|
|
|
| def read_varints(b, n, pos=0): |
| out = [0] * n; i = 0 |
| while i < n: |
| v = sh = 0 |
| while True: |
| c = b[pos]; pos += 1; v |= (c & 0x7F) << sh; sh += 7 |
| if c < 0x80: break |
| out[i] = v; i += 1 |
| return out, pos |
|
|
|
|
| def kmeans1d(x, k=16, it=30): |
| x = np.asarray(x, np.float64); c = np.quantile(x, (np.arange(k) + 0.5) / k) |
| for _ in range(it): |
| a = np.abs(x[:, None] - c[None]).argmin(1) |
| for j in range(k): |
| if (a == j).any(): c[j] = x[a == j].mean() |
| c = np.sort(c); a = np.abs(x[:, None] - c[None]).argmin(1) if len(x) < 2_000_000 else None |
| return c.astype(np.float32), a.astype(np.uint8) |
|
|
|
|
| def pack4(a): |
| a = np.r_[a, np.zeros(len(a) % 2, np.uint8)].astype(np.uint8); return (a[0::2] * 16 + a[1::2]).astype(np.uint8).tobytes() |
|
|
|
|
| def unpack4(b, n): |
| a = np.frombuffer(b, np.uint8); return np.stack([a >> 4, a & 15], 1).reshape(-1)[:n] |
|
|
|
|
| def build(words_by_freq, counts_by_freq, bigram_pairs, bigram_counts, n_lex, maxbi, D=0.75, cmin=2, K=262144, alpha=False): |
| """words_by_freq[i] has count id i+2 in bigram_pairs (0=boundary, 1=unk).""" |
| N = n_lex; ws = words_by_freq[:N]; cs = counts_by_freq[:N] |
| bk = [-int(round(2 * math.log(c))) for c in cs] |
| if alpha: |
| order = sorted(range(N), key=lambda i: ws[i]) |
| else: |
| order = sorted(range(N), key=lambda i: (bk[i], ws[i])) |
| nid = np.empty(N + 2, np.int64); nid[0] = N; nid[1] = -1 |
| for a, i in enumerate(order): nid[i + 2] = a |
| tot = float(sum(cs)) * 1.02 |
| buckets = [] |
| for i in order: |
| if not buckets or buckets[-1][0] != bk[i]: buckets.append([bk[i], 0, 0.0]) |
| buckets[-1][1] += 1; buckets[-1][2] += cs[i] |
| b_lp = [math.log(s / n / tot) for _, n, s in buckets] |
| unk_lp = math.log(0.02 / 50000) |
| puni = np.empty(N + 2); puni[0] = math.log(0.04); puni[1] = unk_lp; j = 0 |
| if alpha: |
| lc = np.log(np.array(cs, np.float64) / tot); NB = 16 |
| ucb = np.linspace(lc.min(), lc.max(), NB).astype(np.float32) |
| ua = np.abs(lc[:, None] - ucb[None].astype(np.float64)).argmin(1).astype(np.uint8) |
| puni[2:] = ucb[ua] |
| else: |
| for (_, n, _), lp in zip(buckets, b_lp): |
| for _ in range(n): puni[[order[j] + 2]] = lp; j += 1 |
| v = bigram_pairs // K; w = bigram_pairs % K; c = bigram_counts.astype(np.float64) |
| inlex = (v - 2 < N) & (w - 2 < N) |
| hist_tot = np.bincount(np.where(v - 2 < N, v, 1), weights=c, minlength=N + 2) |
| keep = inlex & (v != 1) & (w != 1) & (c >= cmin) |
| pk_all = np.log(np.maximum(c - D, 0.1) / np.maximum(hist_tot[np.minimum(v, N + 1)], 1)) |
| gain = np.where(keep, c * (pk_all - puni[np.minimum(w, N + 1)]), -np.inf) |
| if keep.sum() > maxbi: |
| thr = np.partition(gain, len(gain) - maxbi)[len(gain) - maxbi]; keep &= gain >= thr |
| kv, kw, pk = v[keep], w[keep], pk_all[keep] |
| av = nid[kv]; aw = nid[kw] |
| o = np.lexsort((aw, av)); av, aw, pk = av[o], aw[o], pk[o] |
| cb, qa = kmeans1d(pk); pkd = cb[qa] |
| cnt = np.bincount(av, minlength=N + 1) |
| first = np.r_[True, av[1:] != av[:-1]]; delta = np.where(first, aw, aw - np.r_[0, aw[:-1]] - 1) |
| |
| pu_new = np.empty(N + 1); pu_new[nid[np.arange(2, N + 2)]] = puni[2:]; pu_new[N] = puni[0] |
| s_bi = np.bincount(av, weights=np.exp(pkd), minlength=N + 1); s_un = np.bincount(av, weights=np.exp(pu_new[aw]), minlength=N + 1) |
| hs = np.nonzero(cnt)[0] |
| bow = np.log(np.maximum(1 - s_bi[hs], 1e-4) / np.maximum(1 - s_un[hs], 1e-4)) |
| cbb, qb = kmeans1d(bow) |
| fc = bytearray(); prev = b"" |
| for i in order: |
| bb = bytes(ord(ch) - 0x0900 for ch in ws[i]); k = 0 |
| while k < min(len(bb), len(prev), 127) and bb[k] == prev[k]: k += 1 |
| suf = bytearray(bb[k:]) or bytearray([0x7F]); suf[-1] |= 0x80; fc.append(k); fc += suf; prev = bb |
| bh = varints([n for _, n, _ in buckets]) + np.array(b_lp, np.float16).tobytes() |
| if alpha: |
| P = bytearray(); S = bytearray(); prev = b"" |
| for i in order: |
| bb = bytes(ord(ch) - 0x0900 for ch in ws[i]); k = 0 |
| while k < min(len(bb), len(prev), 127) and bb[k] == prev[k]: k += 1 |
| suf = bytearray(bb[k:]) or bytearray([0x7F]); suf[-1] |= 0x80; P.append(k); S += suf; prev = bb |
| lex = b"LX4\0" + struct.pack("<IIfI", N, len(S), unk_lp, NB) + ucb.astype(np.float32).tobytes() + pack4(ua[order]) + bytes(P) + bytes(S) |
| else: |
| lex = struct.pack("<IIIf", N, len(buckets), len(fc), unk_lp) + bh + bytes(fc) |
| b1, b2 = varints(cnt), varints(delta) |
| lm = (struct.pack("<IIIII", N, len(av), len(b1), len(b2), len(hs)) + cb.tobytes() + cbb.tobytes() + b1 + b2 |
| + pack4(qa) + pack4(qb)) |
| return lex, lm, {"lex_words": N, "bigrams": int(len(av)), "histories": int(len(hs)), "buckets": len(buckets)} |
|
|
|
|
| class LexLM: |
| def __init__(self, lex_path, lm_path=None): |
| if lex_path is None: |
| self.words = []; self.N = 0; self.id = {}; self.uni = []; self.unk = 0.0; self.bi = {}; self.bow = {}; return |
| b = open(lex_path, "rb").read() |
| if b[:4] == b"LX4\0": |
| N, nfc, self.unk, NB = struct.unpack("<IIfI", b[4:20]); pos = 20; ucb = np.frombuffer(b, np.float32, NB, pos); pos += 4 * NB |
| band = unpack4(b[pos:pos + (N + 1) // 2], N); pos += (N + 1) // 2; uni = ucb[band].tolist(); bn = None |
| P = b[pos:pos + N]; S = b[pos + N:pos + N + nfc]; words = []; prev = []; j = 0 |
| for k in P: |
| cur = prev[:k] |
| while True: |
| c = S[j]; j += 1; ch = c & 0x7F |
| if ch != 0x7F: cur.append(ch) |
| if c & 0x80: break |
| words.append("".join(chr(0x0900 + x) for x in cur)); prev = cur |
| pos = end = 0 |
| else: |
| N, nb, nfc, self.unk = struct.unpack("<IIIf", b[:16]); pos = 16 |
| bn, pos = read_varints(b, nb, pos); blp = np.frombuffer(b, np.float16, nb, pos).astype(np.float32).tolist(); pos += 2 * nb |
| if bn is not None: words = []; prev = []; end = pos + nfc |
| while pos < end: |
| k = b[pos]; pos += 1; cur = prev[:k] |
| while True: |
| c = b[pos]; pos += 1; ch = c & 0x7F |
| if ch != 0x7F: cur.append(ch) |
| if c & 0x80: break |
| words.append("".join(chr(0x0900 + x) for x in cur)); prev = cur |
| self.words = words; self.N = N; self.id = {w: i for i, w in enumerate(words)} |
| self.uni = uni if bn is None else [lp for n, lp in zip(bn, blp) for _ in range(n)] |
| self.bi = {}; self.bow = {} |
| if lm_path: |
| b = open(lm_path, "rb").read(); N2, nbg, l1, l2, nh = struct.unpack("<IIIII", b[:20]); pos = 20 |
| cb = np.frombuffer(b, np.float32, 16, pos); pos += 64; cbb = np.frombuffer(b, np.float32, 16, pos); pos += 64 |
| cnt, pos = read_varints(b, N + 1, pos); delta, pos = read_varints(b, nbg, pos) |
| q = cb[unpack4(b[pos:pos + (nbg + 1) // 2], nbg)].tolist(); pos += (nbg + 1) // 2 |
| qbw = cbb[unpack4(b[pos:pos + (nh + 1) // 2], nh)].tolist() |
| i = 0; hi = 0 |
| for h, n in enumerate(cnt): |
| if not n: continue |
| self.bow[h] = qbw[hi]; hi += 1; acc = -1 |
| for _ in range(n): |
| acc = delta[i] if acc < 0 else acc + delta[i] + 1 |
| self.bi[h * HK + acc] = q[i]; i += 1 |
|
|
| def unigram(self, wid): |
| return self.unk if wid is None else (math.log(0.04) if wid == self.N else self.uni[wid]) |
|
|
| def logp(self, prev, wid): |
| """ln P(wid | prev). prev None = sentence start; wid None = OOV word; wid N = </s>.""" |
| h = self.N if prev is None else prev |
| if wid is not None: |
| p = self.bi.get(h * HK + wid) |
| if p is not None: return p |
| return self.bow.get(h, 0.0) + self.unigram(wid) |
|
|
|
|
| def write_assets(out_dir, n_lex, maxbi): |
| import os |
| root = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")) |
| rows = [l.rstrip("\n").split("\t") for l in open(f"{root}/data/processed/trackB/hi_wordfreq.tsv", encoding="utf-8")] |
| d = np.load(f"{root}/data/processed/trackB/lm/bigrams.npz") |
| lex, lm, st = build([r[0] for r in rows], [int(r[1]) for r in rows], d["pairs"], d["counts"], n_lex, maxbi) |
| os.makedirs(out_dir, exist_ok=True) |
| open(f"{out_dir}/lex.bin", "wb").write(lex); open(f"{out_dir}/lm.bin", "wb").write(lm) |
| st.update(lex_bytes=len(lex), lex_gz=len(gzip.compress(lex, 9)), lm_bytes=len(lm), lm_gz=len(gzip.compress(lm, 9))) |
| return st |
|
|
