Twinity-1: weights, compiled dictionary, inference code

6cc3500 verified 3 months ago

19.7 kB

"""TWLAT-R（V2）模型（《04 實驗設計》§4）。
任務：對文本中每個 proposal（字典判定「這裡可能要改」）在候選集中選一個。
候選 index 0 **永遠是「維持原樣」**，即文本中實際出現的形式。
**核心約束（V2 的全部重點）**：
    禁止任何 per-candidate / per-site 的 trainable embedding lookup。
    候選只能由「表面字串 + 數值特徵」動態編碼，且與文本共用同一份 char embedding。
    因此新增字典條目不需要新增任何參數，held-out proposal 也不是隨機向量。
    （V1 的 `cand_emb = nn.Embedding(4096, 192)` 佔 13% 參數並阻斷 zero-shot，就是要修掉的。）
結構：
    char_emb（共用）──┬─→ context encoder ──→ span mean-pool ──→ h_i   [B,P,D]
                      │
                      └─→ 候選字串 mean-pool ──→ proj ──────────→ e_c   [B,P,C,D]
    score = MLP([h_i ⊕ e_c ⊕ (h_i * e_c) ⊕ cand_feat]) → [B,P,C]
context encoder 可切換（`TWLATRConfig.encoder`），兩者參數量刻意對齊以便做 scaling curve：
  - "tcn"   ：4 層 dilated depthwise separable Conv1d（dilation 1/2/4/8、kernel 5）
  - "local" ：4 層 local-window Transformer（window 半徑 32、RoPE、pre-LN、GELU）
batch 欄位見 `TWLATR.forward` docstring。
"""

from __future__ import annotations

from dataclasses import dataclass
from typing import Any

import torch
import torch.nn as nn
import torch.nn.functional as F

# Loss 權重（《04》§0 的錯誤分析：73% 的錯是「改了不該改」，故 keep_bias 直接壓它）
W_CAND = 1.0
W_KEEP_BIAS = 0.3
KEEP_MARGIN = 0.5
KEEP_INDEX = 0  # 候選 0 恆為「維持原樣」

ENCODERS = ("tcn", "local")


@dataclass
class TWLATRConfig:
    """TWLAT-R 配置。目標參數量 3.0M–4.5M。
    註：d_model=160 只有約 2.5M（低於下限），故預設起跳為 192；
    tools/param_count_r.py 會印出兩者的對照。
    """

    d_model: int = 192
    n_heads: int = 4
    ffn_dim: int = 768
    dropout: float = 0.1

    # context encoder：可切換，兩種都必須可跑
    encoder: str = "local"
    n_layers: int = 4
    local_window: int = 32  # 半徑；|i-j| <= 32 才可見
    conv_kernel: int = 5
    conv_dilations: tuple[int, ...] = (1, 2, 4, 8)
    # depthwise 之後的 pointwise 通道擴張倍率；設 1 即教科書式 depthwise separable，
    # 設 2 可讓 tcn 與 local 的 mixer 參數量幾乎相等（scaling curve 才公平）
    conv_expansion: int = 2

    # char embedding（文本與候選共用；無任何 per-ID 候選表）
    han_vocab: int = 4000  # 常用字直接查表；id >= han_vocab 走 hash
    n_hash: int = 2
    hash_buckets: int = 2048
    feat_dim: int = 4  # script / 是否在 proposal span 內 / 是否保護段 / 詞界
    feat_vocab: tuple[int, ...] = (16, 4, 4, 4)

    # proposal / candidate 形狀
    max_props: int = 48  # P
    max_cands: int = 8  # C
    cand_len: int = 6  # L，候選表面字串的最大字元數
    cand_feat_dim: int = 12  # K
    score_hidden: int = 512
    # 候選表徵方式：
    #   "dynamic" —— 由表面字串經共用 char_emb 組合而成（V2 預設，可 zero-shot）
    #   "id"      —— per-candidate learned lookup（H-B 的對照組，複製 V1 的失敗模式）
    candidate_encoder: str = "dynamic"
    cand_id_vocab: int = 4096

    seq_len: int = 512
    rope_base: float = 10000.0

    def __post_init__(self) -> None:
        assert self.encoder in ENCODERS, f"encoder 必須是 {ENCODERS}"
        assert self.d_model % self.n_heads == 0
        assert self.d_model % 2 == 0, "hash_dim = d_model // 2，需為偶數"
        assert len(self.feat_vocab) == self.feat_dim
        assert self.conv_kernel % 2 == 1, "kernel 需為奇數才能等長 padding"

    @property
    def head_dim(self) -> int:
        return self.d_model // self.n_heads

    @property
    def hash_dim(self) -> int:
        return self.d_model // 2

    def dilation_at(self, i: int) -> int:
        return self.conv_dilations[i % len(self.conv_dilations)]


# --------------------------------------------------------------------------- #
# RoPE
# --------------------------------------------------------------------------- #


def build_rope_cache(
    seq_len: int, head_dim: int, base: float, device, dtype
) -> tuple[torch.Tensor, torch.Tensor]:
    """回傳 [T, head_dim//2] 的 cos / sin。"""
    half = head_dim // 2
    inv_freq = base ** (-torch.arange(half, device=device, dtype=torch.float32) / half)
    pos = torch.arange(seq_len, device=device, dtype=torch.float32)
    freqs = torch.outer(pos, inv_freq)
    return freqs.cos().to(dtype), freqs.sin().to(dtype)


