如何使用Llama 3构建本地文件的生成式搜索引擎

使用Qdrant、Nvidia NIM API或Llama 3 8B在本地搭建你的GenAI助手

5月23日，我收到了来自Nvidia的一位人士的邮件，邀请我参加NVIDIA与LangChain联合举办的生成式AI代理开发者大赛 。起初，我觉得时间相当紧迫，况且我们家最近添了宝宝，父母也即将来访，我恐怕没有时间参与。但转念一想，我决定可以编写一些代码并提交作品。我思考了几天，一个想法萦绕心头——打造一个开源的生成式搜索引擎，让你能够与本地文件互动。微软的Copilot已经提供了类似功能，但我想我可以制作一个开源版本，纯粹为了乐趣，并分享我在快速开发这个系统过程中积累的一些经验。

系统设计

为了构建一个本地生成式搜索引擎或助手，我们需要以下几个组件：

• 一个包含本地文件内容的索引，以及一个信息检索引擎，用于根据给定的查询/问题检索最相关的文档。
• 一个语言模型，用于从本地文档中选择内容并生成摘要答案。
• 一个用户界面

各组件之间的交互关系如下图所示。

首先，我们需要将本地文件索引化，以便查询本地文件的内容。然后，当用户提出问题时，我们将使用创建的索引，结合一些非对称段落或文档嵌入，检索可能包含答案的最相关文档。这些文档的内容和问题将被传递给部署的大型语言模型，该模型将利用给定文档的内容生成答案。在指令提示中，我们会要求大型语言模型同时返回所用文档的引用。最终，所有内容将在用户界面上可视化呈现给用户。

现在，让我们更详细地看一下每个组件。

语义索引

我们正在构建一个语义索引，该索引将根据文件内容与给定查询的相似度，为我们提供最相关的文档。为了创建这样的索引，我们将使用 Qdrant 作为向量存储。有趣的是，Qdrant 客户端库 不需要完整安装 Qdrant 服务器 ，并且可以处理适合工作内存（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 中的预构建分割器：

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 个 token 的大小，并保留 50 个重叠 token 的窗口。这样我们可以在块的结束或开始处保留一些上下文。在其余代码中，我们创建包含文档路径的元数据，并将这些带有元数据的块添加到索引中。

然而，在我们将文件内容添加到索引之前，我们需要读取它。甚至在读取文件之前，我们需要获取所有需要索引的文件。为了简单起见，在这个项目中，用户可以定义一个他/她想要索引的文件夹。索引器以递归方式从该文件夹及其子文件夹中检索所有文件，并索引支持的文件（我们将探讨如何支持 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 个 token 大小的块，并保留 50 个 token 的重叠。现在，当我们创建了索引后，我们可以创建一个用于查询索引并生成答案的 Web 服务。

生成式搜索API

我们将使用FastAPI创建一个Web服务来托管我们的生成式搜索引擎。该API将通过向量化相似度度量访问我们在上一节中创建的索引数据，使用Qdrant客户端执行搜索，利用前10个片段生成答案，并最终将答案返回给用户。

为了初始化和导入生成式搜索组件所需的库，我们可以使用以下代码：

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）。

我们创建了从Pydantic的BaseModel派生的Item类，用于作为请求函数的参数传递。它将有一个名为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时，可以使用具有70B参数的大型Llama 3 instruct模型。如果未提供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 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],[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[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

最终效果如下：

可用性

该项目完整代码已在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 获取，它是完全开源的！