scripts/46_fsq_code_transfer.py

"""真正有判别力的那个测试：码的「含义」能不能跨序列迁移？

前两个指标都被否掉了：
  · within-code 方差 / 重建 err —— 循环。FSQ 是确定性映射、AE 照重构误差训的，
    within<1 是构造保证的，不是发现。
  · 覆盖率 —— 没有判别力。词表远小于训练位置数时，哈希也会给出 90% 利用率和
    95% 覆盖；词表远大于时样本不够。44_/45_ 两次都栽在这上面。

能区分的只有一个：**同一个码在没见过的序列里，是否指向同一种局部结构。**
靶子用局部 k-mer —— 它直接来自输入序列，不是 encoder 学出来的，所以不循环。

判据：
  在训练序列上学一张 code -> 该码最常见的局部 6-mer 的表，
  然后在**没见过的序列**上考：
      用码去预测该位置的 6-mer 的准确率   vs   把码随机打乱后的对照
  打乱对照保留码的边缘分布，抹掉"码<->位置"的对应关系。
  若准确率 >> 打乱对照，说明码承载了可跨序列迁移的局部结构 información，是词表。
  若 ≈ 对照，说明码到了新序列上就没有意义了 —— 接进 LLM 学不到东西。
"""
from __future__ import annotations

import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[1]
sys.path.insert(0, str(ROOT / "src"))

import torch  # noqa: E402
import torch.nn as nn  # noqa: E402
import torch.nn.functional as F  # noqa: E402

DEV = "cuda" if torch.cuda.is_available() else "cpu"
LATENT, LEVELS = 12, 2                    # vocab 4096
VOCAB = LEVELS ** LATENT
K = 6
STEPS = 1500
DATA = ROOT / "data/processed/rnaqa_n6_merged.jsonl"
N_SEQ, N_TEST = 256, 64


def fsq(z, levels):
    t = (torch.tanh(z) + 1) / 2
    s = t * (levels - 1)
    k = torch.round(s)
    return (k / (levels - 1) * 2 - 1) + (s - s.detach()), k.long()


class AE(nn.Module):
    def __init__(self, dim, latent, levels, hidden=256):
        super().__init__()
        self.enc = nn.Linear(dim, latent)
        self.levels = levels
        self.dec = nn.Sequential(nn.Linear(latent, hidden), nn.GELU(),
                                 nn.Linear(hidden, hidden), nn.GELU(),
                                 nn.Linear(hidden, dim))

    def forward(self, x):
        q, k = fsq(self.enc(x), self.levels)
        c = torch.zeros(x.shape[0], dtype=torch.long, device=x.device)
        for d in range(k.shape[1]):
            c = c * self.levels + k[:, d]
        return self.dec(q), c