def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
    """x: [B, H, T, D]，對相鄰兩維做旋轉。"""
    x_even, x_odd = x[..., 0::2], x[..., 1::2]
    cos = cos[None, None, : x.shape[-2], :]
    sin = sin[None, None, : x.shape[-2], :]
    out = torch.stack([x_even * cos - x_odd * sin, x_even * sin + x_odd * cos], dim=-1)
    return out.flatten(-2)


# --------------------------------------------------------------------------- #
# Embedding：文本與候選共用
# --------------------------------------------------------------------------- #

_HASH_MULT = (2654435761, 40503)
_HASH_ADD = (0, 987654321)


class CharEmbedding(nn.Module):
    """常用字 4000 直接查表；id >= han_vocab 的罕字用 2 組 hash（buckets 2048, dim d/2）
    串接後投影。
    **文本與候選字串共用這一份**：候選只是一串字元 id，沒有自己的 embedding 表，
    所以字典新增條目不會增加任何參數，未見過的字串也落在同一個表徵空間。
    """

    def __init__(self, cfg: TWLATRConfig):
        super().__init__()
        self.cfg = cfg
        self.han = nn.Embedding(cfg.han_vocab, cfg.d_model)
        self.hash = nn.ModuleList(
            nn.Embedding(cfg.hash_buckets, cfg.hash_dim) for _ in range(cfg.n_hash)
        )
        self.hash_proj = nn.Linear(cfg.n_hash * cfg.hash_dim, cfg.d_model)

    def forward(self, ids: torch.Tensor) -> torch.Tensor:
        """ids: 任意形狀 [...]，回傳 [..., d_model]。"""
        cfg = self.cfg
        ids = ids.clamp(min=0)
        rare = ids >= cfg.han_vocab
        han = self.han(ids.clamp(max=cfg.han_vocab - 1))
        parts = []
        for i, emb in enumerate(self.hash):
            m = _HASH_MULT[i % len(_HASH_MULT)]
            a = _HASH_ADD[i % len(_HASH_ADD)]
            parts.append(emb((ids * m + a) % cfg.hash_buckets))
        rare_vec = self.hash_proj(torch.cat(parts, dim=-1))
        return torch.where(rare.unsqueeze(-1), rare_vec, han)


class TextFeatures(nn.Module):
    """文本側的 4 個離散特徵；char embedding 由外部傳入，以免共用的表被重複註冊。"""

    def __init__(self, cfg: TWLATRConfig):
        super().__init__()
        self.cfg = cfg
        self.feat = nn.ModuleList(nn.Embedding(n, cfg.d_model) for n in cfg.feat_vocab)
        self.ln = nn.LayerNorm(cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x: torch.Tensor, feat: torch.Tensor) -> torch.Tensor:
        for i, emb in enumerate(self.feat):
            x = x + emb(feat[..., i].clamp(0, self.cfg.feat_vocab[i] - 1))
        return self.drop(self.ln(x))


# --------------------------------------------------------------------------- #
# Context encoder（可切換）
# --------------------------------------------------------------------------- #


class LocalAttention(nn.Module):
    """local-window self-attention + RoPE；|i-j| <= local_window 才可見。"""

    def __init__(self, cfg: TWLATRConfig):
        super().__init__()
        self.cfg = cfg
        self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
        self.out = nn.Linear(cfg.d_model, cfg.d_model)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, attn_mask, rope):
        b, t, _ = x.shape
        h, d = self.cfg.n_heads, self.cfg.head_dim
        q, k, v = self.qkv(x).view(b, t, 3, h, d).permute(2, 0, 3, 1, 4).unbind(0)
        q, k = apply_rope(q, *rope), apply_rope(k, *rope)
        p = self.cfg.dropout if self.training else 0.0
        y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=p)
        y = y.transpose(1, 2).reshape(b, t, self.cfg.d_model)
        return self.drop(self.out(y))


class ConvMixer(nn.Module):
    """dilated depthwise separable Conv1d：depthwise(k, dilation) → pointwise 擴張 → GELU
    → pointwise 還原。padding 位置先歸零，避免 pad 洩漏進感受野。"""

    def __init__(self, cfg: TWLATRConfig, dilation: int):
        super().__init__()
        d, k = cfg.d_model, cfg.conv_kernel
        pad = dilation * (k - 1) // 2  # 等長輸出
        mid = d * cfg.conv_expansion
        self.dw = nn.Conv1d(d, d, k, padding=pad, dilation=dilation, groups=d)
        self.pw1 = nn.Conv1d(d, mid, 1)
        self.pw2 = nn.Conv1d(mid, d, 1)
        self.drop = nn.Dropout(cfg.dropout)

    def forward(self, x, pad_mask, rope=None):
        z = (x * pad_mask.unsqueeze(-1).to(x.dtype)).transpose(1, 2)
        z = self.pw2(F.gelu(self.pw1(self.dw(z))))
        return self.drop(z.transpose(1, 2))


class EncoderBlock(nn.Module):
    """pre-LN：mixer（local attention 或 dilated conv）+ FFN。"""

    def __init__(self, cfg: TWLATRConfig, layer_idx: int):
        super().__init__()
        self.ln1 = nn.LayerNorm(cfg.d_model)
        if cfg.encoder == "local":
            self.mixer: nn.Module = LocalAttention(cfg)
        else:
            self.mixer = ConvMixer(cfg, cfg.dilation_at(layer_idx))
        self.ln2 = nn.LayerNorm(cfg.d_model)
        self.ffn = nn.Sequential(
            nn.Linear(cfg.d_model, cfg.ffn_dim),
            nn.GELU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(cfg.ffn_dim, cfg.d_model),
            nn.Dropout(cfg.dropout),
        )

    def forward(self, x, mix_arg, rope=None):
        x = x + self.mixer(self.ln1(x), mix_arg, rope)
        x = x + self.ffn(self.ln2(x))
        return x


