"""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"] # v2 casual features (append-only; older params use the first len(mu)) 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) # skel/visr on via params.json (v3 assets) 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) # core-only: no dictionary pack -> no constrained beam / skeleton index / variants 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" # ---------------- candidates with features 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): # Aksharantar convention: literal word-final h -> visarga for s in sorted((x for x in d if d[x]["sc"] is not None), key=lambda x: -d[x]["sc"])[:2]: # beam entries only 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 # vowel gate 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) # untuned fallback z = [(x - m) / s for x, m, s in zip(fv, p["mu"], p["sd"])] # zip truncates: old 24-feature params ignore v2 features if p.get("w1"): # 1-hidden-layer tanh MLP reranker (fit on VALID only): w2 . tanh(W1 z + b1) + b2 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 # ---------------- sentence mode: Viterbi over candidates with the bigram LM 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)]