百行代碼實現三體問答機器人

透過 langchain + gradio 僅用百行 python 代碼即可實現一個簡易的 RAG 應用。本文將以《三體》為例，解讀實現步驟。

什麼是 RAG？#

Retrieval Augmented Generation (RAG)是一種增強 LLM 針對特定信息推理能力的方法。LLM，如 chatgpt，一般是在某個機構使用特定時間的數據集進行訓練得到的模型，而當我們需要對特定領域的信息或私密信息進行問答時，它就難以給出準確答案了。例如直接讓 chatgpt 介紹一下三體中章北海的所作所為，就會得到以下結果：

這種情況下用已有的數據重新微調大模型代價很大，且難以做到即時的效果，而 RAG 很好地解決了該問題。

如何實現 RAG？#

實現一個 RAG 應用一般分為兩步：建立索引和檢索生成。

建立索引#

透過 Embedding 模型將源數據轉換為詞向量保存至向量數據庫中，通常有以下步驟：

Load: 首先使用 DocumentLoader 讀取不同類型的文檔數據。
Split: 然後將文檔按一定規則分割為較小的塊（chunk），以便於模型能夠更好的進行上下文理解。
Store: 最後透過 Embedding 模型將分割出來的塊映射為向量，並存儲在向量數據庫中，以便於檢索。

檢索生成#

實際這是兩個步驟：

Retrieve: 將用戶輸入的問題也轉換為詞向量，並在向量數據庫中檢索相關性最高的塊。
Generate: 使用 ChatModel 或 LLM，如 chatgpt，根據用戶的問題基於特定的 prompt 對檢索到的內容生成摘要。

什麼是 langchain？#

LangChain是一個強大的框架，旨在幫助開發人員使用語言模型構建端到端的應用程序。它提供了一套工具、組件和接口，可簡化創建由大型語言模型 (LLM) 和聊天模型提供支持的應用程序的過程。LangChain 可以輕鬆管理與語言模型的交互，將多個組件鏈接在一起，並集成額外的資源，例如 API 和數據庫。

什麼是 gradio？#

Gradio是一個用於快速構建互動式應用程序的開源 Python 庫。它可以幫助開發者輕鬆地將機器學習模型集成到用戶友好的界面中，從而使模型更易於使用和理解。

使用 langchain+gradio 快速實現一個三體問答機器人#

項目源碼已發布至 GitHub：https://github.com/zivenyang/3body-chatbot
以下是核心代碼（由於當時國內難以申請 openai 賬號，故使用的 Azure Openai 的 api）

from dotenv import load_dotenv
from langchain.chat_models import AzureChatOpenAI
from langchain.memory import ConversationBufferMemory
from langchain.chains import ConversationalRetrievalChain
from langchain.vectorstores import Chroma
from langchain.embeddings import ModelScopeEmbeddings
from langchain.document_loaders import TextLoader, DirectoryLoader
from langchain.text_splitter import RecursiveCharacterTextSplitter
import gradio as gr
import os

# 導入.env中的變量，AZURE_OPENAI_ENDPOINT和AZURE_OPENAI_API_KEY
load_dotenv()

# 詞嵌入模型，因為是中文小說，所以使用的達摩院訓練的中文詞嵌入模型
MODEL_ID = "damo/nlp_gte_sentence-embedding_chinese-base"
# 向量數據庫存儲路徑
PERSIST_DIRECTORY = 'docs/chroma/'

# 為了在控制台同時輸出索引文檔和答案，不然會報錯
class AnswerConversationBufferMemory(ConversationBufferMemory):
    def save_context(self, inputs, outputs) -> None:
        return super(AnswerConversationBufferMemory, self).save_context(inputs,{'response': outputs['answer']})

def create_db():
    """讀取本地文件並生成詞向量存入向量數據庫"""

    # 讀取本地文件，即三體小說
    text_loader_kwargs={'autodetect_encoding': True}
    loader = DirectoryLoader("./docs", glob="**/*.txt", loader_cls=TextLoader, loader_kwargs=text_loader_kwargs)
    pages = loader.load()

    # 文件分塊，chunk_size與也與顯卡性能有關，顯存越大分的越細
    text_splitter = RecursiveCharacterTextSplitter(
    chunk_size = 512,
    chunk_overlap = 0,
    length_function = len,
        )
    splits = text_splitter.split_documents(pages)

    # 生成向量（embedding）並存入數據庫
    embedding = ModelScopeEmbeddings(model_id=MODEL_ID)
    db = Chroma.from_documents(
        documents=splits,
        embedding=embedding,
        persist_directory=PERSIST_DIRECTORY
    )
    # 持久化數據庫
    db.persist()
    return db

def querying(query, history):
    db = None
    if not os.path.exists(PERSIST_DIRECTORY):
        # 向量數據庫不存在則創建
        db = create_db()
    else:
        # 載入已有的向量數據庫
        embedding = ModelScopeEmbeddings(model_id=MODEL_ID)
        db = Chroma(persist_directory=PERSIST_DIRECTORY, embedding_function=embedding)

    # chat模型
    llm = AzureChatOpenAI(
        openai_api_version="2023-05-15",
        azure_deployment="gpt35-16k",
        model_version="0613",
        temperature=0
    )

    # chat緩存，用於保持聊天記錄
    memory = AnswerConversationBufferMemory(memory_key="chat_history", return_messages=True)

    # chat
    qa_chain = ConversationalRetrievalChain.from_llm(
      llm=llm,
      retriever=db.as_retriever(search_kwargs={"k": 7}),
      chain_type='stuff',
      memory=memory,
      return_source_documents=True,
  )
    result = qa_chain({"question": query})
    print(result)
    return result["answer"].strip()

# gradio
iface = gr.ChatInterface(
    fn = querying,
    chatbot=gr.Chatbot(height=1000),
    textbox=gr.Textbox(placeholder="邏輯是誰？", container=False, scale=7),
    title="三體問答機器人",
    theme="soft",
    examples=["簡述一下黑暗森林法則",
              "程心最後和誰在一起了？"],
    cache_examples=True,
    retry_btn="重試",
    undo_btn="撤回",
    clear_btn="清除",
    submit_btn="提交"
    )

iface.launch(share=True)

子文