基于 Langchain 和 FAISS 构建 LLM 知识库

在信息检索和自然语言处理的领域，将文档向量化存储并结合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)