系统设计和架构。Qdrant 用于存储向量数据,Streamlit 用作用户界面。Llama 3 可通过 Nvidia NIM API(70B 版本)或通过 HuggingFace(8B 版本)下载。文档的分块处理由 Langchain 完成。
首先,我们需要将本地文件索引到一个可查询本地文件内容的索引中。然后,当用户提出问题时,我们将使用已创建的索引,通过一些非对称段落或文档嵌入来检索可能包含答案的最相关文档。这些文档的内容和问题会被传送给部署好的大型语言模型,该模型将利用这些文档的内容生成答案。在指令提示中,我们还会要求大型语言模型返回所使用文档的引用。最终,所有内容都将在用户界面上向用户展示。
3、语义索引我们正在构建一个语义索引,它将根据文件内容的相似度和给定查询为我们提供最相关的文档。为此,我们将使用 Qdrant 作为向量存储。有趣的是,使用 Qdrant 客户端库并不需要完整安装 Qdrant 服务器,它可以在工作内存(RAM)中完成文档的相似度匹配。因此,我们只需要通过 pip 安装 Qdrant 客户端即可。

我们可以按照以下方式初始化 Qdrant(注意,由于故事情节的发展,hf参数稍后会定义,但在使用 Qdrant 客户端时,你已经需要定义使用的向量化方法和度量标准):
from qdrant_client import QdrantClientfrom qdrant_client.models import Distance, VectorParamsclient = 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 搜索查询和文档,并由 Microsoft 发布。因此,它非常适合我们正在处理的问题。
对于这个特定的实现,我选择了一个已经经过微调的模型,名为:
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 词元。我们有两个选择:(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') 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 个词元的块,有 50 个词元的重叠。现在,当我们已经创建了索引,我们可以创建一个 web 服务来查询它并生成答案。
4、生成式搜索 API我们将使用 FastAPI 创建一个 web 服务来托管我们的生成式搜索引擎。API 将访问我们在上一节中创建的带索引数据的 Qdrant 客户端,使用向量相似度度量执行搜索,使用前 10 个块生成答案,并使用 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", )
在前几行中,我们加载了基于 MSMARCO 数据微调的 BERT 基础模型的权重,我们也用这些数据来索引我们的文档。
然后,我们检查是否提供了 nvidia_key,如果提供了,我们将使用 OpenAI 库调用 NVIDIA 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 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": "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 square 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],[1])。根据需要使用尽可能多的引用和文档来回答问题。
用户的消息包含一个以 ID(0-9)开头,接着是下一行的文档块的文档列表。为了保持 ID 与文档路径之间的映射,我们创建了一个名为 list_res 的列表,其中包括 ID、路径和内容。用户提示以单词“问题”和用户的查询结束。
响应包含上下文和生成的答案。然而,答案再次由 Llama 3 70B 模型(使用 NVIDIA NIM API)、本地 Llama 3 8B 或本地 Llama 3 8B 量化模型生成,具体取决于传递的参数。
API 可以从包含以下代码行的单独文件启动(假设我们的生成组件在名为 api.py 的文件中,因为 Uvicorn 中的第一个参数映射到文件名):
import uvicornif __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
最终看起来会是这样:
6、结论
本文展示了如何利用生成式 AI 与 Qdrant 进行语义搜索。它通常是一个关于本地文件的检索增强生成(RAG)管道,并指示引用本地文档的声明。整个代码约 300 行长,我们甚至增加了复杂性,让用户在 3 种不同的 Llama 3 模型之间做出选择。对于这个用例,8B 和 70B 参数模型都工作得很好。
资源链接
GitHub :https://github.com/nikolamilosevic86/local-genAI-search