twlat/runtime_v3.py

"""TWLAT V3 推論管線。

    輸入 → protect（PUA 遮罩保護段）→ safe_normalize(twv)
         → lattice（字典硬過濾在此發生）→ 無邊 → fast path
         → clean/masked 雙 pass forward → Viterbi + τ → 最小 splice
         → restore → 輸出

超長輸入：滑動視窗（stride 384），每條邊由「它最居中」的視窗評分——
不再有 V2 的 INPUT_TOO_LONG 降級路徑。
"""
from __future__ import annotations

import pathlib
from dataclasses import dataclass, field

import numpy as np
import torch


from twlat import decoder                                        # noqa: E402
from twlat.features import (FEAT_DIM, L_MAX, LexTables,          # noqa: E402
                            assemble_cands, make_feat, site_arrays,
                            text_arrays)
from twlat.lattice import LatticeBuilder, build_masked_view      # noqa: E402
from twlat.model_v3 import TWLATV3, make_config                  # noqa: E402
from twlat.normalize import safe_normalize                       # noqa: E402
from twlat.paths import data_file                                # noqa: E402
from twlat.protect import protect, restore                       # noqa: E402

import json                                                       # noqa: E402

WIN, STRIDE = 512, 384


@dataclass
class ResultV3:
    output: str
    decisions: list = field(default_factory=list)
    fast_path: bool = False
    error: str | None = None


