""" rag_world.py —— 第 11 讲：花园百科——一个"知识的流行度"已知的世界，和检索增强生成（RAG）的全部零件 ===================================================================== 【这个文件做什么】 RAG 要回答的问题是：语言模型把知识存在参数里（parametric memory），可是它记不住冷门的知识、知识会过时、 不知道时会胡编；把"查资料"（non-parametric memory，一个可检索的文档库）接到生成上，能解决多少？ 要精确地研究这件事，我们需要一个**知道每条知识在预训练语料里出现了几次**的世界： 400 位园丁（名 + 姓，姓是独一无二的造词，例如 Ada Brask），每位有三条事实： Ada Brask lives in Coveton . 住在哪个镇（16 个镇，大镇更常见） Ada Brask grows tulip . 种什么（35 种植物） Ada Brask keeps goat . 养什么（8 种动物） 园丁的"流行度"服从 Zipf 定律：排第 r 位的园丁被提到的概率 ∝ r^(-1.2)。 所以排第 1 的园丁在语料里出现几千次，排第 400 的可能一次都没有——这正是真实网络语料里 "名人 vs 普通人"的样子（Mallen et al. 2023 的 PopQA 用维基百科的月浏览量来度量流行度）。 预训练语料（"网页"）有两种文档： 传记：Ada Brask lives in Coveton . Ada Brask keeps goat . （一位园丁的 1～3 条事实） 新闻：Ada Brask met Pia Holm in Ivydale . （两位园丁在某个镇见面——这个镇 和他们住哪儿无关，是天然的"干扰项"） 文档库（datastore，"百科全书"）：85% 的园丁有一篇百科条目（三条事实都在），外加一批新闻。 这里提供： (1) 世界、语料、文档库、切块（chunking）； (2) 检索：BM25 和稠密检索器（沿用第 10 讲的 retrieval.py），多正例的排序指标（R@k、nDCG@10）； (3) 语言模型：闭卷（只靠参数）和"读者"（reader，读检索到的段落再回答）的训练与推断； (4) kNN-LM 的数据库与检索（04 号脚本）。 【LLM 联系】 这个世界对应的真实研究：Kandpal et al. (2023) 统计了预训练语料里"问题和答案同时出现的文档数"， 发现 LLM 的问答准确率随这个数单调上升；Mallen et al. (2023) 发现检索对冷门知识帮助巨大、对热门知识 反而可能有害；Allen-Zhu & Li (2023, Physics of LMs 3.x) 用合成的人物传记研究模型怎样存储知识。 """ import math import os import random import sys import numpy as np import torch import torch.nn.functional as F HERE = os.path.dirname(os.path.abspath(__file__)) for _d in ("lecture01", "lecture02", "lecture10"): sys.path.insert(0, os.path.join(HERE, "..", _d)) from garden_lm import Vocab, clm_batches, clm_loss, load_or_train, lr_lambda, train # noqa: E402 from toy_world import NOUNS # noqa: E402 import retrieval as R # noqa: E402 CKPT_DIR = os.path.join(HERE, "checkpoints") # ===================================================================== # 1. 世界设定 # ===================================================================== FIRST = ["Ada", "Ben", "Cleo", "Dov", "Esme", "Finn", "Gia", "Hugo", "Ines", "Jude", "Kai", "Lena", "Milo", "Nia", "Otto", "Pia", "Quin", "Rosa", "Saul", "Tova"] TOWNS = ["Ashby", "Brookfield", "Coveton", "Dunmore", "Elmwick", "Fallowby", "Greystone", "Hollins", "Ivydale", "Kestrel", "Larkhill", "Millbrook", "Northam", "Oakridge", "Pebworth", "Redcliff"] TOWN_PRIOR = [1 / (i + 2) for i in range(len(TOWNS))] # 大镇更常见：Ashby 约 20%，Redcliff 约 2.4% PLANTS = NOUNS["fruit"] + NOUNS["veg"] + NOUNS["flower"] # 35 种 ANIMALS = NOUNS["animal"] # 8 种 RELATIONS = ["town", "plant", "animal"] VALUES = {"town": TOWNS, "plant": PLANTS, "animal": ANIMALS} VERB = {"town": ["lives", "in"], "plant": ["grows"], "animal": ["keeps"]} # 用户的提问方式（检索器看到的是问题；读者看到的是"完形填空"式的提示，见 cloze()） QUESTIONS = { "town": ["where does {e} live ?", "which town is home to {e} ?", "in which town does {e} live ?"], "plant": ["what does {e} grow ?", "which plant does {e} grow ?", "what is planted by {e} ?"], "animal": ["which animal does {e} keep ?", "what pet does {e} have ?", "what animal is kept by {e} ?"], } ZIPF_S = 1.2 N_ENT = 400 SPECIAL = ["", "", "", "", "", "", "", "", ""] FUNC = ["lives", "in", "grows", "keeps", "met", "sold", "to", "walked", "with", "."] def make_surnames(n, seed=0): """ 造 n 个互不相同、像姓氏的词（Brask、Tumel……）。 【思路】辅音 + 元音 + 结尾辅音，有一半再加一个音节。造词而不用真实姓氏： 既避免和真实人物重名，也保证每个姓在"世界"里独一无二——知道姓就知道是哪位园丁。 """ rng = random.Random(seed) on = ["b", "br", "d", "dr", "f", "g", "gr", "k", "l", "m", "n", "p", "r", "s", "st", "t", "tr", "v", "w", "z"] vo = ["a", "e", "i", "o", "u"] co = ["n", "r", "l", "s", "m", "k", "x", "sk", "nd", "lt"] out, seen = [], set() banned = {w.lower() for w in FIRST + TOWNS + PLANTS + ANIMALS + FUNC} while len(out) < n: w = rng.choice(on) + rng.choice(vo) + rng.choice(co) if rng.random() < 0.5: w += rng.choice(vo) + rng.choice(["n", "r", "l", "s", "m"]) if w not in seen and w not in banned: seen.add(w) out.append(w.capitalize()) return out class Entity: """一位园丁。rank：流行度排名（1 最热门）；name：[名, 姓]；facts：{关系: 值}。""" def __init__(self, idx, first, surname, facts): self.idx, self.rank = idx, idx + 1 self.name = [first, surname] self.facts = facts def text(self): return " ".join(self.name) class World: """ 400 位园丁和他们的事实。 【参数】seed：决定姓名和事实；园丁按流行度排序（第 0 位最热门）。 【方法】 weights：每位园丁被语料提到的相对概率（Zipf，指数 ZIPF_S）； updated(frac, seed)：返回一个"世界变了"的副本——frac 比例的园丁搬了家（住址换成另一个镇）， 用来研究知识过时（01、05 号脚本）。 """ def __init__(self, seed=0): rng = random.Random(seed) surnames = make_surnames(N_ENT, seed) self.entities = [] for i in range(N_ENT): facts = {"town": rng.choices(TOWNS, weights=TOWN_PRIOR)[0], "plant": rng.choice(PLANTS), "animal": rng.choice(ANIMALS)} self.entities.append(Entity(i, rng.choice(FIRST), surnames[i], facts)) w = [(i + 1) ** -ZIPF_S for i in range(N_ENT)] z = sum(w) self.weights = [x / z for x in w] def updated(self, frac=0.2, seed=7): import copy rng = random.Random(seed) new = copy.deepcopy(self) moved = sorted(rng.sample(range(N_ENT), int(frac * N_ENT))) for i in moved: old = new.entities[i].facts["town"] new.entities[i].facts["town"] = rng.choice([t for t in TOWNS if t != old]) new.moved = moved return new def fact_sentence(ent, rel, value=None): """一条事实的句子：Ada Brask lives in Coveton .（value 可以替换成别的值，用来造反事实）""" return ent.name + VERB[rel] + [value if value is not None else ent.facts[rel], "."] def cloze(ent, rel): """读者的提示（完形填空，和 LAMA 探针一样）：Ada Brask lives in ▢。""" return ent.name + VERB[rel] def question(ent, rel, template=0): """用户的问题（检索器的查询）：where does Ada Brask live ?""" return QUESTIONS[rel][template].format(e=ent.text()) def build_vocab(world): names = FIRST + [e.name[1] for e in world.entities] return Vocab(SPECIAL + FUNC + names + TOWNS + PLANTS + ANIMALS) WORLD = World() VOCAB = build_vocab(WORLD) BOS, EOS, PAD = VOCAB.stoi[""], VOCAB.stoi[""], VOCAB.stoi[""] CTX, ECTX, SEP, REL, IRR = (VOCAB.stoi[w] for w in ["", "", "", "", ""]) # ===================================================================== # 2. 预训练语料与文档库 # ===================================================================== def bio_doc(rng, world, ent): """一篇传记：随机挑 1～3 条事实（每条以 0.6 的概率出现，至少一条），随机顺序。【返回】(词列表, 出现的关系)。""" rels = [r for r in RELATIONS if rng.random() < 0.6] or [rng.choice(RELATIONS)] rng.shuffle(rels) words = [] for r in rels: words += fact_sentence(ent, r) return words, rels def news_doc(rng, world, e1, e2, p_recap=0.0): """ 一篇新闻：两位园丁之间的一件事，三种模板各占三分之一—— Ada Brask met Pia Holm in Ivydale . （镇是随机的，和他们住哪儿无关） Ada Brask sold tulip to Pia Holm . （植物是随机的，和谁种什么无关） Ada Brask walked goat with Pia Holm . （动物是随机的，和谁养什么无关） 这些随机的值是读者的天然"干扰项"：检索到这篇新闻、却把里面的镇名当成住址，就答错了。 以 p_recap 的概率把这句话原样再说一遍（"回顾"）。回顾句里的随机值**只能靠照抄前文**预测（记忆帮不上忙）， 这让预训练产生复制能力——读者（03 号脚本）读资料靠的正是它（第 6、10 讲：复制能力来自预训练里的重复）。 """ kind = rng.choice(RELATIONS) if kind == "town": s = e1.name + ["met"] + e2.name + ["in", rng.choice(TOWNS), "."] elif kind == "plant": s = e1.name + ["sold", rng.choice(PLANTS), "to"] + e2.name + ["."] else: s = e1.name + ["walked", rng.choice(ANIMALS), "with"] + e2.name + ["."] return s + s if rng.random() < p_recap else s def _pair(rng, world): a, b = rng.choices(range(N_ENT), weights=world.weights, k=2) return world.entities[a], world.entities[(b + 1) % N_ENT if a == b else b] def pretrain_corpus(world=WORLD, n_bio=18000, n_news=24000, p_recap=0.5, seed=1): """ 预训练语料：n_bio 篇传记 + n_news 篇新闻，随机打乱。 【参数】p_recap：新闻带回顾句的比例；园丁按 world.weights（Zipf）被挑中。 【返回】(docs, counts)：docs 是词列表（以 结尾）的列表；counts[(园丁下标, 关系)] = 这条事实 出现在几篇传记里——这就是本讲的"流行度"（Kandpal et al. 2023 的做法：数一数相关文档有多少）。 新闻不提供任何事实，所以不计入。 """ rng = random.Random(seed) kinds = ["news"] * n_news + ["bio"] * n_bio rng.shuffle(kinds) docs, counts = [], {} for k in kinds: if k == "news": docs.append(news_doc(rng, world, *_pair(rng, world), p_recap=p_recap) + [""]) else: i = rng.choices(range(N_ENT), weights=world.weights)[0] words, rels = bio_doc(rng, world, world.entities[i]) for r in rels: counts[(i, r)] = counts.get((i, r), 0) + 1 docs.append(words + [""]) return docs, counts CORPUS, COUNTS = pretrain_corpus() def article(ent, world=None, override=None): """百科条目：三条事实按固定顺序。override = {关系: 值} 可以改写其中几条（反事实 / 更新后的知识）。""" words = [] for r in RELATIONS: words += fact_sentence(ent, r, (override or {}).get(r)) return words def build_datastore(world=WORLD, coverage=0.85, n_news=3000, seed=3, covered=None): """ 文档库：百科条目 + 新闻。 【参数】coverage：有百科条目的园丁比例（随机挑，和流行度无关）；covered：直接指定有条目的园丁下标集合； n_news：新闻篇数（另外采样的新闻，和预训练语料同分布）。 【返回】dict：docs（每篇文档的词列表）、kind（"article" / "news"）、owner（条目属于哪位园丁，新闻为 -1）、 covered（有条目的园丁下标集合）。 """ if covered is None: covered = set(random.Random(seed).sample(range(N_ENT), int(coverage * N_ENT))) rng = random.Random(seed + 1) # 新闻的随机数和覆盖哪些园丁无关：换覆盖率时新闻不变 docs, kind, owner = [], [], [] for i in sorted(covered): docs.append(article(world.entities[i])) kind.append("article") owner.append(i) for _ in range(n_news): docs.append(news_doc(rng, world, *_pair(rng, world))) kind.append("news") owner.append(-1) return {"docs": docs, "kind": kind, "owner": owner, "covered": covered} def chunk(store, mode="doc", window=8): """ 切块（chunking）：检索的单位。 【参数】mode："doc" 一篇文档一块；"sentence" 按句号切；"window" 把全部文档首尾相接，每 window 个词切一块 （slides 的 "every k words"，会把句子拦腰切断）。 【返回】list[list[str]]：每块的词列表。 """ if mode == "doc": return [list(d) for d in store["docs"]] if mode == "sentence": out = [] for d in store["docs"]: cur = [] for w in d: cur.append(w) if w == ".": out.append(cur) cur = [] return out stream = [w for d in store["docs"] for w in d] return [stream[i:i + window] for i in range(0, len(stream), window)] def contains(chunk_words, sent): """块里是否完整地包含句子 sent（连续子串）。""" n = len(sent) return any(chunk_words[i:i + n] == sent for i in range(len(chunk_words) - n + 1)) # ===================================================================== # 3. 检索与指标 # ===================================================================== def bm25_index(chunks): return R.BM25([" ".join(c) for c in chunks]) def bm25_topk(bm, query, k): s = bm.scores(query) return [int(i) for i in np.argsort(-s, kind="stable")[:k]] def ranks_of(score_row, relevant): """正确块中排得最靠前的那个的名次（1 起）。relevant 为空时返回无穷大。""" if not relevant: return math.inf order = np.argsort(-score_row, kind="stable") pos = np.nonzero(np.isin(order, list(relevant)))[0] return int(pos[0]) + 1 def multi_metrics(score_matrix, relevant_sets, ks=(1, 5)): """ 多正例的排序指标。 【参数】score_matrix：(查询数, 块数)；relevant_sets：每个查询的正确块下标集合（可能有多个，例如句子被 几个窗口都完整包含；也可能为空——文档库里根本没有答案）。 【返回】{"R@1", "R@5", "nDCG@10", "MRR"}。 nDCG@10 = DCG@10 / 理想 DCG@10，DCG = Σ_{前 10 名里的正确块} 1 / log2(名次 + 1) （BEIR 的主要指标，slides 第 37 页）。 """ out = {f"R@{k}": 0.0 for k in ks} out["nDCG@10"] = 0.0 out["MRR"] = 0.0 for row, rel in zip(score_matrix, relevant_sets): if not rel: continue order = np.argsort(-row, kind="stable") hits = np.isin(order[:10], list(rel)) r = ranks_of(row, rel) for k in ks: out[f"R@{k}"] += r <= k dcg = sum(1 / math.log2(j + 2) for j in np.nonzero(hits)[0]) idcg = sum(1 / math.log2(j + 2) for j in range(min(10, len(rel)))) out["nDCG@10"] += dcg / idcg out["MRR"] += 1 / r n = len(relevant_sets) return {k: v / n for k, v in out.items()} # ===================================================================== # 4. 语言模型：训练 # ===================================================================== ARCHS = {"small": dict(d_model=64, n_heads=4, n_layers=2, d_ff=256), "large": dict(d_model=128, n_heads=4, n_layers=4, d_ff=512)} def new_model(size="large", seed=0): from transformer import TransformerLM torch.manual_seed(seed) return TransformerLM(len(VOCAB), attn="causal", **ARCHS[size]) def base_model(size="large", epochs=1, verbose=False): """ 闭卷模型：只在预训练语料（18000 篇传记 + 24000 篇新闻）上训练的语言模型。知识全部存在参数里。 【参数】size："small"（d=64、2 层，约 10 万参数）或 "large"（d=128、4 层，约 85 万参数）； epochs：遍数。默认 1 遍——和真实的 LLM 预训练一样，每篇文档只看一次， 于是"一条事实在语料里出现几次"就是模型见过它几次。 【返回】模型（缓存名 garden_lm_{size}）。第一次训练 large 约 15 秒。 """ batches = clm_batches(CORPUS, VOCAB, 32) def fit(m): return train(m, batches, clm_loss, epochs=epochs, lr=2e-3, warmup=200, seed=0, verbose=verbose, tag=size) m, _ = load_or_train(f"garden_lm_{size}", {"arch": ARCHS[size], "epochs": epochs, "corpus": "bio18k_news24k_recap0.5", "v": 4}, lambda: new_model(size), fit, verbose=verbose, ckpt_dir=CKPT_DIR) return m def masked_batches(examples, batch_size=32, seed=0): """ examples = [(提示词列表, 目标词列表)]：只在目标的位置上算损失。提示为空时就是普通的语言模型训练（回放）。 【返回】list[(ids, pad, loss_mask)]，按长度排序后切批（同 clm_batches）。 """ seqs = [] for p, t in examples: ids = [BOS] + VOCAB.encode(p + t) m = [0] * (1 + len(p)) + [1] * len(t) seqs.append((ids, m)) seqs.sort(key=lambda x: len(x[0])) out = [] for s in range(0, len(seqs), batch_size): b = seqs[s:s + batch_size] L = max(len(x[0]) for x in b) ids = torch.full((len(b), L), PAD, dtype=torch.long) pad = torch.ones(len(b), L, dtype=torch.bool) mask = torch.zeros(len(b), L) for i, (x, m) in enumerate(b): ids[i, :len(x)] = torch.tensor(x) pad[i, :len(x)] = False mask[i, :len(m)] = torch.tensor(m, dtype=torch.float) out.append((ids, pad, mask)) return out def masked_loss(model, batch): """只在 loss_mask = 1 的目标位置上算交叉熵。【返回】(NLL 之和, 计入的 token 数)。""" ids, pad, mask = batch logits = model(ids[:, :-1], pad=pad[:, :-1]) tgt, keep = ids[:, 1:], mask[:, 1:] nll = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), tgt.reshape(-1), reduction="none") return (nll * keep.reshape(-1)).sum(), keep.sum() # 读者训练用的园丁：每 4 位取 1 位（覆盖从热门到冷门的全部流行度）；其余 300 位只用来测试 READER_TRAIN = [i for i in range(N_ENT) if i % 4 == 1] TEST_ENT = [i for i in range(N_ENT) if i % 4 != 1] def rag_prompt(ctx_chunks, ent, rel): """读者的输入： 块 1 块 2 … Ada Brask lives in""" words = [""] for j, c in enumerate(ctx_chunks): if j: words.append("") words += c return words + [""] + cloze(ent, rel) def reader_examples(store, mode="plain", n_variants=24, p_irrelevant=0.4, seed=11, world=WORLD): """ 读者的训练样例（只用 READER_TRAIN 里的园丁）。 【参数】mode： "plain" 上下文里**总有**正确的百科条目（检索没找到就把它塞进去）——读者学会"照着资料答"； "robust" 以 p_irrelevant 的概率上下文里**没有**正确条目（只剩检索到的其他文档）。 读者先输出一个反思 token：（资料有用）或 （资料没用，到此为止；推断时改用 闭卷回答，见 read()）。 之后把问题复述一遍再给答案： … Ada Brask lives in Ada Brask lives in Coveton 这是 Self-RAG（Asai et al. 2024）"检索、生成、自我批判"的最简版本：先判断段落是否相关 （IsRel），再生成输出。 【为什么 后面不接答案】第一版让读者在 后面直接说出正确答案，结果它把 100 位训练园丁的事实背了下来，读新园丁的资料时反而会说出同名训练园丁的住址。 【为什么 之后要复述问题】第二版在 之后直接给答案，读资料的准确率从 99.6% 掉到 92%：照抄靠的是"找到上文里和当前词相同的位置、抄它后面的词"（induction head）， 答案紧跟在 in 后面时，资料里恰好有 "lives in Coveton" 可以对照；紧跟在 后面时就没有了。 n_variants：每个（园丁, 关系）生成几条（不同的问法、检索深度 k、块的顺序）。 【返回】[(提示, 目标)]。目标是答案词；robust 模式下是 [, 复述的问题…, 答案] 或 []。 【边界】上下文由 BM25 按问题检索文档库得到（k 在 1～3 之间随机）；块的顺序随机打乱，答案不总在第一块。 """ rng = random.Random(seed) chunks = chunk(store, "doc") bm = bm25_index(chunks) out = [] for i in READER_TRAIN: ent = world.entities[i] gold = article(ent) for rel in RELATIONS: for v in range(n_variants): k = rng.choice([1, 2, 3]) top = bm25_topk(bm, question(ent, rel, v % 3), k + 3) cands = [chunks[j] for j in top if chunks[j] != gold] irrelevant = mode == "robust" and rng.random() < p_irrelevant ctx = cands[:k] if irrelevant else [gold] + cands[:k - 1] rng.shuffle(ctx) target = [ent.facts[rel]] if mode == "robust": target = [""] if irrelevant else [""] + cloze(ent, rel) + target out.append((rag_prompt(ctx, ent, rel), target)) return out def reader_model(mode="plain", epochs=4, n_replay=4000, lr=2e-3, verbose=False): """ 读者：从闭卷模型（large）出发，在"检索到的资料 + 问题 → 答案"的样例上微调（7200 条，只在答案上算损失）， 同时回放 n_replay 篇预训练文档（第 6 讲：只在新任务上微调会灾难性遗忘——这里遗忘的就是参数里的知识）。 【参数】mode："plain" 或 "robust"（见 reader_examples）。 【返回】模型（缓存名 reader_{mode}）。第一次约 25 秒。 """ import copy store = build_datastore() def build(): return copy.deepcopy(base_model("large")) def fit(m): ex = reader_examples(store, mode) rng = random.Random(5) replay = [([], d) for d in rng.sample(CORPUS, n_replay)] batches = masked_batches(ex + replay) return train(m, batches, masked_loss, epochs=epochs, lr=lr, warmup=50, seed=0, verbose=verbose, tag=f"reader_{mode}") m, _ = load_or_train(f"reader_{mode}", {"mode": mode, "epochs": epochs, "replay": n_replay, "lr": lr, "v": 3 if mode == "plain" else 5}, build, fit, verbose=verbose, ckpt_dir=CKPT_DIR) return m # ===================================================================== # 5. 推断 # ===================================================================== @torch.no_grad() def next_probs(model, prompts): """ 一批提示（词列表）之后，下一个 token 在整个词表上的分布。 【思路】按长度分组，同一组不需要补齐——补齐和不补齐在浮点上可能有极小的差异，分组让结果和逐条计算完全一致。 【返回】(B, |V|) 的 numpy 数组。 """ ids_list = [[BOS] + VOCAB.encode(p) for p in prompts] out = [None] * len(prompts) groups = {} for i, p in enumerate(ids_list): groups.setdefault(len(p), []).append(i) for _, idx in groups.items(): for s in range(0, len(idx), 256): part = idx[s:s + 256] probs = torch.softmax(model(torch.tensor([ids_list[i] for i in part]))[:, -1], -1).numpy() for j, i in enumerate(part): out[i] = probs[j] return np.stack(out) def answer(probs_row): """概率最高的词和它的概率。""" t = int(probs_row.argmax()) return VOCAB.itos[t], float(probs_row[t]) def closed_book(model, queries, world=WORLD): """ 闭卷回答：只给"Ada Brask lives in"，看模型接下来说哪个词。 【参数】queries：[(园丁下标, 关系)]。 【返回】[(预测的词, 它的概率, 正确答案)]。 """ probs = next_probs(model, [cloze(world.entities[i], r) for i, r in queries]) return [answer(p) + (world.entities[i].facts[r],) for p, (i, r) in zip(probs, queries)] def read(model, ctx_list, queries, world=WORLD, reflect=False): """ 开卷回答：给读者 " 资料 Ada Brask lives in"。 【参数】ctx_list：每个查询的资料（块的列表）；reflect=True 时（robust 读者）先读出反思 token： p_rel = P() / (P() + P())。p_rel ≥ 0.5：接上 " 复述的问题"，照资料回答； 否则判定资料没用，丢掉资料、闭卷回答（同一个模型）。 【返回】[(预测的词, 概率, 正确答案)]；reflect=True 时每条再多一个 p_rel。 """ prompts = [rag_prompt(c, world.entities[i], r) for c, (i, r) in zip(ctx_list, queries)] probs = next_probs(model, prompts) if not reflect: return [answer(p) + (world.entities[i].facts[r],) for p, (i, r) in zip(probs, queries)] p_rel = probs[:, REL] / (probs[:, REL] + probs[:, IRR]) prompts2 = [p + [""] + cloze(world.entities[i], r) if pr >= 0.5 else cloze(world.entities[i], r) for p, pr, (i, r) in zip(prompts, p_rel, queries)] probs2 = next_probs(model, prompts2) return [answer(p) + (world.entities[i].facts[r], float(pr)) for p, pr, (i, r) in zip(probs2, p_rel, queries)] # ===================================================================== # 6. 流行度分桶 # ===================================================================== BUCKETS = [(0, 0, "0 次"), (1, 2, "1～2 次"), (3, 9, "3～9 次"), (10, 29, "10～29 次"), (30, 10 ** 9, "≥30 次")] def bucket(count): for lo, hi, name in BUCKETS: if lo <= count <= hi: return name def test_queries(rels=RELATIONS): """测试用的查询：TEST_ENT 里的 300 位园丁 × 关系。""" return [(i, r) for i in TEST_ENT for r in rels] def by_bucket(queries, correct): """按"这条事实在预训练语料里出现的次数"分桶统计准确率。【返回】{桶名: (准确率, 题数)}。""" acc = {} for (i, r), c in zip(queries, correct): b = bucket(COUNTS.get((i, r), 0)) s, n = acc.get(b, (0, 0)) acc[b] = (s + c, n + 1) return {name: (acc[name][0] / acc[name][1], acc[name][1]) for _, _, name in BUCKETS if name in acc} def bucket_row(queries, correct, width=10): """把 by_bucket 的结果格式化成一行（顺序固定）。""" d = by_bucket(queries, correct) return "".join(f"{d[name][0]:>{width}.1%}" if name in d else f"{'—':>{width}s}" for _, _, name in BUCKETS) def bucket_header(width=10): return "".join(f"{name:>{width - 1}s} " for _, _, name in BUCKETS) # ===================================================================== # 7. kNN-LM（Khandelwal et al. 2020） # ===================================================================== @torch.no_grad() def knn_datastore(model, docs): """ kNN-LM 的数据库：对语料里每个位置，键 = 模型在该位置的最后一层隐藏状态（"上文的表示"），值 = 下一个词。 【返回】(keys (N, d) 的 float32 张量, vals (N,) 的 long 张量, where = [(文档下标, 位置)] 长度 N 的列表)。 【思路】文档按长度分组，同一组一次前向、不需要补齐；where 记录每个键来自哪篇文档的哪个位置（用来展示近邻）。 """ keys, vals, where = [], [], [] groups = {} for n, d in enumerate(docs): groups.setdefault(len(d), []).append(n) for L in sorted(groups): idx = groups[L] for s in range(0, len(idx), 512): part = idx[s:s + 512] ids = torch.tensor([[BOS] + VOCAB.encode(docs[n]) for n in part]) h = model.hidden_states(ids[:, :-1]) keys.append(h.reshape(-1, h.shape[-1])) vals.append(ids[:, 1:].reshape(-1)) where += [(n, t) for n in part for t in range(L)] return torch.cat(keys), torch.cat(vals), where @torch.no_grad() def query_states(model, prompts): """一批提示（词列表）最后一个位置的隐藏状态：kNN-LM 的查询向量。【返回】(B, d) 张量，顺序同 prompts。""" out = [None] * len(prompts) groups = {} for n, p in enumerate(prompts): groups.setdefault(len(p), []).append(n) for _, idx in groups.items(): h = model.hidden_states(torch.tensor([[BOS] + VOCAB.encode(prompts[n]) for n in idx]))[:, -1] for j, n in enumerate(idx): out[n] = h[j] return torch.stack(out) @torch.no_grad() def knn_probs(keys, vals, queries, k=16, temp=1.0, chunk_size=512): """ p_kNN(y | 上文) ∝ Σ_{k 个最近邻里值为 y 的} exp(−‖q − key‖² / temp)。 【参数】queries：(M, d) 的隐藏状态；k：近邻个数；temp：距离的温度（越小越只信最近的那个）。 【返回】(M, |V|) 的张量。 【边界】这里用精确搜索（矩阵乘法算全部距离）；真实系统用 FAISS 等近似最近邻索引（slides 第 31 页）。 """ out = torch.zeros(len(queries), len(VOCAB)) kk = (keys * keys).sum(1) for s in range(0, len(queries), chunk_size): q = queries[s:s + chunk_size] d2 = (q * q).sum(1, keepdim=True) - 2 * q @ keys.T + kk[None] dist, idx = torch.topk(-d2, k, dim=1) w = torch.softmax(dist / temp, dim=1) out[s:s + len(q)].scatter_add_(1, vals[idx], w) return out @torch.no_grad() def knn_neighbors(keys, query, k=5): """一个查询向量的 k 个最近邻：(下标列表, 距离平方列表)。""" d2 = ((keys - query[None]) ** 2).sum(1) dist, idx = torch.topk(-d2, k) return idx.tolist(), (-dist).tolist() __all__ = ["FIRST", "TOWNS", "PLANTS", "ANIMALS", "RELATIONS", "VALUES", "QUESTIONS", "World", "WORLD", "VOCAB", "fact_sentence", "cloze", "question", "pretrain_corpus", "CORPUS", "COUNTS", "article", "build_datastore", "chunk", "contains", "bm25_index", "bm25_topk", "multi_metrics", "ARCHS", "new_model", "base_model", "masked_batches", "masked_loss", "READER_TRAIN", "TEST_ENT", "rag_prompt", "reader_examples", "reader_model", "next_probs", "answer", "closed_book", "read", "BUCKETS", "bucket", "test_queries", "by_bucket", "bucket_row", "bucket_header", "knn_datastore", "knn_probs", "knn_neighbors", "query_states", "R", "CKPT_DIR", "N_ENT", "REL", "IRR", "lr_lambda"]