手搓RAG
RAG(检索增强生成)是结合检索与生成式 AI 的技术框架。核心逻辑是先从外部知识库精准检索相关信息,再将其作为上下文输入大模型生成回答。技术上依赖检索引擎(如向量数据库、BM25)、大语言模型(如 GPT、LLaMA)及数据预处理技术。通过检索增强,解决大模型知识滞后、幻觉问题,提升回答准确性。应用广泛,涵盖智能客服、医疗问答、法律检索、教育辅导等场景,能基于特定领域知识提供精准、可控的生成内容。
wow-RAG 是 Datawhale 推出的 RAG 技术实践项目,网址:datawhalechina/wow-rag: A simple and trans-platform rag framework and tutorial
https://github.com/datawhalechina/wow-rag
目录
安装依赖库
创建并激活虚拟环境
安装 Jupyter Notebook
启动JupyterNotebook服务器
创建新的 Notebook
安装库
获取API
环境变量配置与客户端初始化
文本分割(Chunking)
生成嵌入向量与构建检索索引
文本匹配函数
执行检索与生成答案
流式生成答案函数
参考文章
安装依赖库
创建并激活虚拟环境
在PowerShell里运行:
# 创建虚拟环境(命名为 myenv)
python -m venv myenv# 激活虚拟环境
myenv\Scripts\activate
激活成功后,终端提示符会显示 (myenv) 前缀,表示已进入虚拟环境。
安装 Jupyter Notebook
在激活的虚拟环境中执行:
pip install jupyter notebook
显示:
启动JupyterNotebook服务器
安装完成后,在终端输入:
jupyter notebook
命令执行后,浏览器会自动打开 Jupyter Notebook 的 web 界面(通常是 http://localhost:8888)。终端会保持运行状态,不要关闭,关闭即停止服务。
创建新的 Notebook
在 Jupyter Notebook 的 web 界面中:点击右上角的 New 按钮。在下拉菜单中选择 Python 3(或你安装的 Python 版本)。系统会自动创建并打开一个新的 .ipynb 文件,默认名称为 Untitled.ipynb。
安装库
faiss-cpu 是用于向量相似度搜索的工具。scikit-learn 和 scipy 是数据科学常用的库。openai 是调用 OpenAI API 的官方客户端库。python-dotenv 可用于加载环境变量。
如果在Kaggle里运行,Kaggle 的代码单元采用 IPython 内核,所以安装 Python 包时,应当使用 %pip install 或者 !pip install,而不是把这两个符号组合起来用
%pip install faiss-cpu scikit-learn scipy
%pip install openai
%pip install python-dotenv
可以在Kaggle里选用GPU:
运行以下代码来安装faiss-gpu包:
%pip install faiss-gpu -U
可能会出现如下的bug:
ERROR: Could not find a version that satisfies the requirement faiss-gpu (from versions: none) ERROR: No matching distribution found for faiss-gpu Note: you may need to restart the kernel to use updated packages.
报错No matching distribution found的原因:PyPI 上的 faiss-gpu 包与当前环境的 CUDA 版本不兼容。Kaggle 的 GPU 环境默认使用 Conda 包管理器,且 Faiss 在 Conda 上的支持更为稳定。
可以尝试用conda安装
获取API
首先,在智谱开放平台上注册账号并创建应用,获取 API Key。网址:智谱AI开放平台
Windows 可能会提示 “如果改变文件扩展名,可能会导致文件不可用”,点击 是 确认
在项目根目录创建一个名为 .env 的文件(注意文件名前有个点),打开 .env 文件,添加一行内容:
ZHIPU_API_KEY=你的真实API密钥
将 你的真实API密钥 替换为你从智谱平台获取的实际 API Key。
GLM-4 系列提供了复杂推理、超长上下文、极快推理速度等多款模型,适用于多种应用场景。模型编码:glm-4-plus、glm-4-air-250414、glm-4-airx、glm-4-long 、glm-4-flashx 、glm-4-flash-250414
环境变量配置与客户端初始化
import os
from dotenv import load_dotenv# 加载环境变量(从.env文件读取配置)
load_dotenv()
# 从环境变量中读取智谱API密钥(.env文件中定义的ZHIPU_API_KEY)
api_key = os.getenv('ZHIPU_API_KEY')
# 智谱API的基础URL(接口请求的根路径)
base_url = "https://open.bigmodel.cn/api/paas/v4/"
# 聊天模型名称(智谱的glm-4-flash轻量模型)
chat_model = "glm-4-flash"
# 嵌入模型名称(智谱的embedding-2向量模型)
emb_model = "embedding-2"# 导入OpenAI客户端类(用于对接大模型API)
from openai import OpenAI
# 创建OpenAI客户端实例(适配国内模型)
client = OpenAI(api_key = api_key, # 传入API密钥(身份验证)base_url = base_url # 覆盖默认URL,指向智谱API
)
load_dotenv():读取项目根目录的.env文件,将环境变量加载到系统中,避免硬编码密钥。
os.getenv('ZHIPU_API_KEY'):从环境变量中获取密钥,确保敏感信息不暴露在代码中。
OpenAI()客户端初始化:通过base_url指定智谱 API 地址,实现用 OpenAI 库调用国内模型的兼容。
文本分割(Chunking)
# 定义需要生成嵌入向量的长文本(关于多模态Agent AI的研究内容)
embedding_text = """
Multimodal Agent AI systems have many applications. In addition to interactive AI, grounded multimodal models could help drive content generation for bots and AI agents, and assist in productivity applications, helping to re-play, paraphrase, action prediction or synthesize 3D or 2D scenario. Fundamental advances in agent AI help contribute towards these goals and many would benefit from a greater understanding of how to model embodied and empathetic in a simulate reality or a real world. Arguably many of these applications could have positive benefits.However, this technology could also be used by bad actors. Agent AI systems that generate content can be used to manipulate or deceive people. Therefore, it is very important that this technology is developed in accordance with responsible AI guidelines. For example, explicitly communicating to users that content is generated by an AI system and providing the user with controls in order to customize such a system. It is possible the Agent AI could be used to develop new methods to detect manipulative content - partly because it is rich with hallucination performance of large foundation model - and thus help address another real world problem.For examples, 1) in health topic, ethical deployment of LLM and VLM agents, especially in sensitive domains like healthcare, is paramount. AI agents trained on biased data could potentially worsen health disparities by providing inaccurate diagnoses for underrepresented groups. Moreover, the handling of sensitive patient data by AI agents raises significant privacy and confidentiality concerns. 2) In the gaming industry, AI agents could transform the role of developers, shifting their focus from scripting non-player characters to refining agent learning processes. Similarly, adaptive robotic systems could redefine manufacturing roles, necessitating new skill sets rather than replacing human workers. Navigating these transitions responsibly is vital to minimize potential socio-economic disruptions.Furthermore, the agent AI focuses on learning collaboration policy in simulation and there is some risk if directly applying the policy to the real world due to the distribution shift. Robust testing and continual safety monitoring mechanisms should be put in place to minimize risks of unpredictable behaviors in real-world scenarios. Our “VideoAnalytica" dataset is collected from the Internet and considering which is not a fully representative source, so we already go through-ed the ethical review and legal process from both Microsoft and University Washington. Be that as it may, we also need to understand biases that might exist in this corpus. Data distributions can be characterized in many ways. In this workshop, we have captured how the agent level distribution in our dataset is different from other existing datasets. However, there is much more than could be included in a single dataset or workshop. We would argue that there is a need for more approaches or discussion linked to real tasks or topics and that by making these data or system available.We will dedicate a segment of our project to discussing these ethical issues, exploring potential mitigation strategies, and deploying a responsible multi-modal AI agent. We hope to help more researchers answer these questions together via this paper."""# 设置每个文本块的字符长度(控制 chunk 大小,避免超过模型输入限制)
chunk_size = 150
# 使用列表推导式分割文本:从0开始,每150个字符切分一次,生成多个文本块
chunks = [embedding_text[i:i + chunk_size] for i in range(0, len(embedding_text), chunk_size)]
长文本超过嵌入模型的输入限制(或影响语义准确性),需拆分为较小的chunk。
chunk_size=150:根据模型能力设置(智谱embedding-2支持更长文本,但拆分可提升检索精度)。
range(0, len(embedding_text), chunk_size)生成起始索引,i:i + chunk_size截取子字符串。
生成嵌入向量与构建检索索引
# 导入sklearn的归一化工具、numpy(数值计算)、faiss(向量检索库)
from sklearn.preprocessing import normalize
import numpy as np
import faiss# 初始化空列表存储每个文本块的嵌入向量
embeddings = []# 遍历每个文本块,生成嵌入向量
for chunk in chunks:# 调用智谱嵌入模型API,生成当前文本块的向量response = client.embeddings.create(model=emb_model, # 指定嵌入模型(embedding-2)input=chunk, # 输入文本块)# 提取向量数据(response.data[0].embedding为向量列表),添加到embeddings列表embeddings.append(response.data[0].embedding)# 归一化处理:将向量转换为单位向量(提升余弦相似度计算精度)
# np.array(embeddings).astype('float32'):转换为FAISS支持的float32格式数组
normalized_embeddings = normalize(np.array(embeddings).astype('float32'))# 获取嵌入向量的维度(embedding-2为1536维,embedding-3为2048维)
d = len(embeddings[0])# 创建FAISS索引:IndexFlatIP表示使用内积(Inner Product)计算相似度(等价于归一化后的余弦相似度)
index = faiss.IndexFlatIP(d)# 将归一化后的向量添加到索引中(用于后续快速检索)
index.add(normalized_embeddings)# 获取索引中的向量总数(验证是否所有文本块都已添加)
n_vectors = index.ntotal# 打印向量总数(确认索引构建成功)
print(n_vectors)
for chunk in chunks:response = client.embeddings.create(model=emb_model,input=chunk,)embeddings.append(response.data[0].embedding)
遍历 chunks 文本块列表。使用 client.embeddings.create(...) 调用智谱或其他平台提供的嵌入模型 API。将返回的嵌入向量(response.data[0].embedding)添加到 embeddings 列表中。
normalized_embeddings = normalize(np.array(embeddings).astype('float32'))
将 embeddings 列表转换为 NumPy 数组,并转为 float32 类型(FAISS 所需格式)。使用 normalize() 对向量进行 L2 归一化,使得每个向量变成单位向量。
d = len(embeddings[0])
index = faiss.IndexFlatIP(d)
d 是嵌入向量的维度,例如对于 "embedding-2" 是 1536 维,对于 "embedding-3" 是 2048 维。faiss.IndexFlatIP(d) 创建一个基于内积(Inner Product)的索引结构,适用于归一化后的向量,等效于余弦相似度匹配。
n_vectors = index.ntotal
print(n_vectors)
index.ntotal 表示当前索引中已添加的向量总数。打印该值以确认所有文本块的向量都已正确加入索引。
文本匹配函数
from sklearn.preprocessing import normalize
def match_text(input_text, index, chunks, k=2):"""在文本块中找到与输入文本最相似的前k个文本块。"""# 确保k不超过文本块总数(避免索引越界)k = min(k, len(chunks))# 为输入文本生成嵌入向量response = client.embeddings.create(model=emb_model,input=input_text,)# 提取输入文本的向量input_embedding = response.data[0].embedding# 归一化输入向量(与索引向量格式一致)input_embedding = normalize(np.array([input_embedding]).astype('float32'))# 在FAISS索引中搜索最相似的k个向量# distances:相似度分数(内积值,越大越相似);indices:相似文本块的索引distances, indices = index.search(input_embedding, k)# 初始化字符串存储匹配结果matching_texts = ""# 遍历搜索结果for i, idx in enumerate(indices[0]):# 打印相似度和匹配的文本块(方便调试)print(f"similarity: {distances[0][i]:.4f}\nmatching text: \n{chunks[idx]}\n")# 将结果拼接成字符串返回matching_texts += f"similarity: {distances[0][i]:.4f}\nmatching text: \n{chunks[idx]}\n"return matching_texts
input_text: 待匹配的输入文本;index: FAISS 索引对象,存储了文本块的嵌入向量;chunks: 文本块列表,与索引中的向量一一对应;k: 要返回的最相似文本块数量,默认为 2
k = min(k, len(chunks))
这一行确保 k 不超过文本块总数,避免在索引中搜索时出现越界错误。如果 k 大于文本块数量,则自动调整为最大可用值。
input_embedding = normalize(np.array([input_embedding]).astype('float32'))
将输入向量转换为 numpy 数组并进行 L2 归一化,确保与 FAISS 索引中的向量格式一致。
matching_texts = ""
初始化一个空字符串变量 matching_texts,用于拼接最终要返回的匹配结果信息。
for i, idx in enumerate(indices[0]):print(f"similarity: {distances[0][i]:.4f}\nmatching text: \n{chunks[idx]}\n")matching_texts += f"similarity: {distances[0][i]:.4f}\nmatching text: \n{chunks[idx]}\n"
similarity: {distances[0][i]:.4f} 中,distances 是 FAISS 返回的二维数组,distances[0] 表示输入向量与所有匹配项之间的相似度分数列表,distances[0][i] 是第 i 个匹配项的相似度值,:.4f 表示保留四位小数进行格式化输出。matching text: \n{chunks[idx]}\n 中的 \n 是换行符,用于在输出时增加可读性;chunks[idx] 表示根据索引 idx 取出原始文本块内容。
执行检索与生成答案
# 定义用户查询(需要回答的问题)
input_text = "What are the applications of Agent AI systems ?"# 调用match_text函数,找到最相似的2个文本块
matched_texts = match_text(input_text=input_text, index=index, chunks=chunks, k=2)# 构建提示词:让模型基于检索到的文本生成答案
prompt = f"""
根据找到的文档
{matched_texts}
生成
{input_text}
的答案,尽可能使用文档原文。不要复述问题,直接开始回答。
"""
matched_texts = match_text(input_text=input_text, index=index, chunks=chunks, k=2)
调用之前定义好的 match_text 函数,传入用户的问题、已构建好的 FAISS 索引对象 index、原始文本块列表 chunks 以及返回结果数量 k=2,函数会返回两个最相似文本块及其相似度信息拼接成的字符串,赋值给 matched_texts。
流式生成答案函数
def get_completion_stream(prompt):"""使用流式输出生成模型回复(逐字显示,提升交互体验)。"""# 调用智谱聊天模型API,开启流式输出response = client.chat.completions.create(model=chat_model, # 指定聊天模型(glm-4-flash)messages=[{"role": "user", "content": prompt}, # 用户提示词],stream=True, # 启用流式输出(边生成边返回))# 遍历流式响应的每个chunkif response:for chunk in response:# 提取当前chunk的内容(delta.content为增量文本)content = chunk.choices[0].delta.content# 如果内容存在,立即打印(end=''避免换行,flush=True强制刷新输出)if content:print(content, end='', flush=True)# 注意:此处原代码有递归调用bug,会导致无限循环,实际使用时应删除此行# get_completion_stream(prompt)
response = client.chat.completions.create(model=chat_model,messages=[{"role": "user", "content": prompt},],stream=True,)
调用智谱AI的聊天模型API,创建一个包含用户提示词的消息结构,并启用流式输出(stream=True),使得模型在生成回复时可以逐字返回结果,而不是等待整个回复生成完成后再一次性返回。
if response:for chunk in response:content = chunk.choices[0].delta.contentif content:print(content, end='', flush=True)
对 API 返回的响应进行遍历处理,每个 chunk 表示模型生成的一部分内容。从中提取增量文本 content,如果存在则立即打印到控制台,使用 end='' 避免每次打印自动换行,flush=True 确保内容被立即显示,从而实现类似打字机效果的实时输出。
get_completion_stream(prompt)
输出:
这一行调用了之前定义的 get_completion_stream 函数,并传入了构造好的提示词 prompt。函数内部会向大语言模型(如智谱的 GLM 系列)发起请求,以流式方式逐字接收并打印模型生成的回答内容。由于使用了流式输出(stream=True),回答不会一次性返回,而是一边生成一边显示在终端上,提升交互体验。
参考文章
wow-raghttps://github.com/datawhalechina/wow-rag