"""
=============================================================================
第 6 讲 · 代码 5：模型编辑——只改一个事实，别的都不动
=============================================================================

对应讲义：lectures/lecture06_adaptation.md 第 6 节

-----------------------------------------------------------------------------
【这个文件要回答的问题 What this file answers】
-----------------------------------------------------------------------------
"kari 搬家了，从 varno 搬到了 osku。" 怎样让 base model 知道这件事，而不影响它知道的其他事情？
模型编辑（model editing）的三个要求（MEND、ROME 论文）：

    可靠性 reliability ：用编辑时那句话问，答出新城市
    泛化 generalization：换一种说法问（"kari moved to"、"kari is a baker from"、问答格式……）也答新城市
    局部性 locality    ：别的人住在哪、kari 自己的职业，都不变

    演示 1：事实存在哪里？用 ROME 的"因果追踪"（causal tracing）在 base model 里找"回忆城市"发生的位置：
            把名字的嵌入加上噪声（模型就想不起城市了），再把某一层、某一个位置的"干净"状态放回去，
            看能恢复多少。
    演示 2：编辑 10 个人的城市，每次只编辑一个，比较：
                微调这一句话（全部参数）、只微调第 2 层的 MLP、少走几步的微调、
                ROME 式的秩一编辑（rank-one edit）：只改第 2 层 MLP 输出矩阵的一个秩一分量
    演示 3：连续编辑：一个接一个地改 1、5、10、20 个人，编辑会不会互相干扰、会不会把模型改坏？

-----------------------------------------------------------------------------
【运行方式】
-----------------------------------------------------------------------------
    python3 05_model_editing.py          # 约 10 秒

需要 PyTorch、numpy。依赖本目录的 adapt_lab.py，以及 code/lecture01 的 fm_world.py、tiny_gpt.py、console.py。

【PyTorch 小知识】
    - register_forward_hook(fn)：fn(module, inputs, output) 在模块每次前向之后被调用；
      fn 返回一个张量时，它会**替换**这个模块的输出。演示 1 用它"放回干净的状态"，演示 2 用它读取 MLP 的中间激活（键）。
    - torch.linalg.solve(C, k)：解线性方程组 C·u = k，比先求逆 C⁻¹ 再乘更快、更准。
    - torch.outer(a, b)：外积 a·bᵀ，一个秩一矩阵。
"""

import copy
import math
import os
import random
import sys
import time

import numpy as np
import torch
import torch.nn.functional as F

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))
from adapt_lab import TOK, choice_predictions, encode_pairs, load_base  # noqa: E402
from console import pad, print_header  # noqa: E402
from fm_world import CITIES, JOBS, qa_prompt  # noqa: E402

EDIT_LAYER = 1          # 编辑第 2 层（下标 1）的 MLP —— 演示 1 的因果追踪会说明为什么


# =============================================================================
# 演示 1：因果追踪
# =============================================================================
def run_traced(model, idx, noise=None, noise_pos=None, restore=None):
    """
    手写前向（和 TinyGPT 相同），可以 (1) 给某些位置的字符嵌入加噪声，(2) 把某一层、某种状态、某个位置的值换成给定值。

    【参数】noise：(B, len(noise_pos), d) 的噪声；restore：dict {(层, 种类): (位置, 值)}，种类是
            "resid"（这一层之后的残差流）、"attn"（注意力的输出）、"mlp"（MLP 的输出）
    【返回】(logits, states)：states[(层, 种类)] 是 (B, T, d) 的张量，供"干净"的那次运行保存下来
    """
    B, T = idx.shape
    x = model.tok_emb(idx)
    if noise is not None:
        x = x.clone()
        x[:, noise_pos] = x[:, noise_pos] + noise
    x = x + model.pos_emb(torch.arange(T))
    states = {}

    def maybe_restore(key, val):
        if restore and key in restore:
            pos, clean = restore[key]
            val = val.clone()
            val[:, pos] = clean
        states[key] = val
        return val

    for li, blk in enumerate(model.blocks):
        x = x + maybe_restore((li, "attn"), blk.attn(blk.ln1(x)))
        x = x + maybe_restore((li, "mlp"), blk.mlp(blk.ln2(x)))
        x = maybe_restore((li, "resid"), x)
    return model.head(model.ln_f(x)), states


