基于Llama3，為本地文件創(chuàng)建生成式AI搜索引擎 精華
本文分享一個(gè)開(kāi)源項(xiàng)目——一款創(chuàng)新的生成式搜索引擎，能夠?qū)崿F(xiàn)用戶(hù)與本地文件的智能互動(dòng)。此項(xiàng)目在微軟Copilot等現(xiàn)有工具的基礎(chǔ)上，推出了一種開(kāi)放源代碼的替代方案，旨在推動(dòng)技術(shù)共享與創(chuàng)新。
1 系統(tǒng)設(shè)計(jì)
為構(gòu)建本地生成式搜索引擎或助手，需要幾個(gè)組件：
- 內(nèi)容索引系統(tǒng)：負(fù)責(zé)存儲(chǔ)本地文件內(nèi)容，并配備信息檢索引擎，以便高效地搜索與用戶(hù)查詢(xún)或問(wèn)題最相關(guān)的文檔。
- 語(yǔ)言模型：用于分析選定的本地文檔內(nèi)容，并據(jù)此生成精煉的總結(jié)性答案。
- 用戶(hù)界面：為用戶(hù)提供直觀的操作界面，以便輕松地進(jìn)行查詢(xún)和獲取信息。
組件之間的交互方式如下所示：
系統(tǒng)設(shè)計(jì)和架構(gòu)。使用Qdrant作為向量存儲(chǔ)，Streamlit用于用戶(hù)界面。Llama 3可以通過(guò)Nvidia NIM API（700B版本）使用，也可以通過(guò)HuggingFace下載（80B版本）。文檔分塊使用Langchain完成。
構(gòu)建本地生成式搜索引擎的第一步是創(chuàng)建索引，用以存儲(chǔ)和檢索本地文件內(nèi)容。當(dāng)用戶(hù)提出問(wèn)題，系統(tǒng)會(huì)通過(guò)這個(gè)索引快速定位到最相關(guān)的文檔。隨后，選定的文檔內(nèi)容被送入高級(jí)語(yǔ)言模型，該模型不僅生成答案，還提供對(duì)引用文檔的明確標(biāo)注。最終，用戶(hù)界面將這些信息以清晰、易于理解的方式展示給用戶(hù)。
2 語(yǔ)義索引
語(yǔ)義索引旨在通過(guò)分析文件內(nèi)容與查詢(xún)之間的相似度，提供最相關(guān)的文檔匹配。索引的構(gòu)建采用了Qdrant作為其向量存儲(chǔ)解決方案。Qdrant客戶(hù)端庫(kù)的便利之處在于，它不需要完整的服務(wù)器端安裝，便能在工作內(nèi)存中直接進(jìn)行文檔相似性比較，極大地簡(jiǎn)化了部署流程，僅需通過(guò)pip命令安裝Qdrant客戶(hù)端即可。
Qdrant初始化時(shí)，需要預(yù)先設(shè)定所使用的向量化方法和度量標(biāo)準(zhǔn)（注意，hf參數(shù)稍后定義）。向量化和度量的具體配置應(yīng)在客戶(hù)端初始化階段完成。以下是Qdrant初始化的一個(gè)示例：
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_cnotallow=VectorParams(size=768, distance=Distance.DOT))
qdrant = Qdrant(client, collection_name, hf)為構(gòu)建向量索引，必須對(duì)硬盤(pán)中的文檔進(jìn)行嵌入處理。需要選擇合適的嵌入方法和向量比較度量標(biāo)準(zhǔn)，不同的段落、句子或詞嵌入技術(shù)將產(chǎn)生不同的結(jié)果。在文檔向量搜索中，主要挑戰(zhàn)之一是非對(duì)稱(chēng)搜索問(wèn)題，這在信息檢索領(lǐng)域極為常見(jiàn)，尤其是在處理短查詢(xún)與長(zhǎng)文檔匹配時(shí)。傳統(tǒng)的單詞或句子嵌入技術(shù)通常針對(duì)相似長(zhǎng)度的文檔進(jìn)行優(yōu)化，如果文檔長(zhǎng)度與查詢(xún)長(zhǎng)度差異過(guò)大，就可能導(dǎo)致信息檢索效果不佳。
然而，有一種嵌入方法能夠有效應(yīng)對(duì)非對(duì)稱(chēng)搜索問(wèn)題。以MSMARCO數(shù)據(jù)集為例，該數(shù)據(jù)集基于Bing的搜索查詢(xún)和文檔，由Microsoft發(fā)布，并針對(duì)此類(lèi)問(wèn)題進(jìn)行了優(yōu)化。MSMARCO數(shù)據(jù)集的模型經(jīng)過(guò)微調(diào)，能夠提供出色的搜索效果，非常適合解決當(dāng)前面臨的問(wèn)題。
在本次實(shí)現(xiàn)中，選用了針對(duì)MSMARCO數(shù)據(jù)集進(jìn)行過(guò)微調(diào)的模型，名為：
sentence-transformers/msmarco-bert-base-dot-v5這個(gè)模型基于BERT架構(gòu)，并針對(duì)點(diǎn)積相似性度量進(jìn)行了特別優(yōu)化。在初始化Qdrant客戶(hù)端時(shí)，已明確采用點(diǎn)積作為衡量相似性的方法（注意此模型的維度為768）：
client.create_collection(collection_name,vectors_cnotallow=VectorParams(size=768, distance=Distance.DOT))在選擇相似性度量標(biāo)準(zhǔn)時(shí)，雖然余弦相似性是一個(gè)可行選擇，但鑒于模型已針對(duì)點(diǎn)積優(yōu)化，采用點(diǎn)積能夠?qū)崿F(xiàn)更優(yōu)的性能表現(xiàn)。點(diǎn)積的優(yōu)勢(shì)在于它不僅關(guān)注向量間的角度差異，還包括了向量的大小因素，這在評(píng)估向量整體相似度時(shí)尤為重要。通過(guò)歸一化處理，可以在特定條件下使兩種度量標(biāo)準(zhǔn)達(dá)到相同效果。然而，當(dāng)向量的大小成為一個(gè)關(guān)鍵考量時(shí)，點(diǎn)積顯然是更為合適的度量手段。
模型初始化建議利用GPU以提升計(jì)算效率，具體代碼實(shí)現(xiàn)如下：
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類(lèi)模型受限于其內(nèi)存消耗的二次方增長(zhǎng)特性，只能處理有限長(zhǎng)度的上下文，通常不超過(guò)512個(gè)token。面對(duì)這一局限，有兩種應(yīng)對(duì)策略：一是僅利用文檔前512個(gè)token生成答案，舍棄之后的內(nèi)容；二是將文檔切分為多個(gè)小塊，每塊作為一個(gè)獨(dú)立單元存儲(chǔ)于索引之中。為了保留完整的信息，我們選擇了后者。文檔的分塊工作，計(jì)劃利用LangChain的內(nèi)置分塊工具來(lái)完成：
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)在編寫(xiě)的代碼中，把文本切割成500個(gè)token的段落，并設(shè)置了50個(gè)token的重疊區(qū)域，這樣做是為了在段落的首尾保持上下文的連貫性。接著，為每個(gè)段落創(chuàng)建了包含文檔存儲(chǔ)路徑的元數(shù)據(jù)，并將其與文本段落一并索引。
在將文件內(nèi)容索引之前，必須先讀取這些文件。而在讀取之前，需要先確定哪些文件需要被索引。本項(xiàng)目簡(jiǎn)化了這一流程，允許用戶(hù)指定他們希望索引的文件夾。索引器將遞歸地搜索該文件夾及其子文件夾中的所有文件，并索引那些支持的文件類(lèi)型，如PDF、Word、PPT和TXT格式。
以下是檢索給定文件夾及其子文件夾內(nèi)所有文件的遞歸方法：
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完成文件檢索后，接下來(lái)便是讀取這些文件中的文本內(nèi)容。目前，工具支持的文件格式包括MS Word文檔（.docx）、PDF文檔、MS PowerPoint演示文稿（.pptx）以及純文本文件（.txt）。
對(duì)于MS Word文檔的讀取，采用docx-python庫(kù)來(lái)實(shí)現(xiàn)。以下是將文檔內(nèi)容讀取到字符串變量中的函數(shù)示例：
import docx
def getTextFromWord(filename):
doc = docx.Document(filename)
fullText = []
for para in doc.paragraphs:
fullText.append(para.text)
return '\n'.join(fullText)對(duì)于MS PowerPoint文件的處理，采取相似的方法。為此，需要下載并安裝pptx-python庫(kù)，并編寫(xiě)如下函數(shù)：
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()對(duì)于PDF文件，使用 PyPDF2 庫(kù)：
reader = PyPDF2.PdfReader(file)
for i in range(0,len(reader.pages)):
file_content = file_content + " "+reader.pages[i].extract_text()最后，整個(gè)索引函數(shù)是這樣：
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個(gè)token的段落，并在段落間保留了50個(gè)token的重疊，確保了內(nèi)容的連續(xù)性。在此基礎(chǔ)上，已經(jīng)成功建立了索引。接下來(lái)，將開(kāi)發(fā)一個(gè)Web服務(wù)，它不僅能夠查詢(xún)索引，還能根據(jù)查詢(xún)結(jié)果智能生成答案。
3 生成式搜索API
這里通過(guò)FastAPI框架搭建Web服務(wù)，用于承載生成式搜索引擎。這個(gè)API將連接到之前建立的Qdrant客戶(hù)端索引，通過(guò)向量相似性搜索算法深入挖掘，再借助Llama 3模型對(duì)篩選出的最相關(guān)塊生成精準(zhǔn)答案，并將這些答案反饋給用戶(hù)。
為了配置并引入生成式搜索的關(guān)鍵組件，以下是相應(yīng)的代碼示例：
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 框架用來(lái)創(chuàng)建 API 接口，以實(shí)現(xiàn)數(shù)據(jù)的高效交互。通過(guò) qdrant_client 庫(kù)，能夠訪(fǎng)問(wèn)之前建立的索引數(shù)據(jù)，而 langchain_qdrant 庫(kù)則增強(qiáng)了其功能。在處理模型嵌入和本地化部署 Llama 3 模型時(shí)，分別采用了 PyTorch 和 Transformers 這兩個(gè)業(yè)界領(lǐng)先的庫(kù)。此外，項(xiàng)目還通過(guò) OpenAI 庫(kù)與 NVIDIA NIM API 進(jìn)行了集成，相關(guān)的 API 密鑰被安全地存儲(chǔ)在預(yù)設(shè)的 environment_var 文件中，確保了與 Nvidia 和 HuggingFace 的無(wú)縫對(duì)接。
為了更高效地處理請(qǐng)求參數(shù)，開(kāi)發(fā)了一個(gè)名為 Item 的類(lèi)，它基于 Pydantic 的 BaseModel 進(jìn)行擴(kuò)展，并且包含了一個(gè)關(guān)鍵字段：query，該字段專(zhuān)用于捕獲和傳遞用戶(hù)的查詢(xún)指令。
緊接著，項(xiàng)目將啟動(dòng)機(jī)器學(xué)習(xí)模型的初始化過(guò)程：
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",
)系統(tǒng)已完成對(duì)基于MSMARCO數(shù)據(jù)集優(yōu)化的BERT模型的加載，該模型用于文檔索引工作。
若存在nvidia_key，系統(tǒng)會(huì)調(diào)用NVIDIA NIM API，啟用具有70億參數(shù)的Llama 3 instruct模型。若無(wú)nvidia_key，鑒于本地部署限制，將加載或量化處理后的8億參數(shù)Llama 3模型，使其在減少內(nèi)存占用的同時(shí)，保持模型性能。
接下來(lái)，啟動(dòng)Qdrant客戶(hù)端的初始化過(guò)程，以便進(jìn)行高效的數(shù)據(jù)索引和檢索：
client = QdrantClient(path="qdrant/")
collection_name = "MyCollection"
qdrant = Qdrant(client, collection_name, hf)同時(shí)，使用 FastAPI 創(chuàng)建第一個(gè)模擬 GET 函數(shù)：
app = FastAPI()
@app.get("/")
async def root():
return {"message": "Hello World"}這個(gè)函數(shù)會(huì)返回格式為 {"message":"Hello World"} 的 JSON。
為了確保API能夠正常工作，將設(shè)計(jì)兩個(gè)功能：第一個(gè)功能專(zhuān)門(mén)進(jìn)行語(yǔ)義搜索；第二個(gè)功能則在搜索的基礎(chǔ)上，選取最相關(guān)的前10個(gè)文本塊作為上下文，進(jìn)一步生成答案，并對(duì)使用的文檔進(jìn)行引用。
@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}這兩個(gè)函數(shù)均采用POST方法，并通過(guò)JSON格式利用Item類(lèi)傳遞查詢(xún)參數(shù)。第一個(gè)函數(shù)負(fù)責(zé)返回10個(gè)最相似的文檔片段，同時(shí)提供每個(gè)片段的路徑，并賦予其從0至9的文檔ID。該函數(shù)主要執(zhí)行基礎(chǔ)的語(yǔ)義搜索，使用點(diǎn)積作為相似性度量標(biāo)準(zhǔn)，這一點(diǎn)在Qdrant索引創(chuàng)建期間已設(shè)定——即在定義中包含了distance=Distance.DOT的參數(shù)。
第二個(gè)名為ask_localai的函數(shù)則更為復(fù)雜，它在第一個(gè)函數(shù)的搜索機(jī)制基礎(chǔ)上進(jìn)行了擴(kuò)展，增加了生成答案的功能。該函數(shù)為L(zhǎng)lama 3模型構(gòu)建了一個(gè)包含系統(tǒng)提示消息的提示模板，指示模型如何生成答案：
請(qǐng)使用上下文中給出的文檔回答用戶(hù)的提問(wèn)。上下文中的文檔應(yīng)當(dāng)包含問(wèn)題的答案。在陳述時(shí)，請(qǐng)始終引用用來(lái)提出主張的文檔的ID（用方括號(hào)表示，例如[0]、[1]）。根據(jù)回答問(wèn)題的需要，盡可能多地引用文獻(xiàn)和文檔。
用戶(hù)的消息包含了一個(gè)文檔列表，列表中的每個(gè)文檔都按ID（0-9）編號(hào)，并在下一行顯示文檔內(nèi)容。為了保持ID與文檔路徑之間的映射關(guān)系，我們創(chuàng)建了一個(gè)名為list_res的列表，其中包含了ID、路徑和內(nèi)容。用戶(hù)提示以“Question”一詞結(jié)束，隨后是用戶(hù)的查詢(xún)內(nèi)容。
響應(yīng)包含上下文和生成的答案。然而，答案再次由 Llama 3 70B 模型（使用 NVIDIA NIM API）、本地 Llama 3 8B 或本地量化的 Llama 3 8B 生成，具體取決于傳遞的參數(shù)。
API 可以從包含以下代碼的單獨(dú)文件啟動(dòng)（假設(shè)生成組件在名為 api.py 的文件中，Uvicorn 的第一個(gè)參數(shù)對(duì)應(yīng)文件名）：
import uvicorn
if __name__=="__main__":
uvicorn.run("api:app",host='0.0.0.0', port=8000, reload=False, workers=3)4 簡(jiǎn)單的用戶(hù)界面
本地生成式搜索引擎的用戶(hù)界面是其最后一塊拼圖，采用Streamlit構(gòu)建，界面簡(jiǎn)潔，包含查詢(xún)輸入框、搜索按鈕、結(jié)果展示區(qū)以及可交互的文檔列表。關(guān)鍵代碼不足45行:
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最終：
5 結(jié)語(yǔ)
本文闡述了如何融合Qdrant的語(yǔ)義搜索技術(shù)與生成式人工智能，構(gòu)建了針對(duì)本地文件的檢索增強(qiáng)生成（RAG）流程，這一流程能夠?qū)ξ臋n中的聲明進(jìn)行引用說(shuō)明。代碼總計(jì)約300行，用戶(hù)可選擇三種不同參數(shù)規(guī)模的Llama 3模型以滿(mǎn)足不同場(chǎng)景的需求。在本用例中，無(wú)論是8億還是70億參數(shù)的模型，均能穩(wěn)定運(yùn)行并提供出色的性能。
本文轉(zhuǎn)載自 ??AI科技論談??，作者： AI科技論談
