ling_probe.py

#!/usr/bin/env python3
"""
================================================================================
ling_probe.py — 词性 / 句法深度 / 情感极性 的逐槽、逐头线性探针
================================================================================

一. 测什么
----------
  POS      17 类 UPOS 词性                     线性分类, 报准确率
  DEPTH    到依存树根节点的距离                线性回归, 报 R² 与 Spearman
  VALENCE  词级情感极性 (内置词表, 正/负)      线性分类, 报准确率

二. ★ 最重要的是对照, 不是绝对数字
-----------------------------------
词性绝大部分是【词汇性】的 —— 静态词嵌入本身就编码了大半, 不需要任何上下文。
所以"探针准确率 92%"单独看毫无意义。必须对两条基线:

  基线1  多数类           猜最常见的那一类
  基线2  槽0 (词嵌入)     ★关键★ 只看 token 身份、零上下文能做到多少

  真正有信息的是【超出槽0多少】—— 那部分才是模型从上下文里算出来的。
  bank/run 这类歧义词的消歧, 只能靠上下文, 只会出现在这个增量里。

三. 本架构独有的部分
--------------------
  标准 transformer 只能问"到第几层为止能解码出词性" —— 残差流是个和,
  你只能探测累加后的总量。
  槽位栈可以逐槽、甚至逐头单独探针 ⇒ 得到一张 L×H 的图: 哪个头编码词性。
  再接上 head_test / patch_test 就能做因果确认。

四. 依赖
--------
    pip install spacy scikit-learn tiktoken
    python -m spacy download en_core_web_sm
  (探针本身用 torch 实现, 不需要 sklearn; tiktoken 用于把 token id 还原成文本)

五. 用法
--------
    python ling_probe.py --model delta_d1280_h10_l10.py \\
        --ckpt /workspace/data/d1280_h10_l10_out/latest.pt --seqs 64
    --heads   额外跑逐头探针 (L×H 网格, 慢 H 倍)
    --struct  额外跑结构探针 (Hewitt-Manning, 依存树距离)
================================================================================
"""
import argparse, importlib.util, inspect, math, sys
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F

# ── 内置情感词表 (词级极性, 仅用于 VALENCE 探针) ──
POS_W = """good great excellent wonderful amazing fantastic superb brilliant lovely
beautiful happy joy delight pleased glad love best better perfect success win winner
strong healthy safe fun enjoy enjoyable positive nice friendly kind generous helpful
smart clever wise brave hope hopeful bright clean fresh rich reward praise honor
благо respect trust peace calm comfort easy free fair true right improve improved
gain benefit advantage victory triumph celebrate charming elegant graceful warm""".split()
NEG_W = """bad terrible awful horrible worst worse poor ugly sad angry hate hatred
fear afraid scared danger dangerous risk harmful hurt pain painful suffer suffering
fail failure lose loser weak sick ill disease death dead kill killed murder war
violence violent cruel evil wrong false lie liar cheat steal crime criminal guilty
dirty filthy rotten toxic waste destroy damage broken ruin disaster tragedy crisis
problem trouble difficult hard stress anxiety depressed miserable boring dull""".split()


def load_module(path):
    spec = importlib.util.spec_from_file_location("expmod", path)
    m = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(m)
    return m


# ═══════════════════ 分词对齐 ═══════════════════
def align(ids, dec_bytes, nlp):
    """把 GPT-2 token 对齐到 spaCy token。

    做法: 逐 token 取字节 → 累积字节偏移 → 整段解码成文本 → spaCy 分析 →
    把 spaCy 的字符偏移换算成字节偏移 → 每个 GPT-2 token 归属到覆盖它起始
    字节的那个 spaCy token。

    返回 (有效的 GPT-2 下标数组, 对应的 spaCy token 列表, doc)。
    """
    parts, offs, cur = [], [], 0
    for t in ids:
        b = dec_bytes([int(t)])
        # ★ GPT-2 的 token 通常带前导空格(" quick"), 起始字节比 spaCy 的
        #   "quick" 早一位, 直接比对会全部落空。跳过前导空白再记偏移。
        lead = len(b) - len(b.lstrip())
        offs.append(cur + lead)
        cur += len(b)
        parts.append(b)
    text = b"".join(parts).decode("utf-8", errors="replace")
    if len(text.strip()) < 20:
        return None
    doc = nlp(text)
    # spaCy 字符偏移 → 字节偏移
    sp = []
    for tk in doc:
        s = len(text[:tk.idx].encode("utf-8"))
        e = s + len(tk.text.encode("utf-8"))
        sp.append((s, e, tk))
    idx, toks, j = [], [], 0
    for i, o in enumerate(offs):
        while j < len(sp) and sp[j][1] <= o:
            j += 1
        if j < len(sp) and sp[j][0] <= o < sp[j][1]:
            idx.append(i)
            toks.append(sp[j][2])
    return (np.array(idx, dtype=np.int64), toks, doc) if idx else None


