import math import pathlib import numpy as np np.seterr(all="ignore") # ---------------- boundary scheme ---------------- #----------------------------------# import re LABELS = ["join", "space", "newline", "para"] GAP_STR = {"join": "", "space": " ", "newline": "\n", "para": "\n\n"} PARA = 0x02 # one symbol for "\n\n" GAP_BYTES = {"join": b"", "space": b" ", "newline": b"\n", "para": bytes([PARA])} def parse(text): text = text.strip() return re.split(r"\s+", text), re.findall(r"\s+", text) def gap_label(gap): n= gap.count("\n") return "para" if n>=2 else "newline" if n==1 else "space" if n==0 else "join" def labels_of(text): return [gap_label(g) for g in parse(text)[1]] def render(chunks, labels): out = [chunks[0]] for chunk, lab in zip(chunks[1:], labels): out += [GAP_STR[lab], chunk] return "".join(out) def skeleton(text): return "".join(text.split()) def canonical(text): """Whitespace runs collapsed to canonical bytes, paragraph as one symbol.""" out, run = [], "" for ch in text.strip(): if ch.isspace(): run += ch else: if run: n = run.count("\n") out.append(bytes([PARA]) if n >= 2 else b"\n" if n == 1 else b" ") run = "" out.append(ch.encode("utf-8")) return b"".join(out) # ---------------- n-gram + Viterbi ---------------- #----------------------------------# FNV = np.uint64(1099511628211) def _gram_hashes(arr, k): h = np.zeros(len(arr) - k + 1, dtype=np.uint64) for j in range(k): h = h * FNV + arr[j:j + len(h)].astype(np.uint64) return h class NgramViterbi: def __init__(self, order=7, backoff=0.4): self.order, self.backoff = order, math.log(backoff) self.tables, self.cache = {}, {} def fit(self, text): arr = np.frombuffer(canonical(text), dtype=np.uint8) self.total = len(arr) for k in range(1, self.order + 1): h, c = np.unique(_gram_hashes(arr, k), return_counts=True) self.tables[k] = (h, c.astype(np.float64)) return self def save(self, path): np.savez_compressed(path, total=self.total, order=self.order, **{f"h{k}": v[0] for k, v in self.tables.items()}, **{f"c{k}": v[1] for k, v in self.tables.items()}) @classmethod def load(cls, path): z = np.load(path) m = cls(order=int(z["order"])) m.total = int(z["total"]) m.tables = {k: (z[f"h{k}"], z[f"c{k}"]) for k in range(1, m.order + 1)} return m def _count(self, gram): h = np.uint64(0) for b in gram: h = h * FNV + np.uint64(b) hs, cs = self.tables[len(gram)] i = np.searchsorted(hs, h) return cs[i] if i < len(hs) and hs[i] == h else 0.0 def logp(self, ctx, b): key = (ctx, b) if key in self.cache: return self.cache[key] pen = 0.0 for k in range(min(len(ctx), self.order - 1), -1, -1): num = self._count(ctx[len(ctx) - k:] + bytes([b])) if num: den = self._count(ctx[len(ctx) - k:]) if k else self.total val = pen + math.log(num / den) break pen += self.backoff else: val = pen - math.log(self.total) if len(self.cache) < 2_000_000: self.cache[key] = val return val def _extend(self, state, seq): lp = 0.0 for b in seq: lp += self.logp(state, b) state = (state + bytes([b]))[-(self.order - 1):] return lp, state def _decode(self, text, keep): chunks, _ = parse(text) cbytes = [c.encode("utf-8") for c in chunks] _, start = self._extend(b"", cbytes[0]) beam = {start: [(0.0, [])]} for cb in cbytes[1:]: nxt = {} for state, paths in beam.items(): for lab in LABELS: d, s2 = self._extend(state, GAP_BYTES[lab] + cb) bucket = nxt.setdefault(s2, []) for lp, labs in paths: bucket.append((lp + d, labs + [lab])) beam = {s: sorted(p, reverse=True)[:keep] for s, p in nxt.items()} return chunks, sorted((p for ps in beam.values() for p in ps), reverse=True) def restore(self, text): chunks, paths = self._decode(text, keep=1) return render(chunks, paths[0][1]) def kbest(self, text, n=8): chunks, paths = self._decode(text, keep=n) return chunks, [labs for _, labs in paths[:n]] # ---------------- LSTM + beam ---------------- #----------------------------------# def _torch(): import torch return torch DEV = None def _device(): global DEV if DEV is None: torch = _torch() DEV = "mps" if torch.backends.mps.is_available() else "cpu" return DEV def char_lm(emb=64, hidden=512, layers=2): import torch.nn as nn class CharLM(nn.Module): def __init__(self): super().__init__() self.emb = nn.Embedding(256, emb) self.lstm = nn.LSTM(emb, hidden, layers, batch_first=True) self.head = nn.Linear(hidden, 256) def forward(self, x, state=None): out, state = self.lstm(self.emb(x), state) return self.head(out), state return CharLM() def load_lm(path): torch = _torch() m = char_lm() m.load_state_dict(torch.load(path, map_location=_device())) return m.to(_device()).eval() class LstmBeam: def __init__(self, model, k=8): self.model = model.to(_device()).eval() self.k = k def restore(self, text): import torch import torch.nn.functional as F dev = _device() def advance(logits, state, seq): B = logits.shape[0] x = torch.tensor(list(seq), device=dev).expand(B, -1) step, state = self.model(x, state) prev = torch.cat([logits[:, None, :], step[:, :-1, :]], 1) lp = F.log_softmax(prev, -1).gather(2, x[:, :, None]).squeeze(2).sum(1) return lp, step[:, -1, :], state with torch.no_grad(): chunks, _ = parse(text) cbytes = [c.encode("utf-8") for c in chunks] logits, state = self.model(torch.zeros(1, 1, dtype=torch.long, device=dev)) logits = logits[:, -1, :] _, logits, state = advance(logits, state, cbytes[0]) logps, labels = torch.zeros(1, device=dev), [[]] for cb in cbytes[1:]: B = logps.shape[0] h, c = state eh, ec = h.repeat_interleave(4, 1), c.repeat_interleave(4, 1) el, elp = logits.repeat_interleave(4, 0), logps.repeat_interleave(4, 0) gaps = [GAP_BYTES[lab] for lab in LABELS] lsm = F.log_softmax(el, -1) add = torch.zeros(4 * B, device=dev) xg = torch.zeros(4 * B, 1, dtype=torch.long, device=dev) has = torch.zeros(4 * B, dtype=torch.bool) for j in range(4 * B): g = gaps[j % 4] if g: has[j], xg[j, 0], add[j] = True, g[0], lsm[j, g[0]] gl, (gh, gc) = self.model(xg, (eh, ec)) gl, m = gl[:, -1, :], has.to(dev) el = torch.where(m[:, None], gl, el) eh = torch.where(m[None, :, None], gh, eh) ec = torch.where(m[None, :, None], gc, ec) elp = elp + add d, el, (eh, ec) = advance(el, (eh, ec), cb) elp = elp + d new_labels = [labels[j // 4] + [LABELS[j % 4]] for j in range(4 * B)] top = elp.topk(min(self.k, 4 * B)) logps, logits = top.values, el[top.indices] state = (eh[:, top.indices, :].contiguous(), ec[:, top.indices, :].contiguous()) labels = [new_labels[j] for j in top.indices.tolist()] return render(chunks, labels[0]) # ---------------- rerank ---------------- #----------------------------------# class Rerank: def __init__(self, ngram, lm, k=8): self.ngram, self.lm, self.k = ngram, lm.to(_device()).eval(), k def _score(self, seqs): import torch import torch.nn.functional as F dev = _device() L = max(len(b) for b in seqs) x = torch.zeros(len(seqs), L + 1, dtype=torch.long, device=dev) for r, b in enumerate(seqs): x[r, 1:len(b) + 1] = torch.tensor(list(b)) with torch.no_grad(): logits, _ = self.lm(x[:, :-1]) lp = F.log_softmax(logits, -1).gather(2, x[:, 1:, None]).squeeze(2) return (lp * (x[:, 1:] != 0)).sum(1) def restore(self, text): chunks, paths = self.ngram.kbest(text, self.k) cands = [render(chunks, labs) for labs in paths] return cands[int(self._score([canonical(c) for c in cands]).argmax())] # --- loader --------------------------------------------------------------# #----------------------------------# MODELS = pathlib.Path(__file__).parent / "models" def load(name="rerank", models_dir=MODELS): models_dir = pathlib.Path(models_dir) if name == "ngram_viterbi": return NgramViterbi.load(models_dir / "ngram.npz") if name == "lstm_beam": return LstmBeam(load_lm(models_dir / "lstm.pt")) if name == "rerank": return Rerank(NgramViterbi.load(models_dir / "ngram.npz"), load_lm(models_dir / "lstm.pt")) raise ValueError(f"unknown model: {name}")