实践记录02文档QA系统火山API ｜ 豆包MarsCode AI 刷题

前言

可能有很多小伙伴免费的Token已经用完了，使用的是注册火山引擎的免费50w Token。在这里主要记录一下使用火山引擎API来实现02文档QA系统，问题主要是在创建Embedding层上。

代码流程

首先是遍历文档下的所有文件，根据文件的后缀名采取不同的读取方式

# 加载Documents
base_dir = "./OneFlower"  # 文档的存放目录
documents = []
for file in os.listdir(base_dir):
    # 构建完整的文件路径
    file_path = os.path.join(base_dir, file)
    if file.endswith(".pdf"):
        loader = PyPDFLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith(".docx"):
        loader = Docx2txtLoader(file_path)
        documents.extend(loader.load())
    elif file.endswith(".txt"):
        loader = TextLoader(file_path)
        documents.extend(loader.load())

这段代码的主要目的是将读取到的多个文档分割成多个小块，每个块的大小为 200 个字符，相邻块之间有 10 个字符的重叠。这种分割方式通常用于处理大型文本数据，以便于后续的处理或分析。

from langchain.text_splitter import RecursiveCharacterTextSplitter
text_splitter = RecursiveCharacterTextSplitter(chunk_size=200, chunk_overlap=10)
chunked_documents = text_splitter.split_documents(documents)

接下这段代码与直接调用OpenAI的Embedding层不同，官方提供了一个DoubaoEmbeddings类来调用Doubao的Embedding层。如果你还有10k的免费Token额度，那么你可以顺利运行，但是如果你额度用完了，并且想使用火山方舟提供的免费Token额度，你可以参考第三部分的内容。

class DoubaoEmbeddings(BaseModel, Embeddings):
    client: Ark = None
    api_key: str = ""
    model: str

    def __init__(self, **data: Any):
        super().__init__(**data)
        if self.api_key == "":
            self.api_key = os.environ["OPENAI_API_KEY"]
        self.client = Ark(
            base_url=os.environ["OPENAI_BASE_URL"],
            api_key=self.api_key,
            timeout=2000,
            max_retries=2
        )

    def embed_query(self, text: str) -> List[float]:
        """
        生成输入文本的 embedding.
        Args:
            texts (str): 要生成 embedding 的文本.
        Return:
            embeddings (List[float]): 输入文本的 embedding，一个浮点数值列表.
        """
        embeddings = self.client.embeddings.create(model=self.model, input=text)
        return embeddings.data[0].embedding

    def embed_documents(self, texts: List[str]) -> List[List[float]]:
        return [self.embed_query(text) for text in texts]

    class Config:
        arbitrary_types_allowed = True


vectorstore = Qdrant.from_documents(
    documents=chunked_documents,  # 以分块的文档
    embedding=DoubaoEmbeddings(
        model=os.environ["EMBEDDING_MODELEND"],
    ),  # 用火山引擎的Embedding Model做嵌入
    location=":memory:",  # in-memory 存储
    collection_name="my_documents",
) 

这段代码的主要目的是将分块后的文档存储在一个内存中的向量存储中，并使用指定的嵌入模型生成文档的嵌入向量。向量存储可以用于后续的相似度搜索、聚类等任务。

from langchain_community.vectorstores import Qdrant
vectorstore = Qdrant.from_documents(
    documents=chunked_documents,  # 以分块的文档
    embedding=DoubaoEmbeddings(
        model=os.environ["EMBEDDING_MODELEND"],
    ),  # 用OpenAI的Embedding Model做嵌入
    location=":memory:",  # in-memory 存储
    collection_name="my_documents",
)  # 指定collection_name

这段代码的主要目的是构建一个基于向量存储的问答系统。首先，它使用 ChatOpenAI 模型生成自然语言文本，然后通过 MultiQueryRetriever 从多个角度生成查询，最后使用 RetrievalQA 链结合检索器和模型来生成回答。

import logging  # 导入Logging工具
from langchain_openai import ChatOpenAI  # ChatOpenAI模型
from langchain.retrievers.multi_query import (
    MultiQueryRetriever,
)  # MultiQueryRetriever工具
from langchain.chains import RetrievalQA  # RetrievalQA链

# 设置Logging
logging.basicConfig()
logging.getLogger("langchain.retrievers.multi_query").setLevel(logging.INFO)

# 实例化一个大模型工具 - Doubao-pro-32k
llm = ChatOpenAI(model=os.environ["LLM_MODELEND"], temperature=0)

# 实例化一个MultiQueryRetriever
retriever_from_llm = MultiQueryRetriever.from_llm(
    retriever=vectorstore.as_retriever(), llm=llm
)

# 实例化一个RetrievalQA链
qa_chain = RetrievalQA.from_chain_type(llm, retriever=retriever_from_llm)

Flask 是一个轻量级的 Web 框架，用于构建 Web 应用程序。这段代码允许用户输入问题并通过之前构建的 RetrievalQA 链生成回答。生成的回答结果会显示在网页上。

from flask import Flask, request, render_template
app = Flask(__name__)  # Flask APP
@app.route("/", methods=["GET", "POST"])
def home():
    if request.method == "POST":
        # 接收用户输入作为问题
        question = request.form.get("question")

        # RetrievalQA链 - 读入问题，生成答案
        result = qa_chain({"query": question})

        # 把大模型的回答结果返回网页进行渲染
        return render_template("index.html", result=result)

    return render_template("index.html")

if __name__ == "__main__":
    app.run(host="0.0.0.0", debug=True, port=5000)

运行截图

点击http://10.140.0.115:5000， 会跳转到外网可以访问的地址

创建火山Embedding层

去火山方舟控制台，点击在线推理 -> 创建推理接入点。 如果还没有接触过火山方舟的同学，可以参考README.md文档创建账号免费领取Token，教你如何获取API Key 、模型名以及BaseUrl。 接下来按照图中所示创建一个专用的Embedding接入点 然后复制下来以ep开头的接入点ID 在命令行中输入：export EMBEDDING_MODELEND=你的接入点ID，即可替换掉原来的“Doubao-Embed”。下面是一个简单的使用Embedding层的例子，来判断你是否修改成功。

import os
from volcenginesdkarkruntime import Ark
client = Ark(api_key=os.environ.get("OPENAI_API_KEY"))
print("----- embeddings request -----")
resp = client.embeddings.create(
    model=os.environ["EMBEDDING_MODELEND"],
    input=["花椰菜又称菜花、花菜，是一种常见的蔬菜。"]
)
print(resp)

总结

这篇文章主要介绍了如何使用火山方舟API构建一个文档QA系统，重点在于创建Embedding层的过程。在后续的代码和实践中也会频繁用到Embedding层进行词嵌入，所以在这里也是解决掉了一个大问题。