""" ============================================================================= 第 12 讲 · 工具模块：一个虚构世界的知识库（文档 + 问答对） ============================================================================= 对应讲义：lectures/lecture12_agents_tools_rag.md 第 2 节 被本目录的所有代码使用。 ----------------------------------------------------------------------------- 【为什么要虚构一个世界？】 ----------------------------------------------------------------------------- 检索增强生成（RAG）要解决的问题是：模型**不知道**某些知识（太新、太冷门、私有），需要去查。 如果用真实世界的知识（"法国的首都是哪里？"），我们分不清模型是"查到的"还是"本来就记得"。 所以这里随机生成一个虚构的世界：几十个虚构的国家，各有首都、语言、人口、总统…… 这些"知识"只存在于我们生成的文档里，任何模型都不可能事先知道。 每个国家有 4 篇短文档（一共 4 × 国家数 篇）。主世界的第一个国家（名字都是随机拼出来的）： 概况 norzanul is a country in the lowlands . its capital is maryelvos . the official language of norzanul is karpel . the weather there is mild . the country was founded in 1737 . 人口 about 88 million people live in norzanul . many of them work on farms . norzanul is led by president munul yeltal . the main export of the country is cotton . 城市 maryelvos lies on the river xarsor . the city is famous for its salt museum . 人物 munul yeltal was born in 1975 . before politics , yeltal worked as a doctor . 概况、人口、城市三篇各有一半的概率多一句无关的话（如 "the weather there is mild ."），让文档长短不一。 注意"城市"和"人物"两篇**不提国家的名字**：问"norzanul 的首都边上是哪条河？"必须先查到首都是 maryelvos，再去查 maryelvos —— 这就是**多跳**（multi-hop）问题。 问题有三种风格： direct 问题和答案所在的文档用词相同："what is the capital of norzanul ?" para 换一种说法（同义改写），和文档几乎没有共同的词："where does the government of norzanul sit ?" → 只靠词面匹配的检索（TF-IDF、BM25）会在这里吃亏（02 号文件） multi 需要两跳："which river does the capital of norzanul lie on ?" 另外还有需要计算的问题（04 号文件的智能体用计算器）： "how many years older is the president of A than the president of B ?" 【用法】 from toy_kb import make_world, Vocab kb = make_world(40, seed=0) # 主世界：评估用 train_world = make_world(300, seed=1, exclude=kb.names) # 另一个世界：训练阅读器、稠密检索器用 kb.docs[0] # {"id": 0, "tokens": [...], "kind": "overview", "entity": "norzanul", ...} kb.questions[0] # {"q": [...], "answer": [...], "doc": 0, "relation": "capital", ...} 直接运行本文件会打印一些例子。 """ import random # ============================================================================= # 第 1 部分：词表素材 # ============================================================================= SYLLABLES = ["vel", "mor", "zan", "dri", "quor", "via", "tal", "ben", "sor", "kel", "dra", "nim", "oth", "pel", "ras", "tun", "ul", "wen", "xar", "yel", "gor", "hal", "ist", "jun", "lom", "mar", "nor", "pra", "sel", "ter", "vos", "bri", "cal", "dun", "fen", "gil", "hov", "kar", "lis", "mun"] REGIONS = ["north", "south", "east", "west", "mountains", "lowlands"] EXPORTS = ["copper", "wool", "tea", "salt", "timber", "cotton", "coffee", "glass", "silk", "rice", "iron", "fish"] LANDMARKS = [["glass", "bridge"], ["stone", "tower"], ["old", "harbor"], ["night", "market"], ["red", "cathedral"], ["river", "gardens"], ["clock", "tower"], ["salt", "museum"], ["royal", "library"], ["opera", "house"]] JOBS = ["teacher", "doctor", "farmer", "lawyer", "engineer", "pilot", "journalist", "banker", "nurse", "chef"] # 关系 → (direct 问题模板, para 问题模板, 实体类型)。{e} 会被替换成实体名（可能是多个词） RELATIONS = { "capital": ("what is the capital of {e} ?", "where does the government of {e} sit ?", "country"), "language": ("what is the official language of {e} ?", "what tongue do people in {e} speak ?", "country"), "founded": ("when was {e} founded ?", "in which year did {e} become a nation ?", "country"), "population": ("how many people live in {e} ?", "what is the population of {e} ?", "country"), "leader": ("who is the president of {e} ?", "who governs {e} ?", "country"), "export": ("what is the main export of {e} ?", "what does {e} mostly sell abroad ?", "country"), "river": ("which river does {e} lie on ?", "what waterway flows through {e} ?", "city"), "landmark": ("what is {e} famous for ?", "which sight should a visitor to {e} see ?", "city"), "born": ("when was {e} born ?", "in what year did {e} enter the world ?", "person"), "job": ("what did {e} work as before politics ?", "what was the earlier career of {e} ?", "person"), } # 多跳问题：(模板, 第一跳的关系, 第二跳的关系) MULTI_HOP = { "capital→river": ("which river does the capital of {e} lie on ?", "capital", "river"), "capital→landmark": ("what is the capital of {e} famous for ?", "capital", "landmark"), "leader→born": ("when was the president of {e} born ?", "leader", "born"), "leader→job": ("what did the president of {e} work as before politics ?", "leader", "job"), } # ============================================================================= # 第 2 部分：生成世界 # ============================================================================= class World: """ 一个虚构世界：国家、文档、问题。 【属性】 countries : list[dict]，每个国家的所有属性 docs : list[dict]，{"id", "tokens", "kind", "entity", "country"} questions : list[dict]，单跳问题（direct 和 para 两种风格） multi : list[dict]，多跳问题 names : set[str]，世界里用到的所有名字（生成另一个世界时用来避免重名） """ def __init__(self, countries): self.countries = countries self.docs, self.questions, self.multi = [], [], [] self.names = set() self._build() def _add_doc(self, text, kind, entity, country): self.docs.append({"id": len(self.docs), "tokens": text.split(), "kind": kind, "entity": entity, "country": country}) return len(self.docs) - 1 def _build(self): for c in self.countries: n, cap, ldr = c["name"], c["capital"], c["leader"] self.names |= {n, cap, c["language"], c["river"]} | set(ldr.split()) f = c["fillers"] # 随机加上的"无关句子"：让文档长短不一，稀疏检索不能靠文档长度"蒙对" d_over = self._add_doc( f"{n} is a country in the {c['region']} . its capital is {cap} . the official language of {n} " f"is {c['language']} . " + ("the weather there is mild . " if f[0] else "") + f"the country was founded in {c['founded']} .", "overview", n, n) d_people = self._add_doc( f"about {c['population']} million people live in {n} . " + ("many of them work on farms . " if f[1] else "") + f"{n} is led by president {ldr} . the main export of the country is {c['export']} .", "people", n, n) d_city = self._add_doc( f"{cap} lies on the river {c['river']} . " + ("its streets are narrow . " if f[2] else "") + f"the city is famous for its {' '.join(c['landmark'])} .", "city", cap, n) last = ldr.split()[1] d_person = self._add_doc( f"{ldr} was born in {c['born']} . before politics , {last} worked as a {c['job']} .", "person", ldr, n) c["docs"] = {"overview": d_over, "people": d_people, "city": d_city, "person": d_person} answers = { "capital": ([cap], d_over), "language": ([c["language"]], d_over), "founded": ([str(c["founded"])], d_over), "population": ([str(c["population"]), "million"], d_people), "leader": (ldr.split(), d_people), "export": ([c["export"]], d_people), "river": ([c["river"]], d_city), "landmark": (list(c["landmark"]), d_city), "born": ([str(c["born"])], d_person), "job": ([c["job"]], d_person), } c["answers"] = answers entity_of = {"country": n, "city": cap, "person": ldr} for rel, (t_direct, t_para, etype) in RELATIONS.items(): ans, doc = answers[rel] for style, tmpl in [("direct", t_direct), ("para", t_para)]: self.questions.append({"q": tmpl.format(e=entity_of[etype]).split(), "answer": ans, "doc": doc, "relation": rel, "style": style, "country": n}) for name, (tmpl, r1, r2) in MULTI_HOP.items(): self.multi.append({"q": tmpl.format(e=n).split(), "answer": answers[r2][0], "bridge": answers[r1][0], "docs": [answers[r1][1], answers[r2][1]], "relation": name, "first": r1, "second": r2, "country": n}) def country(self, name): """按名字找国家（找不到返回 None）。""" return next((c for c in self.countries if c["name"] == name), None) def _make_name(rng, used, n_syll=(2, 3)): """用音节拼一个没用过的名字，比如 norzanul、maryelvos。""" for _ in range(1000): name = "".join(rng.choice(SYLLABLES) for _ in range(rng.randint(*n_syll))) if name not in used and len(name) >= 5: used.add(name) return name raise RuntimeError("名字用完了：请增加音节或减少国家数") def make_world(n_countries, seed=0, exclude=()): """ 生成一个有 n_countries 个国家的虚构世界。 【参数】 n_countries : 国家数（每个国家 4 篇文档、20 个单跳问题、4 个多跳问题） seed : 随机种子 exclude : 不允许使用的名字（生成"训练世界"时传入主世界的 names，保证两个世界没有重名） 【返回】World 【边界】所有名字（国家、城市、语言、河流、人名）互不相同；同一个世界里年份、人口可能重复。 """ rng = random.Random(seed) used = set(exclude) countries = [] for _ in range(n_countries): countries.append({ "name": _make_name(rng, used), "capital": _make_name(rng, used), "language": _make_name(rng, used, (2, 2)), "river": _make_name(rng, used, (2, 2)), "leader": f"{_make_name(rng, used, (2, 2))} {_make_name(rng, used, (2, 2))}", "region": rng.choice(REGIONS), "founded": rng.randint(1700, 1990), "population": rng.randint(1, 90), "export": rng.choice(EXPORTS), "landmark": rng.choice(LANDMARKS), "born": rng.randint(1940, 1985), "job": rng.choice(JOBS), "fillers": [rng.random() < 0.5 for _ in range(3)], }) return World(countries) def arithmetic_questions(world, n, seed=0): """ 需要计算的问题（04 号文件的智能体用计算器回答）： "how many years older is the president of A than the president of B ?" → born(B) − born(A) "what is the combined population of A and B in millions ?" → pop(A) + pop(B) 【返回】list[dict]：{"q", "answer"（字符串）, "kind", "countries"} 【边界】年龄差的问题只挑 A 的总统确实比 B 的总统年长的组合（答案为正数）。 """ rng = random.Random(seed) out = [] while len(out) < n: a, b = rng.sample(world.countries, 2) if rng.random() < 0.5: if a["born"] >= b["born"]: continue q = f"how many years older is the president of {a['name']} than the president of {b['name']} ?" out.append({"q": q.split(), "answer": str(b["born"] - a["born"]), "kind": "age_gap", "countries": [a["name"], b["name"]]}) else: q = f"what is the combined population of {a['name']} and {b['name']} in millions ?" out.append({"q": q.split(), "answer": str(a["population"] + b["population"]), "kind": "pop_sum", "countries": [a["name"], b["name"]]}) return out # ============================================================================= # 第 3 部分：词表 # ============================================================================= class Vocab: """ 词 ↔ 编号。特殊符号： 0、 1（问题和文档之间的分隔符）、 2。 【用法】vocab = Vocab([句子1的词列表, 句子2的词列表, ...]) 【边界】为了让阅读器能处理任何世界的文档，建词表时应把所有会遇到的文本都传进来 （训练世界 + 主世界）。主世界的名字在训练时从没出现过，它们的词向量不会被训练 —— 模型只能靠"问题里的这个词和文档里的那个词是同一个"来定位答案，这正是阅读理解要学的本事。 """ PAD, SEP, UNK = "", "", "" def __init__(self, sentences): words = sorted({w for s in sentences for w in s} - {self.PAD, self.SEP, self.UNK}) self.itos = [self.PAD, self.SEP, self.UNK] + words self.stoi = {w: i for i, w in enumerate(self.itos)} def __len__(self): return len(self.itos) def encode(self, words): """词列表 → 编号列表（不认识的词变成 ）。""" return [self.stoi.get(w, self.stoi[self.UNK]) for w in words] def all_sentences(*worlds): """收集若干个世界里所有文档和问题的词列表（建词表用）。""" out = [] for w in worlds: out += [d["tokens"] for d in w.docs] + [q["q"] for q in w.questions] + [q["q"] for q in w.multi] out += [q["q"] for q in arithmetic_questions(w, 50)] return out if __name__ == "__main__": kb = make_world(40, seed=0) print(f"主世界：{len(kb.countries)} 个国家，{len(kb.docs)} 篇文档，{len(kb.questions)} 个单跳问题，" f"{len(kb.multi)} 个多跳问题\n") c = kb.countries[0] for kind in ["overview", "people", "city", "person"]: print(f" [{kind:<8s}] {' '.join(kb.docs[c['docs'][kind]]['tokens'])}") print() for q in kb.questions[:20:3]: print(f" ({q['style']:<6s}) {' '.join(q['q']):<55s} → {' '.join(q['answer'])}") m = kb.multi[0] print(f" (multi ) {' '.join(m['q']):<55s} → {' '.join(m['answer'])}（中间实体：{' '.join(m['bridge'])}）") a = arithmetic_questions(kb, 2, seed=3) for q in a: print(f" (calc ) {' '.join(q['q'])} → {q['answer']}") train_world = make_world(300, seed=1, exclude=kb.names) assert not (train_world.names & kb.names) vocab = Vocab(all_sentences(train_world, kb)) print(f"\n训练世界：{len(train_world.countries)} 个国家（和主世界没有重名）；两个世界合起来的词表 {len(vocab)} 个词")