engine.py

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