这篇文章将介绍可以让你与本地文件交互的开源生成式搜索引擎。Microsoft Copilot 已经提供了类似的东西,但我认为我可以制作一个开源版本,并分享我在快速编写系统代码过程中收集到的一些知识。
系统设计
为了构建本地生成搜索引擎或助手,我们需要几个组件:
下图显示了组件如何交互。
首先,我们需要将本地文件编入索引,以便查询本地文件的内容。然后,当用户提出问题时,我们将使用创建的索引以及一些不对称段落或文档嵌入来检索可能包含答案的最相关文档。这些文档的内容和问题被传递给已部署的大型语言模型,该模型将使用给定文档的内容来生成答案。在指令提示中,我们将要求大型语言模型也返回对所用文档的引用。最终,所有内容都将在用户界面上向用户可视化。
现在,让我们更详细地了解每个组件。
语义索引
我们正在构建一个语义索引,该索引将根据文件内容和给定查询的相似性为我们提供最相关的文档。要创建这样的索引,我们将使用 Qdrant 作为向量存储。有趣的是,Qdrant 客户端库不需要完整安装Qdrant 服务器,并且可以对适合工作内存 (RAM) 的文档进行相似性分析。因此,我们需要做的就是 pip install Qdrant client。
我们可以通过以下方式初始化 Qdrant(请注意,hf 参数稍后会根据故事流程进行定义,但使用 Qdrant 客户端,你已经需要定义正在使用哪种矢量化方法和度量):
from qdrant_client import QdrantClient
from qdrant_client.models import Distance, VectorParams
client = QdrantClient(path="qdrant/")
collection_name = "MyCollection"
if client.collection_exists(collection_name):
client.delete_collection(collection_name)
client.create_collection(collection_name,vectors_config=VectorParams(size=768, distance=Distance.DOT))
qdrant = Qdrant(client, collection_name, hf)
为了创建向量索引,我们必须将文档嵌入硬盘。对于嵌入,我们必须选择正确的嵌入方法和正确的向量比较度量。可以使用几种段落、句子或单词嵌入方法,但效果各不相同。基于文档创建向量搜索的主要问题是非对称搜索问题。非对称搜索问题在信息检索中很常见,当查询较短而文档较长时就会出现。单词或句子嵌入通常会进行微调,以提供基于相似大小文档(句子或段落)的相似性分数。一旦不是这种情况,适当的信息检索就可能失败。
不过,我们可以找到一种能很好地解决非对称搜索问题的嵌入方法。例如,在 MSMARCO 数据集上微调的模型通常效果很好。MSMARCO 数据集基于必应搜索查询和文档,由微软发布。因此,它非常适合我们处理的问题。
对于这个特定的实现,我选择了一个已经微调过的模型,名为 “MSMARCO”:
sentence-transformers/msmarco-bert-base-dot-v5
该模型基于 BERT,并使用点积作为相似性度量进行了微调。我们已将 qdrant 客户端初始化为使用点积作为相似度度量(注意该模型的维度为 768):
client.create_collection(collection_name,vectors_config=VectorParams(size=768, distance=Distance.DOT))768, distance=Distance.DOT))
我们可以使用其他指标,例如余弦相似度,但是,鉴于这个模型是使用点积进行微调的,因此我们将使用这个指标获得最佳性能。此外,我们还可以从几何角度进行思考: 余弦相似度只关注角度的差异,而点积同时考虑了角度和幅度。通过对数据进行归一化处理,使其具有统一的幅度,这两种度量方法就变得等价了。在忽略大小有利的情况下,余弦相似性是有用的。不过,如果幅度很重要,点积则是更合适的相似度量。
初始化 MSMarco 模型的代码如下(如果你有 GPU,请使用它:)
model_name = "sentence-transformers/msmarco-bert-base-dot-v5""sentence-transformers/msmarco-bert-base-dot-v5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
hf = HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
我们需要解决的下一个问题是,由于变换器模型对内存的二次需求,类 BERT 模型的上下文大小有限。在许多类 BERT 模型中,上下文大小被设定为 512 个标记。有两种选择:(1) 我们可以只根据前 512 个词组来回答问题,而忽略文档的其余部分;或者 (2) 创建一个索引,将一个文档分割成多个词组,并以词组的形式存储在索引中。在第一种情况下,我们会丢失很多重要信息,因此我们选择了第二种变体。要对文档进行分块,我们可以使用 LangChain 预建的分块器:
from langchain_text_splitters import TokenTextSplitter
text_splitter = TokenTextSplitter(chunk_size=500, chunk_overlap=50)
texts = text_splitter.split_text(file_content)
metadata = []
for i in range(0,len(texts)):
metadata.append({"path":file})
qdrant.add_texts(texts,metadatas=metadata)
在所提供的部分代码中,我们将文本分块为 500 个标记符的大小,并以 50 个重叠标记符为窗口。这样,我们就能在文本块结束或开始的地方保留一些上下文。在代码的其余部分,我们将根据用户硬盘上的文档路径创建元数据,并将这些带有元数据的文本块添加到索引中。
不过,在将文件内容添加到索引之前,我们需要先读取文件。甚至在读取文件之前,我们还需要获取需要索引的所有文件。为了简单起见,在本项目中,用户可以定义一个自己想要索引的文件夹。索引器将以递归方式检索该文件夹及其子文件夹中的所有文件,并为支持的文件编制索引(我们将了解如何支持 PDF、Word、PPT 和 TXT)。
我们可以以递归方式检索给定文件夹及其子文件夹中的所有文件:
def get_files(dir):
file_list = []
for f in listdir(dir):
if isfile(join(dir,f)):
file_list.append(join(dir,f))
elif isdir(join(dir,f)):
file_list= file_list + get_files(join(dir,f))
return file_list
一旦在列表中检索到所有文件,我们就可以读取包含文本的文件内容了。本工具首先支持 MS Word 文档(扩展名为“.docx”)、PDF 文档、MS PowerPoint 演示文稿(扩展名为“.pptx”)和纯文本文件(扩展名为“.txt”)。
为了读取 MS Word 文档,我们可以使用 docx-python 库。将文档读入字符串变量的函数如下所示:
import docx
def getTextFromWord(filename):
doc = docx.Document(filename)
fullText = []
for para in doc.paragraphs:
fullText.append(para.text)
return '\n'.join(fullText)
MS PowerPoint 文件也能实现类似功能。为此,我们需要下载并安装 pptx-python 库,并编写这样一个函数:
from pptx import Presentation
def getTextFromPPTX(filename):
prs = Presentation(filename)
fullText = []
for slide in prs.slides:
for shape in slide.shapes:
fullText.append(shape.text)
return '\n'.join(fullText)
读取文本文件非常简单:
f = open(file,'r')open(file,'r')
file_content = f.read()
f.close()
对于 PDF 文件,我们将使用 PyPDF2 库:
reader = PyPDF2.PdfReader(file)
for i in range(0,len(reader.pages)):
file_content = file_content + " "+reader.pages[i].extract_text()
最后,整个索引功能看起来就像这样:
file_content = """"
for file in onlyfiles:
file_content = ""
if file.endswith(".pdf"):
print("indexing "+file)
reader = PyPDF2.PdfReader(file)
for i in range(0,len(reader.pages)):
file_content = file_content + " "+reader.pages[i].extract_text()
elif file.endswith(".txt"):
print("indexing " + file)
f = open(file,'r')
file_content = f.read()
f.close()
elif file.endswith(".docx"):
print("indexing " + file)
file_content = getTextFromWord(file)
elif file.endswith(".pptx"):
print("indexing " + file)
file_content = getTextFromPPTX(file)
else:
continue
text_splitter = TokenTextSplitter(chunk_size=500, chunk_overlap=50)
texts = text_splitter.split_text(file_content)
metadata = []
for i in range(0,len(texts)):
metadata.append({"path":file})
qdrant.add_texts(texts,metadatas=metadata)
print(onlyfiles)
print("Finished indexing!")
如前所述,我们使用 LangChain 中的 TokenTextSplitter 创建了 500 个 token 的块,其中有 50 个 token 重叠。现在,当我们创建了一个索引后,就可以创建一个网络服务来查询它并生成答案了。
生成式搜索 API
我们将使用 FastAPI 创建一个网络服务,以托管我们的生成式搜索引擎。该应用程序接口将访问 Qdrant 客户端和我们在上一节中创建的索引数据,使用矢量相似度指标执行搜索,使用 Llama 3 模型生成答案,最后将答案反馈给用户。
为了初始化和导入生成式搜索组件的库,我们可以使用以下代码:
from fastapi import FastAPI
from langchain_community.embeddings import HuggingFaceEmbeddings
from langchain_qdrant import Qdrant
from qdrant_client import QdrantClient
from pydantic import BaseModel
import torch
from transformers import AutoTokenizer, AutoModelForCausalLM
import environment_var
import os
from openai import OpenAI
class Item(BaseModel):
query: str
def __init__(self, query: str) -> None:
super().__init__(query=query)
如前所述,我们正在使用 FastAPI 创建 API 接口。我们将利用 qdrant_client 库访问我们创建的索引数据,并利用 langchain_qdrant 库提供额外支持。对于嵌入和本地加载 Llama 3 模型,我们将使用 PyTorch 和 Transformers 库。此外,我们将使用 OpenAI 库调用英伟达 NIM API,API 密钥存储在我们创建的 environment_var 文件中(适用于 Nvidia 和 HuggingFace)。
我们在 Pydantic 中创建了从 BaseModel 派生的 Item 类,作为参数传递给请求函数。它将有一个名为 query 的字段。
现在,我们可以开始初始化机器学习模型了
model_name = "sentence-transformers/msmarco-bert-base-dot-v5""sentence-transformers/msmarco-bert-base-dot-v5"
model_kwargs = {'device': 'cpu'}
encode_kwargs = {'normalize_embeddings': True}
hf = HuggingFaceEmbeddings(
model_name=model_name,
model_kwargs=model_kwargs,
encode_kwargs=encode_kwargs
)
os.environ["HF_TOKEN"] = environment_var.hf_token
use_nvidia_api = False
use_quantized = True
if environment_var.nvidia_key !="":
client_ai = OpenAI(
base_url="https://integrate.api.nvidia.com/v1",
api_key=environment_var.nvidia_key
)
use_nvidia_api = True
elif use_quantized:
model_id = "Kameshr/LLAMA-3-Quantized"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="auto",
)
else:
model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
model = AutoModelForCausalLM.from_pretrained(
model_id,
torch_dtype=torch.float16,
device_map="auto",
)
在前几行中,我们为基于 BERT 的模型加载权重,该模型在 MSMARCO 数据基础上进行了微调,我们也使用该数据为文档编制索引。
然后,我们检查是否提供了 nvidia_key,如果提供了,我们就使用 OpenAI 库调用英伟达 NIM API。使用 NVIDIA NIM API 时,我们可以使用 Llama 3 指令模型的大版本,其中包含 70B 个参数。如果没有提供 nvidia_key,我们将在本地加载 Llama 3。不过,至少对于大多数消费电子产品而言,本地无法加载 70B 参数模型。因此,我们要么加载 Llama 3 8B 参数模型,要么加载经过额外量化的 Llama 3 8B 参数模型。通过量化,我们可以节省空间,以更少的 RAM 执行模型。例如,Llama 3 8B 通常需要约 14GB 的 GPU 内存,而量化后的 Llama 3 8B 只需 6GB 的 GPU 内存即可运行。因此,我们可以根据参数加载完整模型或量化模型。
现在我们可以初始化 Qdrant 客户端
client = QdrantClient(path="qdrant/")"qdrant/")
collection_name = "MyCollection"
qdrant = Qdrant(client, collection_name, hf)
此外,FastAPI 还创建了第一个模拟 GET 函数
app = FastAPI()
@app.get("/")
async def root():
return {"message": "Hello World"}
该函数将返回格式为 {“message”: “Hello World”} 的 JSON 文件。
不过,为了使该 API 能够发挥作用,我们将创建两个函数,一个函数只执行语义搜索,另一个函数将执行搜索,然后将前 10 个块作为上下文并生成答案,同时引用所使用的文档。
@app.post("/search")
def search(Item:Item):
query = Item.query
search_result = qdrant.similarity_search(
query=query, k=10
)
i = 0
list_res = []
for res in search_result:
list_res.append({"id":i,"path":res.metadata.get("path"),"content":res.page_content})
return list_res
@app.post("/ask_localai")
async def ask_localai(Item:Item):
query = Item.query
search_result = qdrant.similarity_search(
query=query, k=10
)
i = 0
list_res = []
context = ""
mappings = {}
i = 0
for res in search_result:
context = context + str(i)+"\n"+res.page_content+"\n\n"
mappings[i] = res.metadata.get("path")
list_res.append({"id":i,"path":res.metadata.get("path"),"content":res.page_content})
i = i +1
rolemsg = {"role": "system",
"content": "Answer user's question using documents given in the context. In the context are documents that should contain an answer. Please always reference document id (in squere brackets, for example [0],[1]) of the document that was used to make a claim. Use as many citations and documents as it is necessary to answer question."}
messages = [
rolemsg,
{"role": "user", "content": "Documents:\n"+context+"\n\nQuestion: "+query},
]
if use_nvidia_api:
completion = client_ai.chat.completions.create(
model="meta/llama3-70b-instruct",
messages=messages,
temperature=0.5,
top_p=1,
max_tokens=1024,
stream=False
)
response = completion.choices[0].message.content
else:
input_ids = tokenizer.apply_chat_template(
messages,
add_generation_prompt=True,
return_tensors="pt"
).to(model.device)
terminators = [
tokenizer.eos_token_id,
tokenizer.convert_tokens_to_ids("<|eot_id|>")
]
outputs = model.generate(
input_ids,
max_new_tokens=256,
eos_token_id=terminators,
do_sample=True,
temperature=0.2,
top_p=0.9,
)
response = tokenizer.decode(outputs[0][input_ids.shape[-1]:])
return {"context":list_res,"answer":response}
这两个函数都是 POST 方法,我们使用 Item 类通过 JSON 主体传递查询。第一个方法会返回 10 个最相似的文档块,并附带路径和 0-9 的文档 ID。因此,它只是使用点积作为相似度量来执行普通的语义搜索(这是在 Qdrant 中编制索引时定义的--请记住包含 distance=Distance.DOT 的一行)。
第二个名为 ask_localai 的函数稍微复杂一些。它包含了第一个方法中的搜索机制(因此,通过代码来理解语义搜索可能更容易),但增加了一个生成部分。它为 Llama 3 创建了一个提示,在系统提示信息中包含如下指令:
用户信息包含一个文件列表,其结构为 ID(0-9),下一行是文件块。为了维护 ID 和文档路径之间的映射,我们创建了一个名为 list_res 的列表,其中包括 ID、路径和内容。用户提示以 “Question(问题)”结束,随后是用户的询问。
响应包含上下文和生成的答案。不过,答案还是由 Llama 3 70B 模型(使用英伟达 NIM API)、本地 Llama 3 8B 或本地 Llama 3 8B 量化生成,具体取决于传递的参数。
API 可以从包含以下代码行的单独文件中启动(我们的生成组件位于名为 api.py 的文件中,因为 Uvicorn 中的第一个参数映射到了文件名):
import uvicorn
if __name__=="__main__":
uvicorn.run("api:app",host='0.0.0.0', port=8000, reload=False, workers=3)
简单的用户界面
本地生成式搜索引擎的最后一个组成部分是用户界面。我们将使用 Streamlit 创建一个简单的用户界面,其中包括一个输入栏、一个搜索按钮、一个用于显示生成答案的部分,以及一个可打开或下载的参考文档列表。
Streamlit 用户界面的全部代码不到 45 行(准确地说是 44 行):
import re
import streamlit as st
import requests
import json
st.title('_:blue[Local GenAI Search]_ :sunglasses:')
question = st.text_input("Ask a question based on your local files", "")
if st.button("Ask a question"):
st.write("The current question is \"", question+"\"")
url = "http://127.0.0.1:8000/ask_localai"
payload = json.dumps({
"query": question
})
headers = {
'Accept': 'application/json',
'Content-Type': 'application/json'
}
response = requests.request("POST", url, headers=headers, data=payload)
answer = json.loads(response.text)["answer"]
rege = re.compile("\[Document\ [0-9]+\]|\[[0-9]+\]")
m = rege.findall(answer)
num = []
for n in m:
num = num + [int(s) for s in re.findall(r'\b\d+\b', n)]
st.markdown(answer)
documents = json.loads(response.text)['context']
show_docs = []
for n in num:
for doc in documents:
if int(doc['id']) == n:
show_docs.append(doc)
a = 1244
for doc in show_docs:
with st.expander(str(doc['id'])+" - "+doc['path']):
st.write(doc['content'])
with open(doc['path'], 'rb') as f:
st.download_button("Downlaod file", f, file_name=doc['path'].split('/')[-1],key=a
)
a = a + 1
最后都会变成这样:
结论
本文展示了如何利用 Qdrant 将生成式人工智能与语义搜索结合起来。它通常是本地文件上的一个检索-增强生成(RAG)管道,并附有对本地文件进行引用声明的指令。整个代码长约 300 行,我们甚至增加了复杂性,让用户在 3 种不同的 Llama 3 模型之间进行选择。在这种情况下,8B 和 70B 参数模型都能很好地工作。