【Datawhale AI 夏令营】第四期 基于2B源大模型RAG 的三体问答

【定位】：Datawhale AI 夏令营 第四期 Task3 应用练手demo

【学习手册链接】：https://linklearner.com/activity/14/11/25

【练手材料来源】：《三体》 https://github.com/JessyTsui/awesome_LLM_beginner.git

• 本文不会按照baseline和代码顺序讲解代码，而是以项目实践顺序来讲解，推荐先阅读学习手册后有各个模块的简易理解，再阅读本篇笔记。（偷一个流程图充下场面。）

• 这篇baseline很有特点，没有调用什么langchain框架，文档处理那些步骤几乎都是很原始的。

文章目录

• • [1. 数据准备与预处理](#1. 数据准备与预处理)
  • [2. 构建向量库索引](#2. 构建向量库索引)
  • • [2.1 显存爆炸 & 代码解读](#2.1 显存爆炸 & 代码解读)
    • [2.2 引入批处理机制 ≈ chunk_size](#2.2 引入批处理机制 ≈ chunk_size)
  • [3. 实现效果](#3. 实现效果)
  • [4. 其他优化](#4. 其他优化)
  • • [4.1 提高检索文档数量](#4.1 提高检索文档数量)
    • [4.2 文档预处理（分 chunk_size 提高单text长度）](#4.2 文档预处理（分 chunk_size 提高单text长度）)
    • [4.3 更新模型](#4.3 更新模型)
  • [5. 最终效果](#5. 最终效果)
  • • [5.1 Question: 冯·诺伊曼为啥要对秦始皇请求三个门部件却不要更多？](#5.1 Question: 冯·诺伊曼为啥要对秦始皇请求三个门部件却不要更多？)
    • [5.2 Question: 介绍一下汪淼](#5.2 Question: 介绍一下汪淼)

1. 数据准备与预处理

当我们顺利跑通baseline之后，由于baseline的测试knowledge仅为一个三行文本。

我们是否可以直接基于baseline的代码应用于其他材料？

碰巧同组队友推荐了上述练手材料，故下载仓库

# 切换到Task3位置，再新建终端
git clone https://github.com/JessyTsui/awesome_LLM_beginner.git

完成后直接修改文档地址：

我们可以看到此时的文件编码不对，是因为再构建索引的时候是指定utf-8读取的。

解决方案：

• 方案1：确认源文档的编码方式后，在向量库索引构建时指定对应的编码方式，不过遇到其他的编码方式可能得再次修改代码。
  
• 方案2：靠代码去读取并且转为utf-8编码。我们新增一个格式转化的代码块进行处理。

import codecs

file_list = ["./awesome_LLM_beginner/resource/三体1疯狂年代.txt",
            "./awesome_LLM_beginner/resource/三体2黑暗森林.txt",
             "./awesome_LLM_beginner/resource/三体3死神永生.txt",
]

def convert_to_utf8(file_list):
    for file_path in file_list:
        try:
            # 尝试以 utf-8 编码读取文件
            with open(file_path, 'r', encoding='utf-8') as file:
                content = file.read()
            # 如果读取成功，则以 utf-8 编码写入同一个文件
            with open(file_path, 'w', encoding='utf-8') as file:
                file.write(content)
            print(f"文件 '{file_path}' 已经是 UTF-8 编码。")
        
        except UnicodeDecodeError:
            # 如果出现 UnicodeDecodeError，尝试使用其他常见编码读取
            common_encodings = ['gbk', 'iso-8859-1', 'windows-1252']
            for encoding in common_encodings:
                try:
                    with open(file_path, 'r', encoding=encoding) as file:
                        content = file.read()
                    # 如果读取成功，则以 utf-8 编码写入文件
                    with open(file_path, 'w', encoding='utf-8') as file:
                        file.write(content)
                    print(f"文件 '{file_path}' 已从 {encoding} 转换为 UTF-8 编码。")
                    break
                except UnicodeDecodeError:
                    continue
            else:
                print(f"无法将文件 '{file_path}' 转换为 UTF-8 编码。没有找到合适的编码。")

# 运行函数
convert_to_utf8(file_list) 

执行后可见输出文件 xxx 已从 gbk 转换为 UTF-8 编码。

2. 构建向量库索引

2.1 显存爆炸 & 代码解读

成功解决编码问题后，我们暂时以第一本小说进行测试。构建时发现爆显存了。

进入类声明中精读一下初始化代码，逐层往下定位问题。

该__init__ 方法: 逐行读取文档文件，并将每行文本添加到 self.documents 列表中（纯CPU+内存，可以中间print一下验证）。

随后使用 embed_model 计算所有文档的嵌入向量，并存储在 self.vectors 中。

我们继续观察向量模型是如何对文档列表get_embeddings的：

在上述代码中，一次性将所有文本列表 传递给 tokenizer 进行编码，再传递给模型进行计算。这种操作虽然看起来简单直接，但在实际应用中存在显著的风险。

当输入的文本列表非常庞大时，一次性处理所有文本会占用大量显存，可能导致显存不足甚至溢出。这是因为所有文本的编码和模型计算都在一次操作中进行，特别是对长文本或文本数量巨大的情况下，显存的需求会成倍增加。

2.2 引入批处理机制 ≈ chunk_size

我们在这里直接在类上 引入 批处理机制 ，可以有效缓解上述风险，提高模型计算的稳定性和效率：

重点关注下面引进的 batch_size

# 定义向量模型类
class EmbeddingModel:
    """
    class for EmbeddingModel
    """

    def __init__(self, path: str, batch_size: int = 32) -> None:
        """
        初始化向量模型类
        Args:
            path (str): 预训练模型的路径
            batch_size (int): 批处理大小，默认值为32
        """
        # 加载预训练的tokenizer
        self.tokenizer = AutoTokenizer.from_pretrained(path)
        # 加载预训练的模型并将其移到GPU上
        self.model = AutoModel.from_pretrained(path).cuda()
        self.batch_size = batch_size
        print(f'Loading EmbeddingModel from {path}.')

    def get_embeddings(self, texts: List[str]) -> List[List[float]]:
        """
        计算文本列表的嵌入向量，支持批处理
        Args:
            texts (List[str]): 输入的文本列表
        Returns:
            List[List[float]]: 文本嵌入向量的列表
        """
        embeddings = []  # 用于存储所有文本的嵌入向量
        # 按批处理大小分块处理文本
        for i in range(0, len(texts), self.batch_size):
            texts_chunk = texts[i:i + self.batch_size]  # 获取当前批次的文本
            # 对文本进行编码，添加填充和截断，并返回PyTorch张量
            encoded_input = self.tokenizer(
                texts_chunk, padding=True, truncation=True, return_tensors='pt')
            encoded_input = {k: v.cuda()
                             for k, v in encoded_input.items()}  # 将张量移到GPU上
            with torch.no_grad():  # 禁用梯度计算以节省内存和计算资源
                model_output = self.model(**encoded_input)  # 获取模型输出
                # 提取每个句子的CLS嵌入
                sentence_embeddings = model_output.last_hidden_state[:, 0, :]
            # 对嵌入向量进行归一化
            sentence_embeddings = torch.nn.functional.normalize(
                sentence_embeddings, p=2, dim=1)
            embeddings.extend(sentence_embeddings.tolist()
                              )  # 将当前批次的嵌入向量添加到结果列表中
        return embeddings  # 返回所有文本的嵌入向量列表

原理

其实原理和常规RAG里面的chunk_size差不多，核心思想在于 分块 。但是这里的batch_size设定之后，是按照文档数量（1个text视作1个），结合baseline的文档读取来看，batch_size = 32也就是每32行文本进行一次embedding，而chunk_size通常为300，指的是每隔300个字符分块embedding。

3. 实现效果

这下，我们终于初步实现了一个三体RAG，但是由于只能检索片段，所以其实全局性的摘要问题一般是很难回答的。

（需要类似GraphRAG 再针对每个文档块做总结提炼，做成社区摘要再embedding检索。）

从效果来看，检索的片段虽然和汪淼有关，但是不够全面。而没有RAG的回答给他扣了一堆帽子，没有一个中的。

带有RAG的回答虽然扯到了三体，但是其实也是回答错了。

（当然，我们这本材料本身就是人物对话为主的小说类，而非 汪淼人物传记，不会特意介绍他。）

4. 其他优化

4.1 提高检索文档数量

我们将目光转向 VectorStoreIndex的查询函数，其中[::-1]表示按照相似度逆序排序，我们取前k个最相似的文档

所以我们在调用query获取上下文片段的时候，可以增大k值，而不是使用默认的1。

但是结果一看，惊掉下巴 Σ(っ °Д °;)っ

按照初始的每行分隔documents，检索的'文档'就只能是无序的单句对话了。

4.2 文档预处理（分 chunk_size 提高单text长度）

由于单个文档text长度太短，就算检索到了也是缺乏上下文的，所以我们在构建索引阶段引入一个max_chunk_size。

class VectorStoreIndex:
    """
    class for VectorStoreIndex
    用于创建和管理基于向量存储的索引
    """

    def __init__(self, document_path: str, embed_model: EmbeddingModel, max_chunk_size=1200) -> None:
        """
        初始化向量库索引类
        Args:
            document_path (str): 文档路径，包含每行一个文档的文本文件
            embed_model (EmbeddingModel): 嵌入模型，用于计算文本的向量表示
            max_chunk_size (int): 每个文档块的最大字符数
        """
        self.documents = []
        # 打开文档路径下的文件，按行读取
        for line in open(document_path, 'r', encoding='utf-8'):
            # 去除行尾的空格和换行符
            line = line.strip()
            # 如果文档列表不为空，并且最后一个文档块加上当前行的长度小于最大块大小
            if len(self.documents) > 0 and len(self.documents[-1]) + len(line) < max_chunk_size:
                # 将当前行追加到最后一个文档块，并在行首添加换行符
                self.documents[-1] += ('\n' + line)
            else:
                # 否则，将当前行作为一个新的文档块添加到文档列表
                self.documents.append(line)

        print('load over')
        # 将嵌入模型赋值给实例变量
        self.embed_model = embed_model
        # 使用嵌入模型获取文档的向量表示，并赋值给实例变量
        self.vectors = self.embed_model.get_embeddings(self.documents)

        print(f'Loading {len(self.documents)} documents for {document_path}.')

可以观测输出，发现此时整篇小说切割成了169个文档块。

4.3 更新模型

源大模型上新(其实就是3月模型换成7月的)：

https://modelscope.cn/models/IEITYuan/Yuan2-2B-July-hf

# 源大模型下载 
from modelscope import snapshot_download
# model_dir = snapshot_download('IEITYuan/Yuan2-2B-Mars-hf', cache_dir='.')
model_dir = snapshot_download('IEITYuan/Yuan2-2B-July-hf', cache_dir='.')

后续构造LLM的时候也替换一下

model_path = './IEITYuan/Yuan2-2B-July-hf'
llm = LLM(model_path)

顺带像上篇blog提到的那样，把大模型生成的max_length 改为 max_new_tokens 防止输入过长。

5. 最终效果

最后设置检索2个文档块（怕输入过多变傻了）。

我们开始提问：

5.1 Question: 冯·诺伊曼为啥要对秦始皇请求三个门部件却不要更多？

检索的context如下（2段 text）：

Context: ['"伟大的陛下，您刚才提到东方人在科学思维上的缺陷，就是因为你们没有意识到，复杂的宇宙万物其实是由最简单的单元构成的。我只要三个，陛下。"\n秦始皇挥手召来了三名士兵，他们都很年轻，与秦国的其他士兵一样，一举一动像听从命令的机器。\n"我不知道你们的名字，"冯·诺伊曼拍拍前两个士兵的肩，"你们两个负责信号输入，就叫'入1'、'入2 '吧，"他又指指最后一名士兵， "你，负责信号输出，就叫'出'吧。"他伸手拨动三名士兵，"这样，站成一个三角形，出是顶端，入1 和入2 是底边。"\n"哼，你让他们成楔形攻击队形不就行了？"秦始皇轻蔑地看着冯·诺伊曼。\n牛顿不知从什么地方掏出六面小旗，三白三黑，冯·诺伊曼接过来分给三名士兵，每人一白一黑，说："白色代表0，黑色代表1。好，现在听我说，出，你转身看着入1 和入2，如果他们都举黑旗，你就举黑旗，其他的情况你都举白旗，这种情况有三种：入l 白，入2 黑；入l 黑，入2白；入1、入2 都是白。"\n"我觉得你应该换种颜色，白旗代表投降。"秦始皇说。\n兴奋中的冯·诺伊曼没有理睬皇帝，对三名士兵大声命令："现在开始运行！入1 入2，你们每人随意举旗，好，举！好，再举！举！"\n入1 和入2同时举了三次旗，第一次是黑黑，第二次是白黑，第三次是黑白。出都进行了正确反应，分别举起了一次黑和两次白。\n"很好，运行正确，陛下，您的士兵很聪明！"\n"这事儿傻瓜都会，你能告诉联，他们在干什么吗？"秦始皇一脸困惑地问。\n"这三个人组成了一个计算系统的部件，是门部件的一种，叫'与门'。"冯·诺伊曼说完停了一会儿，好让皇帝理解。\n秦始皇面无表情地说："朕是够郁闷的，好，继续。"\n冯·诺伊曼转向排成三角阵的三名士兵："我们构建下一个部件。你，出，只要看到入1 和入2中有一个人举黑旗，你就举黑旗，这种情况有三种组合------黑黑、白黑、黑白，剩下的一种情况------ 白白，你就举白旗。明白了吗？好孩子，你真聪明，门部件的正确运行你是关键，好好干，皇帝会奖赏你的！下面开始运行：举！好，再举！再举！好极了，运行正常，陛下，这个门部件叫或门。"\n然后，冯·诺伊曼又用三名士兵构建了与非门、或非门、异或门、同或门和三态门，最后只用两名士兵构建了最简单的非门，出总是举与入颜色相反的旗。\n冯·诺伊曼对皇帝鞠躬说："现在，陛下，所有的门部件都已演示完毕，这很简单不是吗？任何三名士兵经过一小时的训练就可以掌握。"\n"他们不需要学更多的东西了吗？"秦始皇问。\n"不需要，我们组建一千万个这样的门部件，再将这些部件组合成一个系统，这个系统就能进行我们所需要的运算，解出那些预测太阳运行的微分方程。这个系统，我们把它叫做......嗯，叫做......"\n"计算机。"汪淼说。\n"啊------好！"冯·诺伊曼对汪淼竖起一根指头，"计算机，这个名字好，整个系统实际上就是一部庞大的机器，是有史以来最复杂的机器！"', '游戏时间加快，三个月过去了。\n秦始皇、牛顿、冯·诺伊曼和汪淼站在金字塔顶部的平台上，这个平台与汪淼和墨子相遇时的很相似，架设着大量的天文观测仪器，其中有一部分是欧洲近代的设备。在他们下方，三千万秦国军队宏伟的方阵铺展在大地上，这是一个边长六公里的正方形。在初升的太阳下，方阵凝固了似的纹丝不动，仿佛一张由三千万个兵马俑构成的巨毯，但飞翔的鸟群误入这巨毯上空时，立刻感到了下方浓重的杀气，鸟群顿时大乱，惊慌混乱地散开或绕行。汪淼在心里算了算，如果全人类站成这样一个方阵，面积也不过是上海浦东大小，比起它表现的力量，这方阵更显示了文明的脆弱。\n"陛下，您的军队真是举世无双，这么短的时间，就完成了如此复杂的训练。"冯·诺伊曼对秦始皇赞叹道。\n"虽然整体上复杂，但每个士兵要做的很简单，比起以前为粉碎马其顿方阵进行的训练来，这算不了什么。"秦始皇按着长剑剑柄说。\n"上帝也保佑，连着两个这样长的恒纪元。"牛顿说。\n"即使是乱纪元，朕的军队也照样训练，以后，他们也会在乱纪元完成你们的计算。"秦始皇骄傲地扫视着方阵说。\n"那么，请陛下发出您伟大的号令吧！"冯·诺伊曼用激动得发颤的声音说。\n秦始皇点点头，一名卫士奔跑过来，握住皇帝的剑柄向后退了几步，抽出了那柄皇帝本人无法抽出的青铜长剑，然后上前跪下将剑呈给皇帝，秦始皇对着长空扬起长剑，高声喊道：\n"成计算机队列！"\n金字塔四角的四尊青铜大鼎同时轰地燃烧起来，站满了金字塔面向方阵一面坡墙的士兵用宏大的合唱将始皇帝的号令传诵下去：\n"成计算机队列------"\n下面的大地上，方阵均匀的色彩开始出现扰动，复杂精细的回路结构浮现出来，并渐渐充满了整个方阵，十分钟后，大地上出现了一块三十六平方公里的计算机主板。\n冯·诺伊曼指着下方巨大的人列回路开始介绍："陛下，我们把这台计算机命名为'秦一号'。请看，那里，中心部分，是CPU，是计算机的核心计算元件，由您最精锐的五个军团构成，对照这张图您可以看到里面的加法器、寄存器、堆栈存贮器；外围整齐的部分是内存，构建这部分时我们发现人手不够，好在这部分每个单元的动作最简单，就训练每个士兵拿多种颜色的旗帜，组合起来后，一个人就能同时完成最初二十个人的操作，这就使内存容量达到了运行'秦 1.0'操作系统的最低要求；你再看那条贯穿整个阵列的通道，还有那些在通道上待命的轻骑兵，那是BUS，系统总线，负责在整个系统间传递信息。"']

回答不错。

5.2 Question: 介绍一下汪淼

重返刚才的死亡问题。

Context: ['汪淼摇头笑了起来，"得承认今天我的理解力太差了，您这岂不是说......"\n"是的，整个人类历史也是偶然，从石器时代到今天，都没什么重大变故，真幸运。但既然是幸运，总有结束的一天；现在我告诉你，结束了，做好思想准备吧。"\n汪淼还想问下去，但将军与他握手告别，阻止了他下面的问题。\n上车后，司机开口问汪淼家的地址，汪淼告诉他后，随口问道： "哦，接我来的不是你？我看车是一样的。"\n"不是我，我是去接丁博士的。"\n汪淼心里一动，便向司机打听丁仪的住处，司机告诉了他。当天晚上，他就去找丁仪。\n2.台球\n推开丁仪那套崭新的三居室的房门，汪淼闻到了一股酒味，看到丁仪躺在沙发上，电视开着，他的双眼却望着天花板。汪淼四下打量了一下，看到房间还没怎么装修，也没什么家具和陈设，宽大的客厅显得很空，最显眼的是客厅一角摆放的一张台球桌。\n对汪淼的不请自来，丁仪倒没表示反感，他显然也想找人说话。\n"这套房子是三个月前买的，"丁仪说，"我买房子干什么？难道她真的会走进家庭？"他带着醉意笑着摇摇头。\n"你们......"汪淼想知道杨冬生活中的一切，但又不知该如何问。"她像一颗星星，总是那么遥远，照到我身上的光也总是冷的。"丁仪走到窗前看着夜空，像在寻找那颗已逝去的星辰。\n汪淼也沉默下来。很奇怪，他现在就是想听一听她的声音，一年前那个夕阳西下的时刻，她同他对视的那一瞬间没有说话，他从来没有听到过她的声音。\n丁仪一挥手，像要赶走什么，将自己从这哀婉的思绪中解脱出来。"汪教授，你是对的，别跟军方和警方纠缠到一块儿，那是一群自以为是的白痴。那些物理学家的自杀与'科学边界'没有关系，我对他们解释过，可解释不清。"\n"他们好像也做过一些调查。"\n"是，而且这种调查还是全球范围的，那他们也应该知道，其中的两人与'科学边界'没有任何来往，包括------杨冬。"丁仪说出这个名字时显得很吃力。"丁仪，你知道，我现在也卷进这件事里了。所以，关于使杨冬做出这种选择的原因，我很想知道，我想你一定知道一些。"汪淼笨拙地说道，试图掩盖他真正的心迹。\n"如果知道了，你只会卷得更深。现在你只是人和事卷进来了，知道后连精神也会卷进来，那麻烦就大了。"\n"我是搞应用研究的，没有你们理论派那么敏感。"\n"那好吧，打过台球吗？"丁仪走到了台球桌前。\n"上学时随便玩过几下。"\n"我和她很喜欢打，因为这让我们想到了加速器中的粒子碰撞。"丁仪说着拿起黑白两个球，将黑球放到洞旁，将白球放到距黑球仅十厘米左右的位置，问汪淼，"能把黑球打进去吗？"\n"这么近谁都能。"\n"试试。"\n汪淼拿球杆，轻击白球，将黑球撞入洞内。\n"很好，来，我们把球桌换个位置。"丁仪招呼一脸迷惑的汪淼，两人抬起沉重的球桌，将它搬到客厅靠窗的一角。放稳后，丁仪从球袋内掏出刚才打进去的黑球，将它放到洞边，又拾起那个白球，再次放到距黑球十厘米左右的地方，"这次还能打进去吗？"\n"当然。"', '"汪教授，看到这份名单，您有什么印象？"常伟思看着汪淼问。\n"我知道其中的三人，都是物理学最前沿的著名学者。"汪淼答道，有些心不在焉，他的目光锁定在最后一个名字上，在他的潜意识中，那两个字的色彩与上面几行字是不同的。怎么会在这里看到她的名字？她怎么了？\n"认识？"大史用一根被烟熏黄的粗指头指着文件上的那个名字问，见汪淼没有反应，他迅速做出反应，道："呵，不太认识。想认识？"\n现在，汪淼知道常伟思把他以前的这个战士调来是有道理的，这个外表粗俗的家伙，眼睛跟刀子一样。他也许不是个好警察，但确实是个狠角色。\n那是一年前，汪淼是"中华二号"高能加速器项目纳米构件部分的负责人。那天下午在良湘的工地上，一次短暂的休息中，他突然被眼前的一幅构图吸引了。作为一名风景摄影爱好者，现实的场景经常在他眼中形成一幅幅艺术构图。构图的主体就是他们正在安装的超导线圈，那线圈有三层楼高，安装到一半，看上去是一个由巨大的金属块和乱麻般的超低温制冷剂管道组成的怪物，仿佛一堆大工业时代的垃圾，显示出一种非人性的技术的冷酷和钢铁的野蛮。就在这金属巨怪前面，出现了一个年轻女性纤细的身影。这构图的光线分布也很绝：金属巨怪淹没在临时施工顶棚的阴影里，更透出那冷峻、粗糙的质感；而一束夕阳金色的光，透过顶棚的孔洞正好投在那个身影上，柔和的暖光照着她那柔顺的头发，照着工作服领口上白皙的脖颈，看上去就像一场狂暴的雷雨后，巨大的金属废墟上开出了一朵娇柔的花......\n"看什么看，干活儿！"\n汪淼吓了一跳，然后发现纳米研究中心主任说的不是他，而是一名年轻工程师，后者也和自己一样呆呆地望着那个身影。汪淼从艺术中回到现实，发现那位女性不是一般的工作人员，因为总工程师陪同着她，在向她介绍着什么，一副很尊敬的样子。\n"她是谁？"汪淼问主任。\n"你应该知道她的，"主任说，用手划了一大圈，"这个投资二百亿的加速器建成后，第一次运行的可能就是验证她提出的一个超弦模型。要说在论资排辈的理论研究圈子，本来轮不到她的，可那些老家伙不敢先来，怕丢人，就让她捡了个便宜。"\n"什么？杨冬是......女的？！"\n"是的，我们也是在前天见到她时才知道。"主任说。\n那名工程师问："她这人是不是有什么心理障碍，要不怎么会从来不上媒体呢？别像是钱钟书似的，到死大家也没能在电视上看上一眼。"\n"可我们也不至于不知道钱钟书的性别吧？我觉得她童年一定有什么不寻常的经历，以致得了自闭症。"汪淼说，多少有一些酸葡萄心理。杨冬和总工程师走过来，在经过时她对他们微笑着点点头，没说一句话，但汪淼记住了她那清澈的眼睛。']

原本以为回答不出来的，结果！

不算特别好，但是有起色了。