| """BhashaCore Track B v2 (commercially-usable data only): distilled char Transformer (4/8-bit weights) -> |
| candidates (free beam + lexicon-constrained beam + lexicon-gated ऑ/ः/nukta variants + letter-name acronym) |
| -> linear reranker (weights fit on Aksharantar VALID only) for word mode; Viterbi with the IndicCorp bigram LM |
| for sentence mode. Assets dir: xlit.q4|xlit.q8|xlit.int8, lex.bin, lm.bin, params.json.""" |
| import json, math, os, re, sys |
| import torch |
| ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..")); sys.path.insert(0, ROOT) |
| from bhc.xlit_model import beam_search_fast as beam_search, score_pairs |
| from bhc.xlit_q import load_any |
| from bhc.lexlm import LexLM |
| from bhc.lextrie import LazyTrie |
| from bhc.text import TOKEN_RE |
| from bhc.skel import SkelIndex, rom_skels, dev_skel, dev_rom, lev |
| from bhc.casual import CasualIndex, ABBR, squeeze_in |
|
|
| LATIN = re.compile(r"[A-Za-z]+$") |
| LETTER = dict(a="ए", b="बी", c="सी", d="डी", e="ई", f="एफ", g="जी", h="एच", i="आई", j="जे", k="के", l="एल", m="एम", n="एन", |
| o="ओ", p="पी", q="क्यू", r="आर", s="एस", t="टी", u="यू", v="वी", w="डब्ल्यू", x="एक्स", y="वाई", z="जेड") |
| FEATS = ["sc", "scn", "inlex", "uni", "free0", "infree", "incon", "aw", "visarga", "nukta", "chandra", "acro", "variant", "frank", |
| "skel", "skmatch", "lev", "sc_rel", "scn_rel", "uni_rel", "lenr", "visr", "hfin", "acro_nv", |
| "ck1", "ck2", "ckr", "abbr", "abbr_r", "novow"] |
| VIS_RE = re.compile(r"[aeiou]h$") |
| VAR_SWAPS = [("ॉ", "ो"), ("ो", "ॉ"), ("ा", "ॉ"), ("ॉ", "ा"), ("ः", ""), ("फ़", "फ"), ("फ", "फ़"), ("ज़", "ज"), ("ज", "ज़"), ("ँ", "ं"), ("ं", "ँ")] |
| DEFAULTS = dict(beam=8, nbest=8, cbeam=8, variants=True, acro=True, w=None, mu=None, sd=None, |
| w_lm=0.8, oov_pen_s=4.0, w_ch_s=1.0, use_lm=True, |
| skel=False, skel_k=200, skel_m=16, skel_wu=0.03, visr=False, acro_v=False, |
| casual=False, cas_top=300000, cas_n1=6, cas_n2=4, abbr=False, squeeze=0, k1gate=99, k2gate=99) |
|
|
|
|
| class BhcTranslit2: |
| def __init__(self, asset_dir, threads=1, use_dict=True, **over): |
| """asset_dir: flat dir (model + params.json [+ lex.bin + lm.bin]) or a release bundle: core (model + params.json) |
| with the dictionary pack in asset_dir/dict/ (lex.bin, lm.bin, params.json overriding the core params). |
| use_dict=False (or no lexicon present) -> core-only mode: free beam + rule candidates, core reranker.""" |
| torch.set_num_threads(threads); self.dir = asset_dir; self.p = dict(DEFAULTS) |
| if os.path.exists(f"{asset_dir}/params.json"): self.p.update(json.load(open(f"{asset_dir}/params.json"))) |
| pack = f"{asset_dir}/dict"; self.dict_dir = None |
| if use_dict and os.path.exists(f"{pack}/lex.bin"): |
| self.dict_dir = pack |
| if os.path.exists(f"{pack}/params.json"): self.p.update(json.load(open(f"{pack}/params.json"))) |
| self.p.update(over) |
| mf = [f for f in ("xlit.q4", "xlit.q8", "xlit.int8", "model.pt") if os.path.exists(f"{asset_dir}/{f}")] |
| self.model_file = f"{asset_dir}/{self.p.get('model_file', mf[0])}"; self.model = load_any(self.model_file) |
| ld = self.dict_dir or asset_dir; lm, lx = f"{ld}/lm.bin", f"{ld}/lex.bin" |
| if os.path.exists(lx) and use_dict: self.lm = LexLM(lx, lm if os.path.exists(lm) and self.p["use_lm"] else None) |
| else: self.lm = LexLM(None) |
| if not self.lm.N: self.p.update(cbeam=0, skel=False) |
| self.trie = LazyTrie(self.lm.words) if self.p["cbeam"] else None |
| self.skix = SkelIndex(self.lm.words, self.lm.unigram) if self.p.get("skel") else None |
| self.cix = CasualIndex(self.lm.words, self.lm.unigram, self.p["cas_top"]) if self.p.get("casual") and self.lm.N else None |
| self.name = "bhc-xlit2" |
|
|
| |
| def cand_feats(self, words): |
| p = self.p; lw = [squeeze_in(w.lower(), p["squeeze"]) for w in words]; todo = sorted(set(lw)); n = len(todo) |
| if p["cbeam"]: |
| r = beam_search(self.model, todo + todo, beam=p["beam"], nbest=p["nbest"], trie=self.trie, constrain=[False] * n + [True] * n) |
| free, con = r[:n], r[n:] |
| else: |
| free = beam_search(self.model, todo, beam=p["beam"], nbest=p["nbest"]); con = [[] for _ in todo] |
| res = {}; need = [] |
| for w, f, c in zip(todo, free, con): |
| d = {} |
| for i, (s, sc) in enumerate(f): |
| if s: d[s] = dict(sc=sc, infree=1, frank=i, free0=int(i == 0)) |
| for s, sc in c: |
| if s: |
| e = d.setdefault(s, dict(sc=sc, infree=0, frank=p["nbest"], free0=0)); e["incon"] = 1 |
| if p["variants"]: |
| for s in sorted(d, key=lambda x: -d[x]["sc"])[:4]: |
| for a, b in VAR_SWAPS: |
| if a in s: |
| v = s.replace(a, b) |
| if v not in d and v in self.lm.id: d[v] = dict(sc=None, infree=0, frank=p["nbest"], free0=0, variant=1); need.append((w, v)) |
| if p["acro"] and (re.fullmatch(r"[b-df-hj-np-tv-z]{2,7}", w) or (p.get("acro_v") and 2 <= len(w) <= 4)): |
| v = "".join(LETTER[ch] for ch in w) |
| if v not in d: d[v] = dict(sc=None, infree=0, frank=p["nbest"], free0=0, acro=1); need.append((w, v)) |
| else: d[v]["acro"] = 1 |
| if p.get("visr") and VIS_RE.search(w): |
| for s in sorted((x for x in d if d[x]["sc"] is not None), key=lambda x: -d[x]["sc"])[:2]: |
| v = s[:-1] + "ः" if s.endswith("ह") else (s + "ः" if not s.endswith("ः") else None) |
| if v and v not in d: d[v] = dict(sc=None, infree=0, frank=p["nbest"], free0=0, visr=1); need.append((w, v)) |
| if self.skix is not None: |
| ids = self.skix.lookup(w, p["skel_k"]); L = len(w) + 1 |
| ids = sorted(ids, key=lambda i: lev(w, dev_rom(self.lm.words[i])) / L - p["skel_wu"] * self.lm.unigram(i))[:p["skel_m"]] |
| for i in ids: |
| v = self.lm.words[i] |
| if v not in d: d[v] = dict(sc=None, infree=0, frank=p["nbest"], free0=0, skel=1); need.append((w, v)) |
| else: d[v]["skel"] = 1 |
| if self.cix is not None: |
| r1, r2 = self.cix.lookup(w, p["cas_n1"], p["cas_n2"]); nv = sum(ch in "aeiou" for ch in w) |
| ab = set(ABBR.get(w, ())) if p.get("abbr") else set(); s1 = {self.lm.words[i] for i in r1[:p["cas_n1"]]} |
| for key, ids, n_ in (("ck1", r1, p["cas_n1"]), ("ck2", r2, p["cas_n2"])): |
| for rk, i in enumerate(ids[:n_]): |
| v = self.lm.words[i] |
| if v not in d: |
| if v not in ab and ((v in s1 and nv > p["k1gate"]) or (v not in s1 and nv > p["k2gate"])): continue |
| d[v] = dict(sc=None, infree=0, frank=p["nbest"], free0=0); need.append((w, v)) |
| d[v][key] = 1; d[v]["ckr"] = min(d[v].get("ckr", 1.0), rk / n_) |
| if p.get("abbr") and w in ABBR: |
| for rk, v in enumerate(ABBR[w]): |
| if v not in d: d[v] = dict(sc=None, infree=0, frank=p["nbest"], free0=0); need.append((w, v)) |
| d[v]["abbr"] = 1; d[v]["abbr_r"] = rk |
| res[w] = d |
| if need: |
| for (w, v), sc in zip(need, score_pairs(self.model, [x[0] for x in need], [x[1] for x in need])): res[w][v]["sc"] = sc |
| out = []; memo = {} |
| for w in lw: |
| if w in memo: out.append(memo[w]); continue |
| rows = []; sks = set(rom_skels(w)); L = len(w) + 1; vis = bool(VIS_RE.search(w)) |
| for s, e in res[w].items(): |
| wid = self.lm.id.get(s) |
| fv = [e["sc"], e["sc"] / (len(s) + 1), int(wid is not None), self.lm.unigram(wid) if wid is not None else 0.0, |
| e.get("free0", 0), e.get("infree", 0), e.get("incon", 0), int("ॉ" in s or "ऑ" in s), int("ः" in s), |
| int("़" in s), int("ँ" in s), e.get("acro", 0), e.get("variant", 0), e.get("frank", p["nbest"]) / p["nbest"], |
| e.get("skel", 0), int(dev_skel(s) in sks), lev(w, dev_rom(s)) / L, 0.0, 0.0, 0.0, len(s) / L, |
| e.get("visr", 0), int(vis and s.endswith("ः")), int(e.get("acro", 0) and not re.search("[aeiou]", w)), |
| e.get("ck1", 0), e.get("ck2", 0), e.get("ckr", 1.0), e.get("abbr", 0), e.get("abbr_r", 3), int(not re.search("[aeiouy]", w))] |
| rows.append((s, fv, wid)) |
| if rows: |
| msc = max(r[1][0] for r in rows); mscn = max(r[1][1] for r in rows) |
| mu = max((r[1][3] for r in rows if r[1][2]), default=0.0) |
| for _, fv, _ in rows: fv[17] = fv[0] - msc; fv[18] = fv[1] - mscn; fv[19] = fv[3] - mu if fv[2] else 0.0 |
| memo[w] = rows; out.append(rows) |
| return out |
|
|
| def score(self, fv): |
| p = self.p |
| if not p.get("w1") and p["w"] is None: return fv[0] + 0.2 * (fv[3] if fv[2] else -25.0) |
| z = [(x - m) / s for x, m, s in zip(fv, p["mu"], p["sd"])] |
| if p.get("w1"): |
| return sum(v * math.tanh(sum(a * b for a, b in zip(row, z)) + c) for v, row, c in zip(p["w2"], p["w1"], p["b1"])) + p.get("b2", 0.0) |
| return sum(w * x for w, x in zip(p["w"], z)) |
|
|
| def translit_words(self, words, k=8): |
| out = [] |
| for w, rows in zip(words, self.cand_feats(words)): |
| if not LATIN.match(w) or not rows: out.append([w]); continue |
| out.append([s for s, _, _ in sorted(rows, key=lambda r: -self.score(r[1]))[:k]]) |
| return out |
|
|
| |
| def decode_slots(self, slots): |
| xs = [t for kind, t in slots if kind == "x"]; cl = iter(self.cand_feats(xs)) if xs else iter(()) |
| p = self.p; lm = self.lm; lat = [] |
| for kind, t in slots: |
| if kind == "x": |
| rows = next(cl); sc = [(s, self.score(fv), wid) for s, fv, wid in rows] |
| if sc: |
| mx = max(x[1] for x in sc); lse = mx + math.log(sum(math.exp(x[1] - mx) for x in sc)) |
| sc = sorted(((s, p["w_ch_s"] * (v - lse), wid) for s, v, wid in sc), key=lambda x: -x[1])[:12] |
| lat.append(sc or [(t, 0.0, None)]) |
| else: |
| lat.append([(t, 0.0, lm.id.get(t, -1) if re.match(r"[\u0900-\u0963\u0971-\u097F]+$", t) else -2)]) |
| if not p["use_lm"] or not lm.bi: |
| return [max(o, key=lambda x: x[1])[0] for o in lat] |
| beams = [(0.0, None, [])] |
| for opts in lat: |
| if opts[0][2] == -2 and len(opts) == 1: |
| beams = [(s, h, ch + [opts[0][0]]) for s, h, ch in beams]; continue |
| best = {} |
| for s, h, ch in beams: |
| for cand, csc, wid in opts: |
| w = None if wid in (None, -1) else wid |
| tot = s + csc + p["w_lm"] * (lm.logp(h, w) - (p["oov_pen_s"] if w is None else 0.0)) |
| if cand not in best or tot > best[cand][0]: best[cand] = (tot, w, ch + [cand]) |
| beams = sorted(best.values(), key=lambda x: -x[0])[:16] |
| return max(beams, key=lambda b: b[0] + p["w_lm"] * lm.logp(b[1], lm.N))[2] |
|
|
| def convert(self, text, policy=None): |
| spans = [(m.start(), m.end(), m.group()) for m in TOKEN_RE.finditer(text)] |
| slots = [("x", t) if LATIN.match(t) else ("fix", t) for _, _, t in spans] |
| out = self.decode_slots(slots) |
| if out and out[-1] == ".": out[-1] = "।" |
| s, pos = [], 0 |
| for (a, b, _), o in zip(spans, out): s.append(text[pos:a]); s.append(o); pos = b |
| return "".join(s) + text[pos:] |
|
|
| def asset_files(self): |
| fs = [self.model_file, f"{self.dir}/lex.bin", f"{self.dir}/params.json"] |
| if self.p["use_lm"] and os.path.exists(f"{self.dir}/lm.bin"): fs.append(f"{self.dir}/lm.bin") |
| return [f for f in fs if os.path.exists(f)] |
|
|