def main() -> int:
    recs = []
    with open(DATA, encoding="utf-8") as f:
        for line in f:
            if line.strip():
                recs.append(json.loads(line))
    recs = [r for r in recs if r.get("sequence")]
    recs.sort(key=lambda r: len(r["sequence"]))
    n_all = 256
    seqs = [recs[int(i * (len(recs) - 1) / (n_all - 1))]["sequence"] for i in range(n_all)]

    d = torch.load(ROOT / "data/processed/_fsq_probe_x.pt", map_location="cpu",
                   weights_only=False)
    x, sid = d["x"].to(DEV), d["sid"].to(DEV)

    # 逐条比对：encoder 每条序列给出 L+2 个位置（[CLS] + 核苷酸 + [SEP]）。
    # 这与观测一致：115886 - 115374 = 512 = 2 x 256。
    cnt = torch.bincount(sid).tolist()
    keep = [i for i in range(len(seqs)) if cnt[i] == len(seqs[i]) + 2]
    drop = len(seqs) - len(keep)
    if drop:
        print(f"[filter] 丢弃 {drop} 条长度不匹配的序列，保留 {len(keep)}")
    remap = {o: n for n, o in enumerate(keep)}
    keep_pos = torch.tensor([s in remap for s in sid.tolist()], device=DEV)
    x = x[keep_pos]
    sid = torch.tensor([remap[s] for s in sid[keep_pos].tolist()], device=DEV)
    seqs = [seqs[i] for i in keep]
    N_SEQ = len(seqs)
    N, D = x.shape

    g = torch.Generator().manual_seed(7)
    test_ids = set(torch.randperm(N_SEQ, generator=g)[:max(1, N_SEQ // 4)].tolist())
    seq_tr = torch.tensor([s not in test_ids for s in sid.tolist()], device=DEV)

    torch.manual_seed(0)
    ae = AE(D, LATENT, LEVELS).to(DEV)
    opt = torch.optim.AdamW(ae.parameters(), lr=3e-3)
    idx = torch.nonzero(seq_tr).squeeze(1)
    for _ in range(STEPS):
        p = idx[torch.randint(0, idx.numel(), (512,), device=DEV)]
        F.mse_loss(ae(x[p])[0], x[p]).backward()
        opt.step(); opt.zero_grad()
    ae.eval()
    with torch.no_grad():
        _, code = ae(x)

    # 构造每个位置的局部 6-mer 编号（[CLS]/[SEP] 和末尾不足 K 的位置记为 -1）
    mer_of_pos = torch.full((N,), -1, dtype=torch.long)
    pos = 0
    for s in seqs:
        b = s.encode()
        L = len(b)
        pos += 1                                   # [CLS]
        for j in range(L):
            if j + K <= L:
                v = 0
                for ch in b[j:j + K]:
                    v = v * 4 + ({65: 0, 67: 1, 71: 2, 84: 3, 85: 0}.get(ch, 0))
                mer_of_pos[pos] = v
            pos += 1
        pos += 1                                   # [SEP]
    # 位置必须与 encoder 的 mask 逐条对齐，否则 k-mer 靶子是错位的
    assert pos == N, f"位置对齐失败: 构造 {pos} != encoder 位置 {N}"
    mer_of_pos = mer_of_pos.to(DEV)

    tr_pos = torch.nonzero(seq_tr & (mer_of_pos >= 0)).squeeze(1)
    te_pos = torch.nonzero(~seq_tr & (mer_of_pos >= 0)).squeeze(1)
    print(f"[data] 训练位置 {tr_pos.numel()}  测试位置 {te_pos.numel()}  "
          f"（局部 {K}-mer，词表 {4**K}）")

    tc, tm = code[tr_pos], mer_of_pos[tr_pos].to(DEV)
    ec, em = code[te_pos], mer_of_pos[te_pos].to(DEV)

    # code -> 该码在训练序列里最常见的 6-mer
    lookup = torch.full((VOCAB,), -1, dtype=torch.long, device=DEV)
    for c in torch.unique(tc):
        sel = tm[tc == c]
        counts = torch.bincount(sel, minlength=4 ** K)
        lookup[c] = int(counts.argmax())
    hit = lookup[ec] >= 0
    acc = float((lookup[ec][hit] == em[hit]).float().mean()) if hit.any() else 0.0
    cov = float(hit.float().mean())

    # 对照：打乱 code<->位置 的对应（保留 code 边缘分布）
    g2 = torch.Generator().manual_seed(99)
    accs = []
    for _ in range(5):
        perm = torch.randperm(tc.numel(), generator=g2)
        tc2 = tc[perm]
        lk = torch.full((VOCAB,), -1, dtype=torch.long, device=DEV)
        for c in torch.unique(tc2):
            sel = tm[tc2 == c]
            lk[c] = int(torch.bincount(sel, minlength=4 ** K).argmax())
        h = lk[ec] >= 0
        accs.append(float((lk[ec][h] == em[h]).float().mean()))
    ctrl = sum(accs) / len(accs)

    # 第二个对照：全局最常见的 6-mer 常数预测
    base = float(torch.bincount(tm, minlength=4 ** K).max()) / tm.numel()

    print()
    print("=" * 78)
    print("跨序列含义迁移：用码预测局部 6-mer")
    print("=" * 78)
    print(f"  测试位置中码在训练集见过的比例: {100*cov:.1f}%")
    print(f"  按码查表预测 6-mer 准确率     : {100*acc:.2f}%")
    print(f"  打乱 code<->位置 对照          : {100*ctrl:.2f}%")
    print(f"  常数预测（训练集最常见 6-mer）: {100*base:.2f}%")
    print(f"  相对打乱对照的增益            : {100*(acc-ctrl):+.2f}pp")
    # 信息量：以比特计
    import math
    lift = acc / max(ctrl, 1e-9)
    print(f"  提升倍数                      : {lift:.2f}x")
    print()
    print("=" * 78)
    print("判定")
    print("=" * 78)
    if acc > ctrl * 1.5 and acc > base * 1.3:
        print(f"  ✅ **码携带可跨序列迁移的局部结构信息**（{100*acc:.1f}% vs 打乱 {100*ctrl:.1f}%，"
              f"{lift:.1f}x）。")
        print("     同一个码在没见过的序列里指向同一类 6-mer —— 这是词表，不是哈希。")
        print("     接进 LLM 有可学的东西。")
    elif acc > ctrl * 1.15:
        print(f"  ⚠️ **弱迁移**：{100*acc:.1f}% vs 打乱 {100*ctrl:.1f}%（{lift:.2f}x）。")
        print("     码有结构但信噪比低。可能要靠加复用压力（commitment / 熵正则）才能变干净。")
    else:
        print(f"  ❌ **码到新序列上就没有含义了**（{100*acc:.1f}% vs 打乱 {100*ctrl:.1f}%，"
              f"{lift:.2f}x）。")
        print("     码只是把每条序列切成一串各自不同的片段 —— 对 LLM 而言是新的乱码。")
        print("     要救这条路线，必须换目标函数（加复用压力），而不是换个量化方式。")

    out = ROOT / "reports" / "fsq_code_transfer.json"
    out.write_text(json.dumps(dict(vocab=VOCAB, k=K, acc=acc, ctrl=ctrl, base=base,
                                   cov=cov, lift=lift), ensure_ascii=False, indent=2),
                   encoding="utf-8")
    print(f"\n[save] {out}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())