| """Memory-light lexicon trie for constrained beam search over large (0.5-2M word) lexicons. |
| A node is a (lo, hi, depth) range of the sorted word list; children are found with bisect and cached lazily, so |
| only nodes actually visited by the beam are ever materialised. Same interface as xlit_model.Trie (kids[n], final[n]).""" |
| import bisect |
| from bhc.xlit_model import T2I, TGT_VOCAB |
|
|
|
|
| class _Kids: |
| def __init__(s, t): s.t = t |
| def __getitem__(s, n): return s.t._children(n) |
|
|
|
|
| class _Final: |
| def __init__(s, t): s.t = t |
| def __getitem__(s, n): |
| lo, hi, d = s.t.nodes[n]; return len(s.t.w[lo]) == d |
|
|
|
|
| class LazyTrie: |
| def __init__(s, words): |
| s.w = sorted(w for w in set(words) if w and all(c in T2I for c in w)) |
| s.nodes = [(0, len(s.w), 0)]; s.ids = {s.nodes[0]: 0}; s.cache = {} |
| s.kids = _Kids(s); s.final = _Final(s) |
|
|
| def _children(s, n): |
| c = s.cache.get(n) |
| if c is not None: return c |
| lo, hi, d = s.nodes[n]; w = s.w; out = {} |
| pos = lo |
| if pos < hi and len(w[pos]) == d: pos += 1 |
| while pos < hi: |
| ch = w[pos][d]; pre = w[pos][:d + 1] |
| end = bisect.bisect_left(w, pre[:-1] + chr(ord(ch) + 1), pos, hi) |
| key = (pos, end, d + 1); nid = s.ids.get(key) |
| if nid is None: nid = len(s.nodes); s.nodes.append(key); s.ids[key] = nid |
| out[T2I[ch]] = nid; pos = end |
| if len(s.cache) > 400000: s.cache.clear() |
| s.cache[n] = out; return out |
|
|
