scripts/22_zero_encoder_probe.py

"""P1 验收：**置零编码器探针**（PLAN §3 P1 的完成判据）。

判据原文：
    "置零编码器后，对齐指标明显下降"

这条是整个计划里最重要的自检 —— 它同时挡两个坑（PLAN §5 风险登记）：
  * **投影塌缩 / 编码器被绕过**：loss 降了，但置零毫无影响 → 模型压根没用上序列；
  * **QA 里复述了序列**：指标虚高 → 置零后分数不掉，就是这个症状。

做法：
    对同一批样本跑两次前向（teacher forcing），一次用真实 RNA 表征、一次把表征置零，
    比较 loss / token 准确率。真对齐的模型，置零后 loss 必须**明显**上升。

跑法：
    python scripts/22_zero_encoder_probe.py --ckpt checkpoints/stage1/projection_final.pt \
        --data data/processed/annotations_50.jsonl --limit 20
"""

from __future__ import annotations

import argparse
import json
import sys
from pathlib import Path

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

import rnaqwen  # noqa: E402,F401

import torch  # noqa: E402

from rnaqwen.data.records import read_jsonl  # noqa: E402
from rnaqwen.models.alignment import Stage1Aligner, Stage1Config  # noqa: E402
from rnaqwen.paths import default_encoder, describe_encoder  # noqa: E402


@torch.no_grad()
def evaluate(model: Stage1Aligner, samples: list[dict], device: torch.device,
             zero_encoder: bool) -> dict:
    """返回平均 loss 与答案 token 的逐 token 准确率。"""
    total_loss, total_tok, correct = 0.0, 0, 0
    orig_forward = model.encoder.forward

    if zero_encoder:
        from rnaqwen.models.encoder import EncoderOutput

        def zero_forward(seqs):  # noqa: ANN001
            out = orig_forward(seqs)
            return EncoderOutput(hidden=torch.zeros_like(out.hidden), mask=out.mask, lengths=out.lengths)

        model.encoder.forward = zero_forward  # type: ignore[assignment]

    try:
        for s in samples:
            inputs, attn, labels = model.build_batch([s["sequence"]], [s["target"]], device)
            out = model.llm(inputs_embeds=inputs, attention_mask=attn, labels=labels)
            total_loss += float(out.loss)
            shift_logits = out.logits[:, :-1]
            shift_labels = labels[:, 1:]
            mask = shift_labels != -100
            if mask.any():
                pred = shift_logits[mask].argmax(-1)
                correct += int((pred == shift_labels[mask]).sum())
                total_tok += int(mask.sum())
    finally:
        model.encoder.forward = orig_forward  # type: ignore[assignment]

    n = max(len(samples), 1)
    return {
        "loss": total_loss / n,
        "token_accuracy": correct / total_tok if total_tok else None,
        "n_samples": len(samples),
        "n_answer_tokens": total_tok,
    }