class TWLATV3Runtime:
    def __init__(self, ckpt: str, device: str | None = None,
                 tau: dict[str, float] | None = None,
                 lexicon_path=None, fo_bonus: float = 0.0):
        lex_p = lexicon_path or data_file("dict/lattice_lexicon.json")
        self.lb = LatticeBuilder(lex_p)
        self.lex = LexTables(lex_p)
        self.vocab = json.loads(
            data_file("dict/char_vocab_v3.json").read_text(encoding="utf-8"))
        self.fo_bonus = fo_bonus
        tp = data_file("model/tau.json")
        self.tau = tau if tau is not None else (
            json.loads(tp.read_text()) if tp.exists() else decoder.DEFAULT_TAU)
        self.device = torch.device(
            device or ("cuda" if torch.cuda.is_available()
                       else "mps" if torch.backends.mps.is_available() else "cpu"))
        ck = torch.load(ckpt, map_location="cpu", weights_only=False)
        self.model = TWLATV3(make_config(**ck["mcfg"]))
        self.model.load_state_dict(ck["model"])
        self.model.eval().to(self.device)
        # observed 相依特徵（is_keep/長度差）只有 finetune 階段見過；
        # 對 pretrain-only checkpoint 餵入會是分佈外輸入
        self.reveal = ck.get("phase") == "finetune"
        ck_lex = ck.get("lexicon_version")
        if ck_lex and ck_lex != self.lb.version:
            # 熱更新是功能不是錯誤：記錄但不擋（held-out 評測依賴此路徑）
            self.lexicon_mismatch = (ck_lex, self.lb.version)
        else:
            self.lexicon_mismatch = None

    # ------------------------------------------------------------------ #

    @torch.no_grad()
    def score_batch(self, texts: list[str], batch_size: int = 32):
        """評分階段：→ list[(base, sv, lat, logits[n_edges,C] | None)]。
        τ 校準重用此結果做多組解碼，不必重跑模型。"""
        stages = []
        jobs = []
        for ti, t in enumerate(texts):
            q, sv = protect(t)
            base = safe_normalize(q, do_twv=True)
            lat = self.lb.build(base)
            stages.append([base, sv, lat, None])
            if not lat.edges:
                continue
            wins = self._windows(len(base))
            assign = self._assign(lat, wins)
            for wi, (ws, we) in enumerate(wins):
                eixs = assign[wi]
                if not eixs:
                    continue
                jobs.append(self._make_job(ti, base[ws:we], ws,
                                           [lat.edges[i] for i in eixs], eixs))

        logit_map: dict[int, dict[int, np.ndarray]] = {}
        jobs.sort(key=lambda j: len(j["text"]))
        for s in range(0, len(jobs), batch_size):
            chunk = jobs[s:s + batch_size]
            out = self.model(self._stack(chunk))
            lg = out["cand_logits"].float().cpu().numpy()
            for bi, j in enumerate(chunk):
                m = logit_map.setdefault(j["ti"], {})
                for k, eix in enumerate(j["eixs"]):
                    m[eix] = lg[bi, k]

        for ti, st in enumerate(stages):
            lat = st[2]
            if not lat.edges:
                continue
            lm = logit_map.get(ti, {})
            cmax = max((v.shape[0] for v in lm.values()), default=1)
            logits = np.full((len(lat.edges), cmax), -np.inf, np.float32)
            for eix, row in lm.items():
                logits[eix, :len(row)] = row
            for i, e in enumerate(lat.edges):
                if i not in lm:                  # 不應發生：未覆蓋 → keep
                    logits[i, e.obs_ix] = 0.0
            st[3] = logits
        return stages

    def decode_stage(self, stage, tau=None) -> ResultV3:
        base, sv, lat, logits = stage
        if logits is None:
            return ResultV3(output=restore(base, sv), fast_path=True)
        edits = decoder.decode(self.lb, lat, logits,
                               self.tau if tau is None else tau,
                               fo_bonus=self.fo_bonus)
        return ResultV3(
            output=restore(decoder.splice(base, edits), sv),
            decisions=[{"span": [e.start, e.end], "from": e.observed,
                        "to": e.replacement, "utility": round(e.utility, 4),
                        "rule_type": e.rule_type} for e in edits])

    @torch.no_grad()
    def convert_batch(self, texts: list[str], batch_size: int = 32
                      ) -> list[ResultV3]:
        return [self.decode_stage(s)
                for s in self.score_batch(texts, batch_size)]

    # ------------------------------------------------------------------ #

    @staticmethod
    def _windows(n: int) -> list[tuple[int, int]]:
        if n <= WIN:
            return [(0, n)]
        wins, s = [], 0
        while True:
            wins.append((s, min(s + WIN, n)))
            if s + WIN >= n:
                return wins
            s += STRIDE

    @staticmethod
    def _assign(lat, wins) -> dict[int, list[int]]:
        """每條邊 → 它最居中的視窗。"""
        out: dict[int, list[int]] = {wi: [] for wi in range(len(wins))}
        for i, e in enumerate(lat.edges):
            best_wi, best_d = None, None
            for wi, (ws, we) in enumerate(wins):
                if e.start >= ws and e.end <= we:
                    c = (ws + we) / 2
                    d = abs((e.start + e.end) / 2 - c)
                    if best_d is None or d < best_d:
                        best_wi, best_d = wi, d
            if best_wi is not None:
                out[best_wi].append(i)
        return out

    def _make_job(self, ti, text, offset, edges, eixs):
        ids, script, prot = text_arrays(text, self.vocab)
        # 邊座標平移到視窗座標
        shifted = []
        for e in edges:
            se = type(e)(e.start - offset, e.end - offset, e.gid, e.obs_ix,
                         e.cand_kill, e.word_crossing, e.word_contained)
            shifted.append(se)
        sa = site_arrays(self.lb, shifted, text)
        mids, m_span, _ = build_masked_view(ids, sa["span"], sa["maskable"])
        mscript, _, _ = build_masked_view(script, sa["span"], sa["maskable"],
                                          mask_id=4)
        mprot, _, _ = build_masked_view(prot.astype(np.uint8), sa["span"],
                                        sa["maskable"], mask_id=0)
        cand_tok, cand_mask, cand_kill, cand_feat = assemble_cands(
            self.lex, sa["gid"], sa["obs"], sa["clue"], sa["eng"],
            sa["flags"], sa["kill"], reveal_observed=self.reveal)
        return {"ti": ti, "eixs": eixs, "text": text,
                "ids": ids, "script": script, "prot": prot,
                "feat": make_feat(script, prot, sa["span"], len(text)),
                "mids": mids, "m_span": m_span.astype(np.int64),
                "mfeat": make_feat(mscript, mprot, m_span, len(mids)),
                "c_span": sa["span"], "obs": sa["obs"],
                "cand_tok": cand_tok, "cand_mask": cand_mask,
                "cand_kill": cand_kill, "cand_feat": cand_feat}

    def _stack(self, jobs):
        B = len(jobs)
        T = max(len(j["ids"]) for j in jobs)
        Tm = max(len(j["mids"]) for j in jobs)
        S = max(len(j["c_span"]) for j in jobs)
        C = max(j["cand_tok"].shape[1] for j in jobs)
        out = {
            "ids": np.zeros((B, T), np.int64),
            "feat": np.zeros((B, T, 4), np.int64),
            "pad": np.zeros((B, T), bool),
            "mids": np.zeros((B, Tm), np.int64),
            "mfeat": np.zeros((B, Tm, 4), np.int64),
            "mpad": np.zeros((B, Tm), bool),
            "c_span": np.zeros((B, S, 2), np.int64),
            "m_span": np.zeros((B, S, 2), np.int64),
            "site_mask": np.zeros((B, S), bool),
            "cand_tok": np.zeros((B, S, C, L_MAX), np.int64),
            "cand_mask": np.zeros((B, S, C), bool),
            "cand_kill": np.zeros((B, S, C), bool),
            "cand_feat": np.zeros((B, S, C, FEAT_DIM), np.float32),
        }
        for b, j in enumerate(jobs):
            n, tm = len(j["ids"]), len(j["mids"])
            ns, c = len(j["c_span"]), j["cand_tok"].shape[1]
            out["ids"][b, :n] = j["ids"]
            out["feat"][b, :n] = j["feat"]
            out["pad"][b, :n] = True
            out["mids"][b, :tm] = j["mids"]
            out["mfeat"][b, :tm] = j["mfeat"]
            out["mpad"][b, :tm] = True
            out["c_span"][b, :ns] = j["c_span"]
            out["m_span"][b, :ns] = j["m_span"]
            out["site_mask"][b, :ns] = True
            out["cand_tok"][b, :ns, :c] = j["cand_tok"]
            out["cand_mask"][b, :ns, :c] = j["cand_mask"]
            out["cand_kill"][b, :ns, :c] = j["cand_kill"]
            out["cand_feat"][b, :ns, :c] = j["cand_feat"]
        return {k: torch.from_numpy(v).to(self.device) for k, v in out.items()}