File size: 10,944 Bytes
843b037 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 | """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
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"]
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) # 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.name = "bhc-xlit2"
# ---------------- candidates with features
def cand_feats(self, words):
p = self.p; lw = [w.lower() 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
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))]
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"])]
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)]
|