Blame for python/bhc/translit2.py

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