def city_prob(logits, city):
    """
    最后一个位置上，正确城市的概率（只在 8 个城市的首字母之间归一化）。

    【思路】提示以 "lives in " 结尾，下一个字符就是城市的首字母；8 个城市的首字母互不相同（v t o b q d s p），
            所以"首字母的分布"就是"城市的分布"，一次前向就够，不用逐个候选打分。
    【参数】logits：(B, T, V)；city：正确城市
    【返回】(B,) 的张量
    """
    firsts = torch.tensor([TOK.stoi[c[0]] for c in CITIES])
    p = F.softmax(logits[:, -1][:, firsts], dim=-1)
    return p[:, CITIES.index(city)]


def demo_causal_trace(world, base):
    print_header("演示 1：事实存在哪里？—— 因果追踪：名字加噪声，再把某处的干净状态放回去")
    n_noise = 10
    emb_std = base.tok_emb.weight.std().item()
    kinds = [("resid", "残差流"), ("attn", "注意力输出"), ("mlp", "MLP 输出")]
    where = ["名字第一个字符", "名字最后一个字符", "lives 的 v", "最后一个位置（in 后的空格）"]
    effects = {(li, k, w): [] for li in range(2) for k, _ in kinds for w in where}
    clean_p, corrupt_p = [], []
    with torch.no_grad():
        for i, p in enumerate(world.people):
            prompt = f"\n{p.name} lives in "
            idx = torch.tensor([TOK.encode(prompt)])
            T = idx.shape[1]
            name_pos = list(range(1, 1 + len(p.name)))
            pos = {where[0]: name_pos[0], where[1]: name_pos[-1], where[2]: name_pos[-1] + 4, where[3]: T - 1}
            logits, clean = run_traced(base, idx)
            clean_p.append(city_prob(logits, p.city).item())
            g = torch.Generator().manual_seed(i)
            noise = torch.randn(n_noise, len(name_pos), base.cfg.d_model, generator=g) * 3 * emb_std
            idxb = idx.expand(n_noise, -1)
            lc, _ = run_traced(base, idxb, noise=noise, noise_pos=name_pos)
            pc = city_prob(lc, p.city).mean().item()
            corrupt_p.append(pc)
            for li in range(2):
                for k, _ in kinds:
                    for w in where:
                        lr_, _ = run_traced(base, idxb, noise=noise, noise_pos=name_pos,
                                            restore={(li, k): ([pos[w]], clean[(li, k)][:, pos[w]])})
                        effects[(li, k, w)].append(city_prob(lr_, p.city).mean().item() - pc)
    print(f"60 个人，提示 \"<名字> lives in \"。正确城市的概率（在 8 个城市里归一化）：干净 {np.mean(clean_p):.3f}，"
          f"名字加噪声后 {np.mean(corrupt_p):.3f}（每人 {n_noise} 次噪声的平均）。")
    print("下表：把这一处的干净状态放回去之后，正确城市的概率**回升**了多少（ROME 论文叫 indirect effect）：\n")
    print(pad("层 · 状态", 22) + "".join(pad(w, 30, ">") for w in where))
    for li in range(2):
        for k, kname in kinds:
            print(pad(f"第 {li + 1} 层 · {kname}", 22)
                  + "".join(pad(f"{np.mean(effects[(li, k, w)]):+.3f}", 30, ">") for w in where))
    print("""
观察 Observations：
  · 放回名字所在位置的任何一个状态，几乎都没用；放回**最后一个位置**第 2 层的状态，概率就回来了大半，
    其中第 2 层 MLP 的输出单独就能恢复一大半。在这个 2 层的字符级模型里，"回忆城市"发生在**要说出城市的那个位置**：
    第 2 层的注意力把名字的各个字符"搬"过来，第 2 层的 MLP 把"这个名字"变成"这个城市"。
  · 这和 ROME 在 GPT-2 XL 上找到的位置不同：那里是"主语最后一个 token"处的**中间层** MLP。我们的模型只有 2 层、
    没有"一个 token 就是一个名字"的分词，名字的信息要到最后才被汇总起来。方法相同，答案取决于模型。
  · 所以演示 2 把编辑放在第 2 层 MLP、最后一个位置上。""")