def dep_depth(tk):
    """到依存树根的距离。
    ★ 不能用 `tk.head is not tk`: spaCy 每次访问 .head 都新建 Token 对象,
      `is` 恒为假, 会一路数到上限。必须比 .i (在 doc 中的下标)。"""
    d, seen = 0, 0
    while tk.head.i != tk.i and seen < 64:
        tk = tk.head
        d += 1
        seen += 1
    return d


# ═══════════════════ 线性探针 ═══════════════════
def fit_probe(Xtr, ytr, Xte, yte, n_cls, dev, epochs=300, lr=2e-2):
    """n_cls>1 分类(返回准确率); n_cls==0 回归(返回 R²)。特征按训练集标准化。

    ★ 回归走【闭式岭回归】而不是梯度下降: 梯度下降在几百轮内会明显欠拟合,
      把"表示里没有这个信息"和"探针没训练好"混为一谈 —— 探针实验里这是
      致命的, 因为结论正好取决于前者。闭式解没有这个风险。
    """
    mu, sd = Xtr.mean(0, keepdim=True), Xtr.std(0, keepdim=True).clamp_min(1e-6)
    Xtr, Xte = (Xtr - mu) / sd, (Xte - mu) / sd
    if n_cls <= 1:                                   # 闭式岭回归
        A = torch.cat([Xtr, torch.ones_like(Xtr[:, :1])], 1).double()
        b = ytr.double()
        G = A.T @ A + 1e-2 * torch.eye(A.shape[1], device=A.device, dtype=A.dtype)
        w = torch.linalg.solve(G, A.T @ b)
        Ae = torch.cat([Xte, torch.ones_like(Xte[:, :1])], 1).double()
        p = Ae @ w
        yt = yte.double()
        ss = ((yt - p) ** 2).sum()
        st = ((yt - yt.mean()) ** 2).sum().clamp_min(1e-9)
        return (1 - ss / st).item()
    W = nn.Linear(Xtr.shape[1], n_cls).to(dev)
    opt = torch.optim.AdamW(W.parameters(), lr=lr, weight_decay=1e-3)
    sch = torch.optim.lr_scheduler.CosineAnnealingLR(opt, epochs)
    for _ in range(epochs):
        opt.zero_grad()
        F.cross_entropy(W(Xtr), ytr).backward()
        opt.step()
        sch.step()
    with torch.no_grad():
        return (W(Xte).argmax(-1) == yte).float().mean().item()


@torch.no_grad()
def run_ablated(model, x, k, h):
    """把第 k 个槽(1-based)的第 h 个头(1-based)置零, 并让影响向下游传播。

    ★ 只把探针特征里那一段抹掉是【错的】—— 那只切断了"读", 没切断"算":
      下游各层仍然看得到这个头的输出, 表示照样被它塑造过。必须在前向里
      写完槽就置零, 后面的 att/ffn 才是真的没有它。
    """
    hd = model.d // model.H
    sl = slice((h - 1) * hd, h * hd)
    pos = model.__class__.__module__ and None
    from importlib import import_module
    pp = model.embed.weight.new_zeros(0)          # 占位, 下面用模块里的函数
    idxpos = model_probe_positions(x)
    xx = model.embed(x).float()
    slots, att_in = [xx], xx
    for i in range(model.L):
        out = model.atts[i](att_in, model.rope_cos, model.rope_sin,
                            model.sharp, idxpos)
        if i + 1 == k:
            out = out.clone()
            out[..., sl] = 0.0
        slots.append(out)
        att_in = model.ffns[i](slots[:i + 1], slots[i + 1])
    return att_in, slots