# --------------------------------------------------------------------------- #
# TWLAT-R
# --------------------------------------------------------------------------- #


class TWLATR(nn.Module):
    def __init__(self, cfg: TWLATRConfig | None = None):
        super().__init__()
        self.cfg = cfg = cfg or TWLATRConfig()

        self.char_emb = CharEmbedding(cfg)
        self.text_feat = TextFeatures(cfg)

        self.layers = nn.ModuleList(
            EncoderBlock(cfg, i) for i in range(cfg.n_layers)
        )
        self.enc_ln = nn.LayerNorm(cfg.d_model)

        # 候選側：mean-pool 後只有一個 LN + 一個線性投影，沒有任何 per-ID 參數
        self.cand_ln = nn.LayerNorm(cfg.d_model)
        self.cand_proj = nn.Linear(cfg.d_model, cfg.d_model)
        # H-B 對照組：per-candidate learned lookup。未見過的候選只會取到
        # 一列未訓練的隨機向量——這正是 V1 `cand_emb = nn.Embedding(4096, 192)`
        # 的行為，用來檢驗動態編碼是否真的帶來 held-out 優勢。
        if cfg.candidate_encoder == "id":
            self.cand_id_emb = nn.Embedding(cfg.cand_id_vocab, cfg.d_model)

        score_in = 3 * cfg.d_model + cfg.cand_feat_dim
        self.score = nn.Sequential(
            nn.Linear(score_in, cfg.score_hidden),
            nn.GELU(),
            nn.Dropout(cfg.dropout),
            nn.Linear(cfg.score_hidden, 1),
        )

        self.apply(self._init_weights)
        self._rope_cache: dict[Any, tuple[torch.Tensor, torch.Tensor]] = {}
        self._window_cache: dict[Any, torch.Tensor] = {}

    @staticmethod
    def _init_weights(m: nn.Module) -> None:
        if isinstance(m, nn.Linear):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)
        elif isinstance(m, nn.Embedding):
            nn.init.normal_(m.weight, std=0.02)
        elif isinstance(m, nn.Conv1d):
            nn.init.normal_(m.weight, std=0.02)
            if m.bias is not None:
                nn.init.zeros_(m.bias)

    # -- caches ------------------------------------------------------------ #

    def _rope(self, t: int, device, dtype):
        key = (t, str(device), dtype)
        if key not in self._rope_cache:
            self._rope_cache[key] = build_rope_cache(
                t, self.cfg.head_dim, self.cfg.rope_base, device, dtype
            )
        return self._rope_cache[key]

    def _window_mask(self, t: int, device) -> torch.Tensor:
        key = (t, str(device))
        if key not in self._window_cache:
            idx = torch.arange(t, device=device)
            self._window_cache[key] = (idx[:, None] - idx[None, :]).abs() <= self.cfg.local_window
        return self._window_cache[key]

    # -- 子步驟 ------------------------------------------------------------- #

    def encode_text(self, input_ids, feat, pad_mask) -> torch.Tensor:
        """[B,T] → [B,T,D]。"""
        t, device = input_ids.shape[1], input_ids.device
        x = self.text_feat(self.char_emb(input_ids), feat)
        if self.cfg.encoder == "local":
            ar = torch.arange(t, device=device)
            eye = ar[:, None] == ar[None, :]  # 保留對角線，避免整列被遮而產生 NaN
            mask = ((self._window_mask(t, device) & pad_mask[:, None, :]) | eye).unsqueeze(1)
            rope = self._rope(t, device, x.dtype)
            for layer in self.layers:
                x = layer(x, mask, rope)
        else:
            for layer in self.layers:
                x = layer(x, pad_mask)
        return self.enc_ln(x)

    @staticmethod
    def span_pool(h: torch.Tensor, spans: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
        """對每個 proposal 的 [start,end) 做 mean-pool。h:[B,T,D] spans:[B,P,2] → [B,P,D]。"""
        t = h.shape[1]
        pos = torch.arange(t, device=h.device)
        start = spans[..., 0].clamp(0, t).unsqueeze(-1)
        end = spans[..., 1].clamp(0, t).unsqueeze(-1)
        w = ((pos >= start) & (pos < end) & valid).to(h.dtype)  # [B,P,T]
        return torch.matmul(w, h) / w.sum(-1, keepdim=True).clamp(min=1.0)

    def encode_cands(self, cand_tok: torch.Tensor) -> torch.Tensor:
        """候選表面字串 → 向量。[B,P,C,L] → [B,P,C,D]。
        走的是與文本同一份 char_emb，且只有 mean-pool + proj：
        任何未見過的字串都能得到有限且有梯度的表徵（zero-shot 的前提）。
        """
        if self.cfg.candidate_encoder == "id":
            # 把字元序列雜湊成單一 id：同字串 → 同 id，不同字串 → 不同 id。
            # 語意上等同 per-candidate 查表，且不需要重新產生資料。
            mult = torch.tensor([1, 131, 131 ** 2, 131 ** 3, 131 ** 4, 131 ** 5],
                                device=cand_tok.device, dtype=torch.long)
            mult = mult[: cand_tok.shape[-1]]
            cid = (cand_tok * mult).sum(-1) % self.cfg.cand_id_vocab
            return self.cand_id_emb(cid)
        valid = (cand_tok > 0).unsqueeze(-1)  # id 0 = 右側 padding
        e = self.char_emb(cand_tok)
        e = e * valid.to(e.dtype)
        pooled = e.sum(-2) / valid.sum(-2).clamp(min=1).to(e.dtype)
        return self.cand_proj(self.cand_ln(pooled))

    # -- forward ------------------------------------------------------------ #

    def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
        """batch：
            input_ids  [B,T] int64      原文字元 id
            feat       [B,T,4] int64    script / 在 proposal span 內 / 保護段 / 詞界
            pad_mask   [B,T] bool       （可省，缺省全 True）
            prop_spans [B,P,2] int64    proposal 的 [start,end)
            prop_mask  [B,P] bool
            cand_tok   [B,P,C,L] int64  候選表面字串（右側 0 padding）
            cand_mask  [B,P,C] bool
            cand_feat  [B,P,C,K] float
        回傳 {"cand_logits": [B,P,C]}，padding 候選為 -inf。
        """
        input_ids = batch["input_ids"]
        b, t = input_ids.shape
        device = input_ids.device

        pad_mask = batch.get("pad_mask")
        pad_mask = (
            torch.ones(b, t, dtype=torch.bool, device=device)
            if pad_mask is None
            else pad_mask.bool()
        )
        prop_mask = batch["prop_mask"].bool()
        cand_mask = batch["cand_mask"].bool()

        h_text = self.encode_text(input_ids, batch["feat"], pad_mask)
        h = self.span_pool(  # [B,P,D]
            h_text, batch["prop_spans"], (pad_mask[:, None, :] & prop_mask[..., None])
        )
        e = self.encode_cands(batch["cand_tok"])  # [B,P,C,D]

        c = e.shape[2]
        h_exp = h.unsqueeze(2).expand(-1, -1, c, -1)
        z = torch.cat([h_exp, e, h_exp * e, batch["cand_feat"].to(e.dtype)], dim=-1)
        logits = self.score(z).squeeze(-1)  # [B,P,C]
        return {"cand_logits": logits.masked_fill(~cand_mask, float("-inf"))}

    def compute_loss(self, outputs, batch):
        return compute_loss(outputs, batch)


# --------------------------------------------------------------------------- #
# Loss：L = L_cand + 0.3 · L_keep_bias
# --------------------------------------------------------------------------- #


def compute_loss(
    outputs: dict[str, torch.Tensor],
    batch: dict[str, torch.Tensor],
    weights: dict[str, float] | None = None,
) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
    """masked cross-entropy over candidates，外加 keep_bias 正則。
    keep_bias：對 gold == 0（維持原樣）的 proposal，任何非 0 候選只要分數不比
    候選 0 低 0.5 以上就受罰 —— 直接壓「改了不該改」（佔 V1 錯誤的 73%）。
    """
    w = {"cand": W_CAND, "keep_bias": W_KEEP_BIAS}
    if weights:
        w.update(weights)

    logits = outputs["cand_logits"]
    prop_mask = batch["prop_mask"].bool()
    cand_mask = batch["cand_mask"].bool()
    n_cands = logits.shape[-1]
    dtype = logits.dtype

    # padding 候選的 -inf 會污染算術：CE 用 finfo.min，hinge 用 0 並顯式遮罩
    neg = torch.finfo(dtype).min
    safe = torch.where(cand_mask, logits, torch.full_like(logits, neg))
    finite = torch.where(cand_mask, logits, torch.zeros_like(logits))

    gold = batch["gold_cand"].clamp(0, n_cands - 1)
    prop_w = prop_mask.to(dtype)
    n_prop = prop_w.sum().clamp(min=1.0)

    # L_cand：每個 proposal 對候選集的 cross-entropy，padding proposal 不計
    logp = torch.log_softmax(safe, dim=-1)
    nll = -logp.gather(-1, gold.unsqueeze(-1)).squeeze(-1)
    l_cand = (nll * prop_w).sum() / n_prop

    # L_keep_bias：gold 為「維持原樣」時，對每個非 0 候選各罰一次 hinge
    is_keep = (gold == KEEP_INDEX) & prop_mask
    s_keep = finite[..., KEEP_INDEX : KEEP_INDEX + 1]
    other = cand_mask & (
        torch.arange(n_cands, device=logits.device)[None, None, :] != KEEP_INDEX
    )
    hinge = F.relu(finite - s_keep + KEEP_MARGIN) * other.to(dtype)
    n_keep = is_keep.to(dtype).sum().clamp(min=1.0)
    l_keep = (hinge.sum(-1) * is_keep.to(dtype)).sum() / n_keep

    total = w["cand"] * l_cand + w["keep_bias"] * l_keep
    with torch.no_grad():
        acc = ((safe.argmax(-1) == gold).to(dtype) * prop_w).sum() / n_prop
    return total, {
        "loss": total.detach(),
        "cand": l_cand.detach(),
        "keep_bias": l_keep.detach(),
        "acc": acc,
        "n_prop": n_prop.detach(),
    }

	"""TWLAT-R（V2）模型（《04 實驗設計》§4）。

	任務：對文本中每個 proposal（字典判定「這裡可能要改」）在候選集中選一個。
	候選 index 0 永遠是「維持原樣」，即文本中實際出現的形式。

	核心約束（V2 的全部重點）：
	禁止任何 per-candidate / per-site 的 trainable embedding lookup。
	候選只能由「表面字串 + 數值特徵」動態編碼，且與文本共用同一份 char embedding。
	因此新增字典條目不需要新增任何參數，held-out proposal 也不是隨機向量。
	（V1 的 `cand_emb = nn.Embedding(4096, 192)` 佔 13% 參數並阻斷 zero-shot，就是要修掉的。）

	結構：

	char_emb（共用）──┬─→ context encoder ──→ span mean-pool ──→ h_i [B,P,D]
	│
	└─→ 候選字串 mean-pool ──→ proj ──────────→ e_c [B,P,C,D]

	score = MLP([h_i ⊕ e_c ⊕ (h_i * e_c) ⊕ cand_feat]) → [B,P,C]

	context encoder 可切換（`TWLATRConfig.encoder`），兩者參數量刻意對齊以便做 scaling curve：
	- "tcn" ：4 層 dilated depthwise separable Conv1d（dilation 1/2/4/8、kernel 5）
	- "local" ：4 層 local-window Transformer（window 半徑 32、RoPE、pre-LN、GELU）

	batch 欄位見 `TWLATR.forward` docstring。
	"""

	from __future__ import annotations

	from dataclasses import dataclass
	from typing import Any

	import torch
	import torch.nn as nn
	import torch.nn.functional as F

	# Loss 權重（《04》§0 的錯誤分析：73% 的錯是「改了不該改」，故 keep_bias 直接壓它）
	W_CAND = 1.0
	W_KEEP_BIAS = 0.3
	KEEP_MARGIN = 0.5
	KEEP_INDEX = 0 # 候選 0 恆為「維持原樣」

	ENCODERS = ("tcn", "local")


	@dataclass
	class TWLATRConfig:
	"""TWLAT-R 配置。目標參數量 3.0M–4.5M。

	註：d_model=160 只有約 2.5M（低於下限），故預設起跳為 192；
	tools/param_count_r.py 會印出兩者的對照。
	"""

	d_model: int = 192
	n_heads: int = 4
	ffn_dim: int = 768
	dropout: float = 0.1

	# context encoder：可切換，兩種都必須可跑
	encoder: str = "local"
	n_layers: int = 4
	local_window: int = 32 # 半徑；\|i-j\| <= 32 才可見
	conv_kernel: int = 5
	conv_dilations: tuple[int, ...] = (1, 2, 4, 8)
	# depthwise 之後的 pointwise 通道擴張倍率；設 1 即教科書式 depthwise separable，
	# 設 2 可讓 tcn 與 local 的 mixer 參數量幾乎相等（scaling curve 才公平）
	conv_expansion: int = 2

	# char embedding（文本與候選共用；無任何 per-ID 候選表）
	han_vocab: int = 4000 # 常用字直接查表；id >= han_vocab 走 hash
	n_hash: int = 2
	hash_buckets: int = 2048
	feat_dim: int = 4 # script / 是否在 proposal span 內 / 是否保護段 / 詞界
	feat_vocab: tuple[int, ...] = (16, 4, 4, 4)

	# proposal / candidate 形狀
	max_props: int = 48 # P
	max_cands: int = 8 # C
	cand_len: int = 6 # L，候選表面字串的最大字元數
	cand_feat_dim: int = 12 # K
	score_hidden: int = 512
	# 候選表徵方式：
	# "dynamic" —— 由表面字串經共用 char_emb 組合而成（V2 預設，可 zero-shot）
	# "id" —— per-candidate learned lookup（H-B 的對照組，複製 V1 的失敗模式）
	candidate_encoder: str = "dynamic"
	cand_id_vocab: int = 4096

	seq_len: int = 512
	rope_base: float = 10000.0

	def __post_init__(self) -> None:
	assert self.encoder in ENCODERS, f"encoder 必須是 {ENCODERS}"
	assert self.d_model % self.n_heads == 0
	assert self.d_model % 2 == 0, "hash_dim = d_model // 2，需為偶數"
	assert len(self.feat_vocab) == self.feat_dim
	assert self.conv_kernel % 2 == 1, "kernel 需為奇數才能等長 padding"

	@property
	def head_dim(self) -> int:
	return self.d_model // self.n_heads

	@property
	def hash_dim(self) -> int:
	return self.d_model // 2

	def dilation_at(self, i: int) -> int:
	return self.conv_dilations[i % len(self.conv_dilations)]


	# --------------------------------------------------------------------------- #
	# RoPE
	# --------------------------------------------------------------------------- #


	def build_rope_cache(
	seq_len: int, head_dim: int, base: float, device, dtype
	) -> tuple[torch.Tensor, torch.Tensor]:
	"""回傳 [T, head_dim//2] 的 cos / sin。"""
	half = head_dim // 2
	inv_freq = base ** (-torch.arange(half, device=device, dtype=torch.float32) / half)
	pos = torch.arange(seq_len, device=device, dtype=torch.float32)
	freqs = torch.outer(pos, inv_freq)
	return freqs.cos().to(dtype), freqs.sin().to(dtype)


	def apply_rope(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
	"""x: [B, H, T, D]，對相鄰兩維做旋轉。"""
	x_even, x_odd = x[..., 0::2], x[..., 1::2]
	cos = cos[None, None, : x.shape[-2], :]
	sin = sin[None, None, : x.shape[-2], :]
	out = torch.stack([x_even * cos - x_odd * sin, x_even * sin + x_odd * cos], dim=-1)
	return out.flatten(-2)


	# --------------------------------------------------------------------------- #
	# Embedding：文本與候選共用
	# --------------------------------------------------------------------------- #

	_HASH_MULT = (2654435761, 40503)
	_HASH_ADD = (0, 987654321)


	class CharEmbedding(nn.Module):
	"""常用字 4000 直接查表；id >= han_vocab 的罕字用 2 組 hash（buckets 2048, dim d/2）
	串接後投影。

	文本與候選字串共用這一份：候選只是一串字元 id，沒有自己的 embedding 表，
	所以字典新增條目不會增加任何參數，未見過的字串也落在同一個表徵空間。
	"""

	def __init__(self, cfg: TWLATRConfig):
	super().__init__()
	self.cfg = cfg
	self.han = nn.Embedding(cfg.han_vocab, cfg.d_model)
	self.hash = nn.ModuleList(
	nn.Embedding(cfg.hash_buckets, cfg.hash_dim) for _ in range(cfg.n_hash)
	)
	self.hash_proj = nn.Linear(cfg.n_hash * cfg.hash_dim, cfg.d_model)

	def forward(self, ids: torch.Tensor) -> torch.Tensor:
	"""ids: 任意形狀 [...]，回傳 [..., d_model]。"""
	cfg = self.cfg
	ids = ids.clamp(min=0)
	rare = ids >= cfg.han_vocab
	han = self.han(ids.clamp(max=cfg.han_vocab - 1))
	parts = []
	for i, emb in enumerate(self.hash):
	m = _HASH_MULT[i % len(_HASH_MULT)]
	a = _HASH_ADD[i % len(_HASH_ADD)]
	parts.append(emb((ids * m + a) % cfg.hash_buckets))
	rare_vec = self.hash_proj(torch.cat(parts, dim=-1))
	return torch.where(rare.unsqueeze(-1), rare_vec, han)


	class TextFeatures(nn.Module):
	"""文本側的 4 個離散特徵；char embedding 由外部傳入，以免共用的表被重複註冊。"""

	def __init__(self, cfg: TWLATRConfig):
	super().__init__()
	self.cfg = cfg
	self.feat = nn.ModuleList(nn.Embedding(n, cfg.d_model) for n in cfg.feat_vocab)
	self.ln = nn.LayerNorm(cfg.d_model)
	self.drop = nn.Dropout(cfg.dropout)

	def forward(self, x: torch.Tensor, feat: torch.Tensor) -> torch.Tensor:
	for i, emb in enumerate(self.feat):
	x = x + emb(feat[..., i].clamp(0, self.cfg.feat_vocab[i] - 1))
	return self.drop(self.ln(x))


	# --------------------------------------------------------------------------- #
	# Context encoder（可切換）
	# --------------------------------------------------------------------------- #


	class LocalAttention(nn.Module):
	"""local-window self-attention + RoPE；\|i-j\| <= local_window 才可見。"""

	def __init__(self, cfg: TWLATRConfig):
	super().__init__()
	self.cfg = cfg
	self.qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model)
	self.out = nn.Linear(cfg.d_model, cfg.d_model)
	self.drop = nn.Dropout(cfg.dropout)

	def forward(self, x, attn_mask, rope):
	b, t, _ = x.shape
	h, d = self.cfg.n_heads, self.cfg.head_dim
	q, k, v = self.qkv(x).view(b, t, 3, h, d).permute(2, 0, 3, 1, 4).unbind(0)
	q, k = apply_rope(q, rope), apply_rope(k, rope)
	p = self.cfg.dropout if self.training else 0.0
	y = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=p)
	y = y.transpose(1, 2).reshape(b, t, self.cfg.d_model)
	return self.drop(self.out(y))


	class ConvMixer(nn.Module):
	"""dilated depthwise separable Conv1d：depthwise(k, dilation) → pointwise 擴張 → GELU
	→ pointwise 還原。padding 位置先歸零，避免 pad 洩漏進感受野。"""

	def __init__(self, cfg: TWLATRConfig, dilation: int):
	super().__init__()
	d, k = cfg.d_model, cfg.conv_kernel
	pad = dilation * (k - 1) // 2 # 等長輸出
	mid = d * cfg.conv_expansion
	self.dw = nn.Conv1d(d, d, k, padding=pad, dilation=dilation, groups=d)
	self.pw1 = nn.Conv1d(d, mid, 1)
	self.pw2 = nn.Conv1d(mid, d, 1)
	self.drop = nn.Dropout(cfg.dropout)

	def forward(self, x, pad_mask, rope=None):
	z = (x * pad_mask.unsqueeze(-1).to(x.dtype)).transpose(1, 2)
	z = self.pw2(F.gelu(self.pw1(self.dw(z))))
	return self.drop(z.transpose(1, 2))


	class EncoderBlock(nn.Module):
	"""pre-LN：mixer（local attention 或 dilated conv）+ FFN。"""

	def __init__(self, cfg: TWLATRConfig, layer_idx: int):
	super().__init__()
	self.ln1 = nn.LayerNorm(cfg.d_model)
	if cfg.encoder == "local":
	self.mixer: nn.Module = LocalAttention(cfg)
	else:
	self.mixer = ConvMixer(cfg, cfg.dilation_at(layer_idx))
	self.ln2 = nn.LayerNorm(cfg.d_model)
	self.ffn = nn.Sequential(
	nn.Linear(cfg.d_model, cfg.ffn_dim),
	nn.GELU(),
	nn.Dropout(cfg.dropout),
	nn.Linear(cfg.ffn_dim, cfg.d_model),
	nn.Dropout(cfg.dropout),
	)

	def forward(self, x, mix_arg, rope=None):
	x = x + self.mixer(self.ln1(x), mix_arg, rope)
	x = x + self.ffn(self.ln2(x))
	return x


	# --------------------------------------------------------------------------- #
	# TWLAT-R
	# --------------------------------------------------------------------------- #


	class TWLATR(nn.Module):
	def __init__(self, cfg: TWLATRConfig \| None = None):
	super().__init__()
	self.cfg = cfg = cfg or TWLATRConfig()

	self.char_emb = CharEmbedding(cfg)
	self.text_feat = TextFeatures(cfg)

	self.layers = nn.ModuleList(
	EncoderBlock(cfg, i) for i in range(cfg.n_layers)
	)
	self.enc_ln = nn.LayerNorm(cfg.d_model)

	# 候選側：mean-pool 後只有一個 LN + 一個線性投影，沒有任何 per-ID 參數
	self.cand_ln = nn.LayerNorm(cfg.d_model)
	self.cand_proj = nn.Linear(cfg.d_model, cfg.d_model)
	# H-B 對照組：per-candidate learned lookup。未見過的候選只會取到
	# 一列未訓練的隨機向量——這正是 V1 `cand_emb = nn.Embedding(4096, 192)`
	# 的行為，用來檢驗動態編碼是否真的帶來 held-out 優勢。
	if cfg.candidate_encoder == "id":
	self.cand_id_emb = nn.Embedding(cfg.cand_id_vocab, cfg.d_model)

	score_in = 3 * cfg.d_model + cfg.cand_feat_dim
	self.score = nn.Sequential(
	nn.Linear(score_in, cfg.score_hidden),
	nn.GELU(),
	nn.Dropout(cfg.dropout),
	nn.Linear(cfg.score_hidden, 1),
	)

	self.apply(self._init_weights)
	self._rope_cache: dict[Any, tuple[torch.Tensor, torch.Tensor]] = {}
	self._window_cache: dict[Any, torch.Tensor] = {}

	@staticmethod
	def _init_weights(m: nn.Module) -> None:
	if isinstance(m, nn.Linear):
	nn.init.normal_(m.weight, std=0.02)
	if m.bias is not None:
	nn.init.zeros_(m.bias)
	elif isinstance(m, nn.Embedding):
	nn.init.normal_(m.weight, std=0.02)
	elif isinstance(m, nn.Conv1d):
	nn.init.normal_(m.weight, std=0.02)
	if m.bias is not None:
	nn.init.zeros_(m.bias)

	# -- caches ------------------------------------------------------------ #

	def _rope(self, t: int, device, dtype):
	key = (t, str(device), dtype)
	if key not in self._rope_cache:
	self._rope_cache[key] = build_rope_cache(
	t, self.cfg.head_dim, self.cfg.rope_base, device, dtype
	)
	return self._rope_cache[key]

	def _window_mask(self, t: int, device) -> torch.Tensor:
	key = (t, str(device))
	if key not in self._window_cache:
	idx = torch.arange(t, device=device)
	self._window_cache[key] = (idx[:, None] - idx[None, :]).abs() <= self.cfg.local_window
	return self._window_cache[key]

	# -- 子步驟 ------------------------------------------------------------- #

	def encode_text(self, input_ids, feat, pad_mask) -> torch.Tensor:
	"""[B,T] → [B,T,D]。"""
	t, device = input_ids.shape[1], input_ids.device
	x = self.text_feat(self.char_emb(input_ids), feat)
	if self.cfg.encoder == "local":
	ar = torch.arange(t, device=device)
	eye = ar[:, None] == ar[None, :] # 保留對角線，避免整列被遮而產生 NaN
	mask = ((self._window_mask(t, device) & pad_mask[:, None, :]) \| eye).unsqueeze(1)
	rope = self._rope(t, device, x.dtype)
	for layer in self.layers:
	x = layer(x, mask, rope)
	else:
	for layer in self.layers:
	x = layer(x, pad_mask)
	return self.enc_ln(x)

	@staticmethod
	def span_pool(h: torch.Tensor, spans: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:
	"""對每個 proposal 的 [start,end) 做 mean-pool。h:[B,T,D] spans:[B,P,2] → [B,P,D]。"""
	t = h.shape[1]
	pos = torch.arange(t, device=h.device)
	start = spans[..., 0].clamp(0, t).unsqueeze(-1)
	end = spans[..., 1].clamp(0, t).unsqueeze(-1)
	w = ((pos >= start) & (pos < end) & valid).to(h.dtype) # [B,P,T]
	return torch.matmul(w, h) / w.sum(-1, keepdim=True).clamp(min=1.0)

	def encode_cands(self, cand_tok: torch.Tensor) -> torch.Tensor:
	"""候選表面字串 → 向量。[B,P,C,L] → [B,P,C,D]。

	走的是與文本同一份 char_emb，且只有 mean-pool + proj：
	任何未見過的字串都能得到有限且有梯度的表徵（zero-shot 的前提）。
	"""
	if self.cfg.candidate_encoder == "id":
	# 把字元序列雜湊成單一 id：同字串 → 同 id，不同字串 → 不同 id。
	# 語意上等同 per-candidate 查表，且不需要重新產生資料。
	mult = torch.tensor([1, 131, 131 2, 131 3, 131 4, 131 5],
	device=cand_tok.device, dtype=torch.long)
	mult = mult[: cand_tok.shape[-1]]
	cid = (cand_tok * mult).sum(-1) % self.cfg.cand_id_vocab
	return self.cand_id_emb(cid)
	valid = (cand_tok > 0).unsqueeze(-1) # id 0 = 右側 padding
	e = self.char_emb(cand_tok)
	e = e * valid.to(e.dtype)
	pooled = e.sum(-2) / valid.sum(-2).clamp(min=1).to(e.dtype)
	return self.cand_proj(self.cand_ln(pooled))

	# -- forward ------------------------------------------------------------ #

	def forward(self, batch: dict[str, torch.Tensor]) -> dict[str, torch.Tensor]:
	"""batch：
	input_ids [B,T] int64 原文字元 id
	feat [B,T,4] int64 script / 在 proposal span 內 / 保護段 / 詞界
	pad_mask [B,T] bool （可省，缺省全 True）
	prop_spans [B,P,2] int64 proposal 的 [start,end)
	prop_mask [B,P] bool
	cand_tok [B,P,C,L] int64 候選表面字串（右側 0 padding）
	cand_mask [B,P,C] bool
	cand_feat [B,P,C,K] float
	回傳 {"cand_logits": [B,P,C]}，padding 候選為 -inf。
	"""
	input_ids = batch["input_ids"]
	b, t = input_ids.shape
	device = input_ids.device

	pad_mask = batch.get("pad_mask")
	pad_mask = (
	torch.ones(b, t, dtype=torch.bool, device=device)
	if pad_mask is None
	else pad_mask.bool()
	)
	prop_mask = batch["prop_mask"].bool()
	cand_mask = batch["cand_mask"].bool()

	h_text = self.encode_text(input_ids, batch["feat"], pad_mask)
	h = self.span_pool( # [B,P,D]
	h_text, batch["prop_spans"], (pad_mask[:, None, :] & prop_mask[..., None])
	)
	e = self.encode_cands(batch["cand_tok"]) # [B,P,C,D]

	c = e.shape[2]
	h_exp = h.unsqueeze(2).expand(-1, -1, c, -1)
	z = torch.cat([h_exp, e, h_exp * e, batch["cand_feat"].to(e.dtype)], dim=-1)
	logits = self.score(z).squeeze(-1) # [B,P,C]
	return {"cand_logits": logits.masked_fill(~cand_mask, float("-inf"))}

	def compute_loss(self, outputs, batch):
	return compute_loss(outputs, batch)


	# --------------------------------------------------------------------------- #
	# Loss：L = L_cand + 0.3 · L_keep_bias
	# --------------------------------------------------------------------------- #


	def compute_loss(
	outputs: dict[str, torch.Tensor],
	batch: dict[str, torch.Tensor],
	weights: dict[str, float] \| None = None,
	) -> tuple[torch.Tensor, dict[str, torch.Tensor]]:
	"""masked cross-entropy over candidates，外加 keep_bias 正則。

	keep_bias：對 gold == 0（維持原樣）的 proposal，任何非 0 候選只要分數不比
	候選 0 低 0.5 以上就受罰 —— 直接壓「改了不該改」（佔 V1 錯誤的 73%）。
	"""
	w = {"cand": W_CAND, "keep_bias": W_KEEP_BIAS}
	if weights:
	w.update(weights)

	logits = outputs["cand_logits"]
	prop_mask = batch["prop_mask"].bool()
	cand_mask = batch["cand_mask"].bool()
	n_cands = logits.shape[-1]
	dtype = logits.dtype

	# padding 候選的 -inf 會污染算術：CE 用 finfo.min，hinge 用 0 並顯式遮罩
	neg = torch.finfo(dtype).min
	safe = torch.where(cand_mask, logits, torch.full_like(logits, neg))
	finite = torch.where(cand_mask, logits, torch.zeros_like(logits))

	gold = batch["gold_cand"].clamp(0, n_cands - 1)
	prop_w = prop_mask.to(dtype)
	n_prop = prop_w.sum().clamp(min=1.0)

	# L_cand：每個 proposal 對候選集的 cross-entropy，padding proposal 不計
	logp = torch.log_softmax(safe, dim=-1)
	nll = -logp.gather(-1, gold.unsqueeze(-1)).squeeze(-1)
	l_cand = (nll * prop_w).sum() / n_prop

	# L_keep_bias：gold 為「維持原樣」時，對每個非 0 候選各罰一次 hinge
	is_keep = (gold == KEEP_INDEX) & prop_mask
	s_keep = finite[..., KEEP_INDEX : KEEP_INDEX + 1]
	other = cand_mask & (
	torch.arange(n_cands, device=logits.device)[None, None, :] != KEEP_INDEX
	)
	hinge = F.relu(finite - s_keep + KEEP_MARGIN) * other.to(dtype)
	n_keep = is_keep.to(dtype).sum().clamp(min=1.0)
	l_keep = (hinge.sum(-1) * is_keep.to(dtype)).sum() / n_keep

	total = w["cand"] * l_cand + w["keep_bias"] * l_keep
	with torch.no_grad():
	acc = ((safe.argmax(-1) == gold).to(dtype) * prop_w).sum() / n_prop
	return total, {
	"loss": total.detach(),
	"cand": l_cand.detach(),
	"keep_bias": l_keep.detach(),
	"acc": acc,
	"n_prop": n_prop.detach(),
	}