# =============================================================================
# 演示 2：编辑一个事实
# =============================================================================
TEMPLATES = {
    "编辑句 lives in": lambda p: f"\n{p.name} lives in",
    "moved to": lambda p: f"\n{p.name} moved to",
    "is a JOB from": lambda p: f"\n{p.name} is a {p.job} from",
    "the JOB X lives in": lambda p: f"\nthe {p.job} {p.name} lives in",
    "问答格式": lambda p: "\n" + qa_prompt(p.name, "city"),
}


def evaluate_edit(model, person, new_city, others):
    """
    一次编辑的评分（全部用受限选择）：
        可靠性：编辑句答新城市；泛化：另外 3 种叙述句 + 问答格式答新城市；
        局部性：其他人的 "lives in" 仍然答对的比例；这个人的职业仍然答对。
    """
    out = {}
    for name, t in TEMPLATES.items():
        out[name] = float(choice_predictions(model, [t(person)], CITIES)[0] == new_city)
    preds = choice_predictions(model, [f"\n{q.name} lives in" for q in others], CITIES)
    out["别人的城市"] = float(np.mean([a == q.city for a, q in zip(preds, others)]))
    out["别人也搬去新城"] = float(np.mean([a == new_city and q.city != new_city for a, q in zip(preds, others)]))
    out["自己的职业"] = float(choice_predictions(model, [f"\n{person.name} works as a"], JOBS)[0] == person.job)
    return out


def finetune_edit(model, person, new_city, steps, lr, only_mlp=False, replay=None):
    """
    最朴素的编辑：在"<名字> lives in <新城市>."这一句话上微调（只在答案字符上算损失）。

    【参数】only_mlp：True 时只训练第 2 层 MLP 的参数（和 ROME 改同一个地方，但用梯度下降）
            replay   ：None，或 list[Person]——把这些人的"<名字> lives in <模型当前的答案>."一起训练，
                       让模型在他们身上"保持原样"（MEND 论文的 FT+KL 基线的简化版；需要额外的数据）
    【返回】编辑后的模型（拷贝）
    """
    m = copy.deepcopy(model)
    params = list(m.blocks[EDIT_LAYER].mlp.parameters()) if only_mlp else list(m.parameters())
    for p in params:
        p.requires_grad_(True)
    pairs = [(f"\n{person.name} lives in", f" {new_city}.")]
    if replay:
        current = choice_predictions(model, [f"\n{q.name} lives in" for q in replay], CITIES)
        pairs += [(f"\n{q.name} lives in", f" {c}.") for q, c in zip(replay, current)]
    x, y, mask = encode_pairs(pairs)
    opt = torch.optim.AdamW(params, lr=lr, weight_decay=0.0)
    for _ in range(steps):
        logp = F.log_softmax(m(x), -1).gather(-1, y[..., None])[..., 0]
        loss = -(logp * mask).sum() / mask.sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    for p in m.parameters():
        p.requires_grad_(False)
    return m


def mlp_keys(model, idx):
    """
    编辑层的 MLP 里，输出矩阵 W_out 的输入（GELU 之后的 256 维激活）——ROME 说的"键" k。

    【思路】用前向钩子在 mlp[2]（第二个线性层）被调用时把它的输入存下来。
    【参数】model：TinyGPT；idx：(B, T) 的字符编号
    【返回】(B, T, 256) 的张量
    """
    store = {}
    lin = model.blocks[EDIT_LAYER].mlp[2]
    h = lin.register_forward_hook(lambda mod, inp, out: store.__setitem__("k", inp[0].detach()))
    with torch.no_grad():
        model(idx)
    h.remove()
    return store["k"]