model_probe_positions = None      # 由 main 注入 (模块级函数, 与训练脚本一致)


@torch.no_grad()
def collect(m, model, nlp, dec_bytes, n_seq, dev, ablate=None):
    """跑数据, 收集 (每槽激活, POS标签, 深度标签, 极性标签)。"""
    sig = inspect.signature(m.ShardLoader.__init__).parameters
    kw = {"insert_p": 0.0} if "insert_p" in sig else {}
    ld = m.ShardLoader(m.DATA_DIR, "val", m.MICRO_BATCH, m.SEQ_LEN, **kw)
    ld.reset()
    upos, feats = {}, None
    Y = {"pos": [], "depth": [], "val": [], "tokid": []}
    seq_id, done = [], 0
    pw, nw = set(POS_W), set(NEG_W)
    while done < n_seq:
        x, _ = ld.next_batch(dev)
        with m.amp_ctx():
            if ablate is None:
                _, slots = model.run(x, keep_slots=True)
            else:
                _, slots = run_ablated(model, x, ablate[0], ablate[1])
        slots = [s.float() for s in slots]
        if feats is None:
            feats = [[] for _ in slots]
        for b in range(x.shape[0]):
            if done >= n_seq:
                break
            a = align(x[b].tolist(), dec_bytes, nlp)
            if a is None:
                continue
            idx, toks, _ = a
            keep, yp, yd, yv = [], [], [], []
            for i, tk in zip(idx, toks):
                if tk.pos_ not in upos:
                    upos[tk.pos_] = len(upos)
                low = tk.text.lower().strip()
                v = 1 if low in pw else (0 if low in nw else -1)
                keep.append(i)
                yp.append(upos[tk.pos_])
                yd.append(min(dep_depth(tk), 15))
                yv.append(v)
            if not keep:
                continue
            kt = torch.tensor(keep, device=dev)
            for si, s in enumerate(slots):
                feats[si].append(s[b][kt].cpu())
            Y["tokid"].append(x[b][kt].cpu().numpy())
            Y["pos"].append(np.array(yp))
            Y["depth"].append(np.array(yd))
            Y["val"].append(np.array(yv))
            seq_id.append(np.full(len(keep), done))
            done += 1
    F_ = [torch.cat(f) for f in feats]
    Yc = {k: np.concatenate(v) for k, v in Y.items()}
    return F_, Yc, np.concatenate(seq_id), upos


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--model", required=True)
    ap.add_argument("--ckpt", default=None)
    ap.add_argument("--seqs", type=int, default=64)
    ap.add_argument("--heads", action="store_true")
    ap.add_argument("--ablate", default=None,
                    help="消融一个头再跑探针, 格式 '槽,头' 如 4,2 (1-based)。"
                         "与不加此参数的结果相减 = 该头对各语言学能力的因果贡献")
    ap.add_argument("--ambig-min", type=float, default=0.15,
                    help="歧义词判定: 次高词性占比 ≥ 该值才算歧义 token")
    ap.add_argument("--epochs", type=int, default=300)
    a = ap.parse_args()

    try:
        import spacy
        nlp = spacy.load("en_core_web_sm", disable=["ner", "lemmatizer"])
    except Exception as e:
        sys.exit(f"需要 spaCy: pip install spacy && python -m spacy download en_core_web_sm\n{e}")
    try:
        import tiktoken
        enc = tiktoken.get_encoding("gpt2")
        dec_bytes = enc.decode_bytes
    except Exception as e:
        sys.exit(f"需要 tiktoken (用于把 token id 还原成文本): pip install tiktoken\n{e}")

    m = load_module(a.model)
    dev = m.device
    model = m.DeltaLM().to(dev)
    tag = "随机初始化(无结论)"
    if a.ckpt:
        ck = torch.load(a.ckpt, map_location=dev, weights_only=False)
        model.load_state_dict(ck["model"])
        tag = f"step={ck.get('step','?')}"
    model.eval()

    print("=" * 104)
    print(f"  语言学探针   {a.model}   {tag}")
    print("=" * 104)
    global model_probe_positions
    model_probe_positions = m.probe_positions
    abl = None
    if a.ablate:
        abl = tuple(int(v) for v in a.ablate.split(","))
        print(f"  ★ 消融: 槽{abl[0]} 的第 {abl[1]} 个头 (置零并向下游传播)")
        print(f"    与不加 --ablate 的结果相减, 即该头对各项能力的因果贡献")
    print(f"  收集 {a.seqs} 条序列的激活与 spaCy 标注 ...")
    Feats, Y, sid, upos = collect(m, model, nlp, dec_bytes, a.seqs, dev, abl)
    n = len(sid)
    cut = np.quantile(sid, 0.75)
    tr, te = sid <= cut, sid > cut          # 按序列切分, 防泄漏
    print(f"  样本 {n:,} 个 token  训练 {tr.sum():,} / 测试 {te.sum():,}"
          f"  (按序列切分, 无泄漏)   UPOS 类别 {len(upos)}")

    # ── 歧义 token 集合 ──
    # ★ 词性绝大部分是词汇性的: 线性探针能直接从词嵌入把 token→词性查表背下来,
    #   于是"槽0 拿 82.8%"这种数字测的是查表, 不是句法。只有【同一个 token 有
    #   多个词性】时查表才失效, 必须靠上下文。把探针限制到这些 token 上, 剩下
    #   的才是真本事。
    tid, yp_np = Y["tokid"], Y["pos"]
    amb = np.zeros(len(tid), dtype=bool)
    from collections import defaultdict
    cnt = defaultdict(lambda: defaultdict(int))
    for t, pp in zip(tid, yp_np):
        cnt[int(t)][int(pp)] += 1
    ambset = set()
    for t, dd in cnt.items():
        v = sorted(dd.values(), reverse=True)
        if len(v) > 1 and sum(v) >= 8 and v[1] / sum(v) >= a.ambig_min:
            ambset.add(t)
    for i, t in enumerate(tid):
        amb[i] = int(t) in ambset
    print(f"  歧义 token: {len(ambset)} 种 / 覆盖 {amb.sum():,} 个位置 "
          f"({amb.mean()*100:.1f}%)   [次高词性占比 ≥ {a.ambig_min}]")

    # ── 身份保留对照: 直接预测 token 自己是哪个词 (取最常见的 N 个) ──
    # ★ 一个只盯着自己位置的头会原样保留 token 身份, POS 自然高 —— 与懂不懂
    #   句法无关。有了这一列, 才能把"真句法"和"只是在传身份"分开。
    freq = defaultdict(int)
    for t in tid:
        freq[int(t)] += 1
    top = [t for t, _ in sorted(freq.items(), key=lambda kv: -kv[1])[:60]]
    tmap = {t: i for i, t in enumerate(top)}
    im = np.array([int(t) in tmap for t in tid])
    yid = np.array([tmap.get(int(t), 0) for t in tid])

    tasks = []
    yp = torch.tensor(Y["pos"], device=dev)
    tasks.append(("POS(全体)", yp, len(upos), None))
    if amb.sum() > 300:
        tasks.append(("POS(歧义)", torch.tensor(yp_np[amb], device=dev),
                      len(upos), amb))
    yd = torch.tensor(Y["depth"], device=dev)
    tasks.append(("DEPTH(R²)", yd, 0, None))
    tasks.append(("身份保留", torch.tensor(yid[im], device=dev), len(top), im))
    vm = Y["val"] >= 0
    if vm.sum() > 200:
        print("  (VALENCE 已移除: 内置词表是词汇性的, 探针只是把词表背下来,")
        print("   槽0 就能拿 98%, 测的是身份记忆不是情感。要测情感需句子级标注如 SST-2)")

    # 基线
    print(f"\n  基线")
    print("  " + "-" * 100)
    for nm, y, nc, msk in tasks:
        yy = y.cpu().numpy()
        m_ = te if msk is None else te[msk]
        if nc > 1:
            maj = np.bincount(yy[(tr if msk is None else tr[msk])]).argmax()
            print(f"    {nm:<12}多数类基线 {(yy[m_] == maj).mean()*100:.1f}%")
        else:
            print(f"    {nm:<12}常数预测基线 R² = 0.000")

    def run_row(X, label):
        cells = []
        for nm, y, nc, msk in tasks:
            Xa = X if msk is None else X[msk]
            ta = tr if msk is None else tr[msk]
            ea = te if msk is None else te[msk]
            v = fit_probe(Xa[ta].to(dev), y[ta], Xa[ea].to(dev), y[ea],
                          nc, dev, epochs=a.epochs)
            cells.append(v * 100 if nc > 1 else v)
        return cells

    names = ["槽0(词嵌入)"] + [f"槽{i}(att{i})" for i in range(1, len(Feats))]
    print(f"\n  【逐槽探针】")
    print("  " + "-" * 100)
    hdr = f"  {'槽':<16}" + "".join(f"{t[0]:<16}" for t in tasks) + "  相对槽0的增量"
    print(hdr)
    base_cells = None
    for si, X in enumerate(Feats):
        c = run_row(X, names[si])
        if si == 0:
            base_cells = c
            delta = ""
        else:
            delta = "  " + " ".join(
                f"{t[0]}{c[i]-base_cells[i]:+.1f}" if t[2] > 1
                else f"{t[0]}{c[i]-base_cells[i]:+.3f}" for i, t in enumerate(tasks))
        cells = "".join(f"{v:<16.1f}" if t[2] > 1 else f"{v:<16.3f}"
                        for v, t in zip(c, tasks))
        print(f"  {names[si]:<16}{cells}{delta}")

    print(f"\n  ★ 关键不是绝对值, 是【相对槽0的增量】。词性绝大部分是词汇性的,")
    print("    槽0(纯词嵌入, 零上下文)就能做到大半。超出槽0的那部分才是模型从")
    print("    上下文里算出来的 —— bank/run 这类歧义词的消歧只会出现在增量里。")

    if a.heads:
        H, hd = model.H, model.d // model.H
        print(f"\n  【逐头探针 · POS 准确率】(标准 transformer 做不了: 残差流是个和)")
        print("  " + "-" * 100)
        print("  " + "槽\\头".ljust(10) + "".join(f"h{h+1:<7}" for h in range(H)))
        sub = [t for t in tasks if t[0] in ("POS(歧义)", "DEPTH(R²)", "身份保留")]
        grids = {t[0]: np.zeros((len(Feats) - 1, H)) for t in sub}
        for nm, y, nc, msk in sub:
            print(f"\n  ── {nm} ──")
            print("  " + "槽\\头".ljust(10) + "".join(f"h{h+1:<7}" for h in range(H)))
            for si in range(1, len(Feats)):
                row = []
                for h in range(H):
                    Xh = Feats[si][:, h*hd:(h+1)*hd]
                    Xa = Xh if msk is None else Xh[msk]
                    ta = tr if msk is None else tr[msk]
                    ea = te if msk is None else te[msk]
                    v = fit_probe(Xa[ta].to(dev), y[ta], Xa[ea].to(dev), y[ea],
                                  nc, dev, epochs=a.epochs)
                    row.append(v * 100 if nc > 1 else v)
                grids[nm][si-1] = row
                print(f"  槽{si:<8}" + "".join(
                    (f"{v:<8.1f}" if nc > 1 else f"{v:<8.3f}") for v in row))
        if "POS(歧义)" in grids and "身份保留" in grids:
            r = grids["POS(歧义)"] / np.maximum(grids["身份保留"], 1e-6)
            fl = sorted(((r[i, j], grids["POS(歧义)"][i, j], grids["身份保留"][i, j],
                          i + 1, j + 1) for i in range(r.shape[0]) for j in range(H)),
                        reverse=True)
            print(f"\n  【句法特异性】POS(歧义)/身份保留 —— 高 = 真在做句法, 不只是传身份")
            for v, p_, d_, k, h in fl[:5]:
                print(f"    槽{k}h{h:<4} 比值 {v:.2f}  (POS歧义 {p_:.1f}% / 身份 {d_:.1f}%)")
        print("  ★ 接上 head_test/patch_test 就能做因果确认: 挖掉它探针掉多少,")
        print("    换成另一句的, 模型对当前词的判断跟不跟着翻。")
    print("=" * 104)


if __name__ == "__main__":
    main()