def main() -> int:
    ap = argparse.ArgumentParser()
    ap.add_argument("--ckpt", type=Path, required=True, help="Stage1 产出的 projection_*.pt")
    ap.add_argument("--data", type=Path, required=True)
    ap.add_argument("--llm", default="models/Qwen3-8B")
    ap.add_argument("--encoder", default=default_encoder(),
                    help="不填就走 E 轨默认（ERNIE-RNA）。R 轨复现请显式给 models/rnafm")
    ap.add_argument("--target-field", default="annotation", choices=["annotation", "literature"])
    ap.add_argument("--limit", type=int, default=20)
    ap.add_argument("--device", default="cuda")
    ap.add_argument("--out", type=Path, default=ROOT / "reports" / "zero_encoder_probe.json")
    # E 轨 checkpoint 要对齐压缩配置（默认 None = 读 checkpoint 里存的）
    ap.add_argument("--compress-latent", type=int, default=None,
                    help="覆盖 checkpoint 里的 compress_latent（E 轨 checkpoint 一般不用填，会自动读到）")
    ap.add_argument("--compress-method", default=None,
                    help="覆盖 checkpoint 里的 compress_method")
    # 判据阈值。
    # ⚠️ 实测：E 轨 run1 的 Δloss 从峰值 0.273 衰减到 0.047，但 ROUGE-1 反而从 0.208 涨到 0.303。
    #    说明 Δloss 是**逐 token 平均**，答案变长/变准会把它摊薄，低 Δloss ≠ 编码器没用。
    #    所以这个探针的结论要和 31_stage1_rouge_curve.py 一起读，别单看这里的 PASS/FAIL。
    ap.add_argument("--pass-delta", type=float, default=0.5)
    ap.add_argument("--weak-delta", type=float, default=0.1)
    args = ap.parse_args()

    rows = [r.as_dict() for r in read_jsonl(args.data)]
    samples = []
    for r in rows:
        t = r.get(args.target_field) or ""
        if r.get("sequence") and t.strip():
            samples.append({"sequence": r["sequence"], "target": t})
        if args.limit and len(samples) >= args.limit:
            break
    if not samples:
        print("[fail] 没有样本")
        return 1

    blob = torch.load(args.ckpt, map_location="cpu")

    # checkpoint 里存的 config 是**训练时**的完整快照，字段集合和 Stage1Config 不完全一致
    # （多出 collapse_delta 之类的旧字段，且 llm_path/encoder_path 会和命令行参数撞名）。
    # 以前直接 **blob["config"] 会 TypeError；这里只挑 dataclass 认识的字段，
    # 且命令行显式给的路径优先。
    import dataclasses

    known = {f.name for f in dataclasses.fields(Stage1Config)}
    saved = {k: v for k, v in (blob.get("config") or {}).items() if k in known}
    dropped = sorted(set(blob.get("config") or {}) - known)
    if dropped:
        print(f"[zero-probe] 忽略 checkpoint 里的未知配置字段：{dropped}")
    saved["llm_path"] = args.llm
    saved["encoder_path"] = args.encoder
    # 允许命令行强制覆盖（E 轨 checkpoint 靠这个对齐压缩配置）
    if args.compress_latent is not None:
        saved["compress_latent"] = args.compress_latent
    if args.compress_method:
        saved["compress_method"] = args.compress_method
    cfg = Stage1Config(**saved)

    model = Stage1Aligner(cfg)
    model.projection.load_state_dict(blob["projection"])

    # E 轨 checkpoint 必须把压缩器一起装回去，否则 build_batch 的 token 数对不上。
    if cfg.compress_latent:
        if "compression" not in blob:
            print(f"[fail] 配置要求 compress_latent={cfg.compress_latent}，"
                  f"但 checkpoint 里没有 compression 权重")
            return 2
        model.compression.load_state_dict(blob["compression"])
        print(f"[zero-probe] 已装回压缩器 {blob.get('compression_spec')}")

    device = torch.device(args.device)
    model.encoder.to(device)
    model.llm.to(device)
    model.projection.to(device)
    if model.compression is not None:
        model.compression.to(device)
    model.eval()

    print(f"[zero-probe] {len(samples)} 条样本；投影 {model.projection.describe()}")
    real = evaluate(model, samples, device, zero_encoder=False)
    print(f"[zero-probe] 真实编码器 : loss={real['loss']:.4f}  tok_acc={real['token_accuracy']}")
    zero = evaluate(model, samples, device, zero_encoder=True)
    print(f"[zero-probe] 置零编码器 : loss={zero['loss']:.4f}  tok_acc={zero['token_accuracy']}")

    d_loss = zero["loss"] - real["loss"]
    acc_drop = ((real["token_accuracy"] or 0) - (zero["token_accuracy"] or 0))
    if d_loss > args.pass_delta:
        verdict = "PASS"
    elif d_loss > args.weak_delta:
        verdict = "WEAK"
    else:
        verdict = "FAIL"
    print(f"\n[zero-probe] Δloss = +{d_loss:.4f}   Δtoken_acc = -{acc_drop:.4f}   -> {verdict}")
    print("            ⚠️ 判据提醒：Δloss 是逐 token 均值，答案越准越长它越低。")
    print("               最终以 scripts/31_stage1_rouge_curve.py 的 ROUGE 为准。")
    if verdict == "FAIL":
        print("           置零几乎没影响 = 模型没在用序列（投影塌缩 / 抄文本），回去查。")

    report = {"checkpoint": str(args.ckpt), "data": str(args.data), "n": len(samples),
              "real": real, "zeroed": zero, "delta_loss": d_loss,
              "delta_token_accuracy": acc_drop, "verdict": verdict,
              "compress_latent": cfg.compress_latent, "compress_method": cfg.compress_method,
              "criterion": "PLAN §3 P1：置零编码器后对齐指标必须明显下降"
                           "（Δloss 仅作参考，最终判据见 ROUGE 曲线）"}
    args.out.parent.mkdir(parents=True, exist_ok=True)
    args.out.write_text(json.dumps(report, indent=2, ensure_ascii=False), encoding="utf-8")
    print(f"[zero-probe] 报告 -> {args.out}")
    return 0 if verdict == "PASS" else 1


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