def key_covariance(world, base, n_chars=60000):
    """
    C = E[k kᵀ]：在训练语料的前 n_chars 个字符上，所有位置的键的二阶矩（256 × 256）。

    【思路】ROME 在一大批维基百科文本上估计它，代表"平时的键长什么样"；秩一编辑用 C⁻¹ 把改动
            推向"平时少见"的方向，尽量不碰平时常见的输入。
    【返回】(256, 256) 的张量
    【边界】样本太少时 C 会接近奇异，rome_edit 给它加了 λI。
    """
    data = torch.tensor(TOK.encode(world.train_text()[:n_chars + 1]))
    T = 96
    xs = torch.stack([data[i:i + T] for i in range(0, len(data) - T, T)])
    K = mlp_keys(base, xs).reshape(-1, 256)
    return K.t() @ K / K.shape[0]


def rome_edit(model, person, new_city, C, reg=0.1, lam=1e-3, steps=100):
    """
    ROME 式的秩一编辑（Meng et al. 2022 的简化版），改第 2 层 MLP 的输出矩阵 W（64 × 256）：

        1. 键 k*：编辑句 "<名字> lives in " 最后一个位置上，W 的输入（256 维）
        2. 值 v*：W 在那个位置"应该"输出的 64 维向量——用梯度下降找：把 MLP 输出换成 v，
                  让模型在那里说出新城市的概率最大，同时 ‖v − v0‖² 不要太大（reg 控制）
        3. 改 W，使 W' k* = v*，并让"别的键"的输出尽量不变：
                  W' = W + Λ (C⁻¹ k*)ᵀ，   Λ = (v* − W k* − b) / ((C⁻¹ k*)ᵀ k*)
           C = E[k kᵀ] 描述了"平时的键长什么样"；C⁻¹ k* 挑出的是 k* 里"平时少见"的方向，
           于是这个秩一的改动主要作用在像 k* 的输入上，对平时常见的输入影响小。
    【参数】C：key_covariance 的结果；reg：v 的 L2 正则；lam：给 C 加 λI，防止求解不稳定
    【返回】编辑后的模型（拷贝）
    【FM 联系】ROME 把 MLP 看成一个"键 → 值"的线性联想记忆：键是"主语"，值是"关于它的事实"。
            秩一编辑就是往这个记忆里写一条新的键值对，同时尽量不碰其他键值对（论文附录 A 的闭式解）。
    """
    m = copy.deepcopy(model)
    lin = m.blocks[EDIT_LAYER].mlp[2]
    prompt = f"\n{person.name} lives in "
    idx = torch.tensor([TOK.encode(prompt)])
    last = idx.shape[1] - 1
    k = mlp_keys(m, idx)[0, last]
    with torch.no_grad():
        v0 = lin.weight @ k + lin.bias
    v = v0.clone().requires_grad_(True)
    x, y, mask = encode_pairs([(prompt, f"{new_city}.")])
    opt = torch.optim.Adam([v], lr=0.5)

    def put_v(mod, inp, out):
        out = out.clone()
        out[:, last] = v
        return out

    for _ in range(steps):
        h = lin.register_forward_hook(put_v)
        logp = F.log_softmax(m(x), -1).gather(-1, y[..., None])[..., 0]
        h.remove()
        loss = -(logp * mask).sum() + reg * ((v - v0) ** 2).sum()
        opt.zero_grad()
        loss.backward()
        opt.step()
    with torch.no_grad():
        u = torch.linalg.solve(C + lam * torch.eye(C.shape[0]), k)       # C⁻¹ k*
        Lam = (v.detach() - v0) / (u @ k)
        lin.weight += torch.outer(Lam, u)
    return m


