5月23日，我收到了NVIDIA的一封邮件，邀请我参加NVIDIA和LangChain联合举办的生成式AI代理开发竞赛。我的第一反应是时间有点紧，因为我们最近刚有了一个宝宝，而且我父母也打算来，我可能没有时间参加。但后来我又想了想，决定还是可以编写一些代码并提交。我思考了几天，有一个想法一直在我脑海中浮现——一个开源的生成式搜索引擎，可以让你与本地文件进行互动。微软的Copilot已经提供了类似的功能，但我认为我可以做一个开源版本，好玩一点，也可以分享一下我在快速开发系统过程中学到的一些东西。
系统设计为了构建一个本地的生成式搜索引擎或助手，我们需要几个组件：
- 使用内容检索引擎从本地文件中检索与给定查询/问题最相关的文档。
- 使用本地文档中选定的内容生成摘要回答的语言模型
- 用户界面
下面的图表展示了各个组件之间的交互。
系统设计和架构。Qdrant 用于向量存储，而 Streamlit 用于用户界面。Llama 3 可以通过 Nvidia NIM API（70B 版本）使用，或者通过 HuggingFace（8B 版本）下载。文档分块使用 Langchain 完成。作者供图
首先，我们需要将本地文件索引到可以查询本地文件内容的索引中。然后，当用户提问时，我们会使用创建的索引，并利用一些不对称的段落或文档嵌入来检索最相关的文档，这些文档可能包含答案。这些文档的内容和问题会被传递给部署的大语言模型，该模型会使用给定文档的内容生成答案。在指令提示中，我们会要求大语言模型也返回所用文档的引用。最终，所有内容都会在用户界面上展示给用户。
现在，我们来详细看看每个组件。
索引构建我们正在构建一个语义索引，该索引将根据文件内容与给定查询的相似性为我们提供最相关的文档。为了创建这样的索引，我们将使用 Qdrant 作为向量存储。有趣的是，Qdrant 客户端库不需要安装完整的 Qdr ant 服务器，并且可以对适合工作内存（RAM）的文档进行相似性比较。因此，我们只需要 pip 安装 Qdrant 客户端。
我们可以按照以下方式初始化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数据集基于Bing搜索的查询和文档，由微软发布。因此，它非常适合我们处理的问题。
对于这个特定的实现，我选择了一个已经经过微调的模型，称为：
sentence-transformers/msmarco-bert-base-dot-v5此模型基于 BERT，并使用点积作为相似度度量进行了微调。我们已经初始化了 qdrant 客户端，以在行中使用点积作为相似度度量（请注意，此模型的维度为 768）。
client.create_collection(collection_name, vectors_config=VectorParams(size=768, distance=Distance.DOT))我们可以使用其他指标，例如余弦相似度，但是鉴于该模型是通过点积进行微调的，使用该指标可以获得最佳性能。从几何角度考虑：余弦相似度仅关注角度差异，而点积则同时考虑角度和大小。通过将数据归一化为统一的大小，这两种度量方法就变得等效了。在忽略大小有益的情况下，余弦相似度是有用的。然而，如果大小很重要，点积是更合适的相似度度量。
初始化MSMarco模型的代码是（如果有可用的GPU，请务必使用）：
model_name = "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 类似的模型由于 transformer 模型的二次内存需求，其上下文大小有限。在许多 BERT 类似的模型中，这个上下文大小被设定为 512 个 token。有两种选择：（1）我们只能基于前 512 个 token 来构建答案，并忽略文档的其余部分，或者（2）创建一个索引，在该索引中，一个文档会被分割成多个片段并作为片段存储。在第一种情况下，我们会丢失很多重要信息，因此我们选择了第二种变体。为了分割文档，我们可以使用 LangChain 提供的预构建分块器：
从 langchain_text_splitters 导入 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个token的块，并且每块之间有50个token的重叠。这样可以在块的开始和结束位置保留一些上下文。在代码的其余部分，我们使用用户的硬盘上的文档路径创建元数据，并将这些带有元数据的块添加到索引中。
然而，在我们将文件内容添加到索引之前，我们需要先读取这些文件。甚至在读取文件之前，我们需要获取所有需要索引的文件。为了简化起见，在这个项目中，用户可以定义一个希望索引的文件夹。索引器会递归地从该文件夹及其子文件夹中获取所有文件，并索引受支持的文件（我们将看看如何支持PDF、Word、PPT和TXT文件）。
我们可以递归地检索给定文件夹及其子文件夹中的所有文件：
def 获取文件(dir):
文件列表 = []
for f in listdir(dir):
if isfile(join(dir,f)):
文件列表.append(join(dir,f))
elif isdir(join(dir,f)):
文件列表 = 文件列表 + 获取文件(join(dir,f))
return 文件列表一旦所有文件都被列出来，我们就可以读取包含文本的文件的内容。在这个工具中，首先我们将支持 Microsoft Word 文档（扩展名为“.docx”）、PDF 文档、Microsoft PowerPoint 演示文稿（扩展名为“.pptx”）和纯文本文件（扩展名为“.txt”）。
为了读取 MS Word 文档，我们可以使用 docx-python 库。读取文档到字符串变量的函数看起来像这样：
import docx
def 从Word获取文本(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 从PPTX获取文本(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')
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("索引 " + 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("索引 " + file)
f = open(file, 'r')
file_content = f.read()
f.close()
elif file.endswith(".docx"):
print("索引 " + file)
file_content = getTextFromWord(file)
elif file.endswith(".pptx"):
print("索引 " + 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("索引完成!")正如我们所说，我们使用 LangChain 的 TokenTextSplitter 将文本切分成每个包含 500 个 token 的块，并且每块之间有 50 个 token 的重叠。现在，当我们创建了一个索引之后，我们可以创建一个网络服务来查询该索引并生成答案。
生成式搜索API我们将使用 FastAPI 创建一个网络服务来托管我们的生成式搜索引擎。该 API 将访问我们在上一节中创建的包含索引数据的 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 库调用 NVIDIA NIM API，API 密钥存储在我们创建的 environment_var（用于 Nvidia 和 HuggingFace）文件中。
我们创建了一个名为 Item 的类，该类继承自 Pydantic 的 BaseModel，用作请求函数的参数传递。它将包含一个名为 query 的字段。
现在，我们可以开始初始化我们的机器学习模型了。
model_name = "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 库调用 NVIDIA NIM API。当我们使用 NVIDIA NIM API 时，可以使用具有 70B 参数的 Llama 3 指令模型的大版本。如果未提供 nvidia_key，我们将本地加载 Llama 3。然而，在本地，至少对于大多数消费电子设备来说，加载 70B 参数的模型是不可能的。因此，我们将加载 Llama 3 8B 参数模型或额外量化后的 Llama 3 8B 参数模型。通过量化，我们可以节省空间并使模型在较少的 RAM 上运行。例如，Llama 3 8B 通常需要大约 14GB 的 GPU RAM，而量化后的 Llama 3 8B 可以在 6GB 的 GPU RAM 上运行。因此，我们将根据参数加载完整模型或量化模型。
我们现在可以初始化Qdrant客户端
client = QdrantClient(path="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": "使用给定的文档回答用户的问题。在上下文中包含的文档应该包含答案。请始终引用用于提出主张的文档的文档ID（例如[0],[1]）。使用尽可能多的引用和文档来回答问题。"}
messages = [
rolemsg,
{"role": "user", "content": "文档:\n"+context+"\n\n问题: "+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("")
]
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 个最相似的文档片段，并附上路径，分配文档 ID 从 0 到 9。因此，它只是执行了简单的语义搜索，使用点积作为相似度度量（这在 Qdrant 索引时定义过 —— 记得包含 distance=Distance.DOT 的那行）。
第二个函数叫做 ask_localai，稍微复杂一些。它包含第一个方法中的搜索机制（因此可能更容易通过查看该代码来理解语义搜索），但增加了生成部分。它为 Llama 3 创建了一个提示，包含一个系统提示消息中的指令，说明如下：
用给定的文档回答用户的问题。上下文中包含的文档应该包含答案。请始终引用用于提出主张的文档的文档ID（用方括号括起来，例如[0],[1]）。使用尽可能多的引用和文档来回答问题。
用户的消息包含一个文档列表，这些文档的结构是以一个ID（0–9）开头，接着是文档的内容。为了保持ID和文档路径之间的映射关系，我们创建了一个名为list_res的列表，其中包括ID、路径和内容。用户的提示以单词“Question”加上用户的查询结束。
响应包含上下文和生成的答案。然而，答案是由 Llama 3 70B 模型（使用 NVIDIA 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[本地GenAI搜索]_ :sunglasses:')
question = st.text_input("基于你的本地文件提出一个问题", "")
if st.button("提出问题"):
st.write("当前问题是 \"", 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("下载文件", f, file_name=doc['path'].split('/')[-1],key=a
)
a = a + 1它最终会看起来像这样：
一个在构建的用户界面中回答的问题示例。作者截图。
可用性该项目的全部代码可在 GitHub 上找到，地址为 https://github.com/nikolamilosevic86/local-genAI-search。过去，我参与过几个生成式搜索项目，并且这些项目也有一些相关发表。你可以查看 https://www.thinkmind.org/library/INTERNET/INTERNET_2024/internet_2024_1_10_48001.html 或 https://arxiv.org/abs/2402.18589。
结论本文展示了如何利用生成式AI和语义搜索结合Qdrant。这通常是一个基于本地文件的检索增强生成（RAG）管道，其中包含指示用户引用本地文档的说明。整个代码大约有300行，我们甚至增加了复杂性，让用户在三种不同的Llama 3模型中进行选择。对于这个用例，无论是8B参数模型还是70B参数模型都表现得相当好。
我想解释一下我所做的步骤，以防将来对其他人有所帮助。不过，如果你想使用这个特定的工具，最简单的方法就是从 GitHub 获取它，它是开源的！
共同學習，寫下你的評論
評論加載中...
作者其他優質文章
