在信息检索和自然语言处理的领域，将文档向量化存储并结合LLMs可以显著提升检索效率和答案质量。
最近需要构建一个医疗相关的知识库来增强LLMs的表现，于是在自己搭建完成之后写了这篇文章，介绍如何使用 Langchain 和 FAISS 来一步一步构建一个（相对）高效的知识库系统。
1. 环境准备
首先需要安装一些必要的python库：
pip install langchain transformers faiss-cpu
这其中：
langchain：用于构建智能化的语言模型应用的框架，支持文档加载、向量存储、LLM交互等功能。transformers：HuggingFace 提供的用于加载预训练模型的库。faiss-cpu：FAISS 是 Meta 开发的高效相似性搜索库，用于从大规模数据中进行检索。
2. 加载文档
我们可以将知识库中的文档存储在本地文件夹中，这里用的是 .txt 文件格式。以下代码将文档从给定文件夹中加载进来：
import os
from langchain.schema import Document# Step 1: 从文件夹加载文档
def load_documents_from_folder(folder_path):"""从给定文件夹加载所有txt文件，并将文件名作为title，内容作为content"""documents = []for filename in os.listdir(folder_path):if filename.endswith(".txt"): # 只处理.txt文件file_path = os.path.join(folder_path, filename)# 读取文件内容with open(file_path, "r", encoding="utf-8") as file:content = file.read()# 使用文件名去除扩展名作为文档标题title = os.path.splitext(filename)[0]# 创建 Document 对象保存每个文件的标题和内容doc = Document(page_content=content, metadata={"title": title})documents.append(doc)return documents
该函数遍历指定文件夹，将所有 .txt 文件的内容作为 page_content，文件名作为 title，并使用 Langchain 的 Document 对象存储每篇文档。
3. 使用 HuggingFace 模型生成嵌入
文档加载完成后，我们需要使用 HuggingFace 的预训练模型对文档进行向量化。这里选择 all-MiniLM-L6-v2 模型，实际上随便用一个效果相对好一点的Embedding模型都可以的。
from langchain.embeddings.huggingface import HuggingFaceEmbeddings# Step 2: 创建嵌入模型并生成向量
embeddings = HuggingFaceEmbeddings(model_name="all-MiniLM-L6-v2", cache_folder='./hf-all-MiniLM-L6-v2')
4. 使用 FAISS 构建向量存储库
接下来，我们使用 FAISS 来创建向量存储库，并将生成的文档向量存储其中。
from langchain.vectorstores import FAISS# 使用 FAISS 生成向量存储库
vectorstore = FAISS.from_documents(documents, embeddings)# Step 3: 保存向量存储库
vectorstore.save_local("./faiss_index")
print("向量存储已保存到本地")
FAISS.from_documents() 方法将文档转化为向量，并构建向量存储库。save_local() 方法将生成的索引文件保存到本地，以便之后可以重新加载。
5. 加载 LLM 模型
在完成文档向量化并存储后，我们可以加载 ChatGLM3 模型(当然也可以是其他LLMs)，用于处理用户的自然语言查询，并根据检索到的文档生成答案。
from transformers import AutoTokenizer, AutoModel# Step 4: 加载 ChatGLM 模型
tokenizer = AutoTokenizer.from_pretrained("checkpoints", trust_remote_code=True)
model = AutoModel.from_pretrained("checkpoints", device='cuda', trust_remote_code=True)
model = model.eval()
6. 查询并检索文档
一旦向量存储库和 LLM 模型都准备就绪，用户可以输入问题，系统将通过 FAISS 检索最相关的文档，并使用检索到的文档作为上下文，由 ChatGLM3 生成答案。
# Step 5: 输入查询并检索相关文档
query = "气虚体质有哪些表现？"
retrieved_docs = vectorstore.similarity_search(query, k=1) # 返回最相关的1个文档，k代表返回相关文档的数量# 打印检索到的文档
print(f'检索到的相关文档：')
for doc in retrieved_docs:print(f"标题: {doc.metadata['title']}\n内容: {doc.page_content}\n")# 将检索到的文档作为上下文进行模型生成
context = "\n".join([doc.page_content for doc in retrieved_docs])
prompt = f"根据以下内容回答问题：\n\n{context}\n\n问题: {query}"
在这里，vectorstore.similarity_search() 函数用于进行相似性搜索，返回最相关的文档列表，随后将这些文档的内容拼接为上下文，输入到GLM3 中。
7. 生成回答
最后，利用检索到的文档上下文和用户查询，通过 ChatGLM3 生成最终的答案。
# 生成答案，使用 model.chat()
response, history = model.chat(tokenizer, prompt, history=[], temperature=0.1)print("模型生成的回答：", response)