def demo_single_edits(world, base, C):
    print_header("演示 2：编辑一个事实 —— 10 个人，每人改一个城市（每次从原模型出发，只编辑一个）")
    rng = random.Random(0)
    targets = rng.sample(world.people, 10)
    edits = [(p, rng.choice([c for c in CITIES if c != p.city])) for p in targets]
    def replay_people(p):
        """回放用的 40 个"别人"（每次编辑固定一组，和被评测的其余 59 人有重叠——这对微调是有利的）。"""
        others = [q for q in world.people if q.name != p.name]
        return random.Random(len(p.name) * 7 + ord(p.name[0])).sample(others, 40)

    methods = {
        "不编辑（base model）": lambda p, c: base,
        "微调这一句：lr 1e-3，30 步": lambda p, c: finetune_edit(base, p, c, 30, 1e-3),
        "微调这一句：lr 1e-4，10 步": lambda p, c: finetune_edit(base, p, c, 10, 1e-4),
        "微调这一句：lr 1e-4，5 步": lambda p, c: finetune_edit(base, p, c, 5, 1e-4),
        "只微调第 2 层 MLP：lr 1e-3，30 步": lambda p, c: finetune_edit(base, p, c, 30, 1e-3, only_mlp=True),
        "微调 + 回放 40 个别人：30 步": lambda p, c: finetune_edit(base, p, c, 30, 1e-3, replay=replay_people(p)),
        "ROME 式秩一编辑（正则 0.1）": lambda p, c: rome_edit(base, p, c, C, reg=0.1),
        "ROME 式秩一编辑（正则 0.01）": lambda p, c: rome_edit(base, p, c, C, reg=0.01),
    }
    cols = list(TEMPLATES) + ["别人的城市", "别人也搬去新城", "自己的职业"]
    print(pad("方法", 34) + "".join(pad(c, 20, ">") for c in cols))
    summary = {}
    for name, fn in methods.items():
        rows = []
        for p, new in edits:
            others = [q for q in world.people if q.name != p.name]
            rows.append(evaluate_edit(fn(p, new), p, new, others))
        avg = {c: np.mean([r[c] for r in rows]) for c in cols}
        summary[name] = avg
        print(pad(name, 34) + "".join(pad(f"{avg[c]:.2f}", 20, ">") for c in cols))
    print("""
（前 5 列：用这种问法问被编辑的人，答出**新**城市的比例——第 1 列是可靠性，第 2–5 列是泛化；
  "别人的城市"：其余 59 人仍然答对的比例（局部性）；"别人也搬去新城"：其余的人里被答成新城市的比例；
  "自己的职业"：被编辑的人的职业仍然答对。每格是 10 次编辑的平均。）

观察 Observations：
  · 微调这一句话（lr 1e-3、30 步）：可靠性、泛化都是满分，但别人的城市只剩一成多答对——大部分人被答成了**新城市**。
    模型学到的不是"kari 住在 osku"，而是"lives in 后面写 osku"。只微调第 2 层 MLP（和 ROME 改同一处）也一样。
  · 把学习率调小、步数减少，局部性回来一些，新事实也写得不那么牢、泛化变差：这是一个此消彼长的旋钮。
    把 40 个别人的事实一起训练（回放）能保住大部分局部性，但需要额外的数据，而且被编辑的人的职业照样被带歪了。
  · ROME 式的秩一编辑只改一个 64 × 256 矩阵的一个秩一分量：可靠性满分，局部性最好（0.85–0.88，不编辑时 0.93），
    这个人的职业也基本保住；泛化到其他说法只有六到九成。正则越弱（v* 走得越远），泛化略好、局部性略差。
  · 对照 ROME 论文在 GPT-2 XL 上的数字（CounterFact）：微调 效能 100.0 / 改写 87.9 / 邻居 40.4，
    ROME 100.0 / 96.4 / 75.4——同样的模式：微调在"邻居"上崩，秩一编辑两头兼顾。我们的玩具里 ROME 的泛化弱一些，
    一个可能的原因：这里的编辑位置是句子的最后一个位置，键 k* 里混着 "lives in" 这个句式的信息，换一种说法键就变了；
    ROME 原文编辑的是主语所在的位置，键基本只取决于主语（还在 50 种随机前缀上取了平均）。""")
    return edits, summary


# =============================================================================
# 演示 3：连续编辑
# =============================================================================
def demo_sequential(world, base, C):
    print_header("演示 3：连续编辑 —— 一个接一个地改 20 个人的城市，前面的编辑还在吗？模型还好吗？")
    rng = random.Random(1)
    targets = rng.sample(world.people, 20)
    edits = [(p, rng.choice([c for c in CITIES if c != p.city])) for p in targets]
    edited_names = {p.name for p, _ in edits}
    untouched = [q for q in world.people if q.name not in edited_names]
    print(pad("方法", 26) + pad("已编辑人数", 12, ">") + pad("所有已编辑的事实（可靠性）", 28, ">")
          + pad("没编辑的人（局部性）", 22, ">") + pad("已编辑者的职业", 16, ">"))
    for name, step_fn in [("ROME 式（正则 0.1）", lambda m, p, c: rome_edit(m, p, c, C, reg=0.1)),
                          ("微调这一句（lr 1e-4，5 步）", lambda m, p, c: finetune_edit(m, p, c, 5, 1e-4))]:
        model = base
        for i, (p, new) in enumerate(edits, start=1):
            model = step_fn(model, p, new)
            if i in (1, 5, 10, 20):
                done = edits[:i]
                rel = np.mean([choice_predictions(model, [f"\n{q.name} lives in"], CITIES)[0] == c for q, c in done])
                preds = choice_predictions(model, [f"\n{q.name} lives in" for q in untouched], CITIES)
                loc = np.mean([a == q.city for a, q in zip(preds, untouched)])
                jobs = choice_predictions(model, [f"\n{q.name} works as a" for q, _ in done], JOBS)
                job_ok = np.mean([a == q.job for a, (q, _) in zip(jobs, done)])
                print(pad(name, 26) + pad(i, 12, ">") + pad(f"{rel:.2f}", 28, ">") + pad(f"{loc:.2f}", 22, ">")
                      + pad(f"{job_ok:.2f}", 16, ">"))
    base_loc = np.mean([a == q.city for a, q in zip(
        choice_predictions(base, [f"\n{q.name} lives in" for q in untouched], CITIES), untouched)])
    print(f"（没编辑的 {len(untouched)} 人在 base model 上的准确率是 {base_loc:.2f}。）")
    print("""
观察 Observations：
  · 一个一个地改，编辑会互相干扰：ROME 式编辑改到第 20 个时，前面写进去的事实只剩一半多，没编辑的人也只剩一半左右答对。
    小学习率的微调退化得更快：改 5 个之后，已编辑的事实只剩 0.40、没编辑的人 0.33。
    每次编辑单独看都"成功"，累积起来却在慢慢把模型改坏——这是把编辑当作日常维护手段时最大的风险。
  · MEND 论文的批量编辑实验（BART-base，zsRE）里，同时编辑 1 个时成功率 0.98，125 个时降到 0.67；
    MEMIT（2022）专门为"一次写入成千上万个事实"重新设计了更新公式，把改动分摊到多层 MLP 上。
  · 更根本的问题（讲义 §6.6）：改了"kari 住在 osku"，那么"kari 的邻居""osku 的居民"这些**相关**的事实要不要跟着变？
    RippleEdits（Cohen et al. 2023）专门测这种"涟漪效应"，发现现有的编辑方法大多做不到。""")


if __name__ == "__main__":
    T0 = time.time()
    world, base = load_base()
    demo_causal_trace(world, base)
    C = key_covariance(world, base)
    demo_single_edits(world, base, C)
    demo_sequential(world, base, C)
    print(f"\n（总用时约 {time.time() - T0:.0f} 秒；计时每次运行略有不同）")