LLM學習(四)——構建 RAG 應用

zddkk發表於2024-04-24

4.1 接入Langchain

首先我們匯入所有需要的庫檔案

import google.generativeai as genai
import os
from langchain_google_genai import ChatGoogleGenerativeAI
from langchain_google_genai import GoogleGenerativeAIEmbeddings
from langchain.prompts import PromptTemplate
from langchain.chains import create_retrieval_chain
from langchain.chains.combine_documents import create_stuff_documents_chain
import sys
from langchain.vectorstores.chroma import Chroma
import gradio as gr

我們使用langchain_google_genai中的ChatGoogleGenerativeAI匯入模型另外Langchain現在還沒有把genai併入所有要想安裝相應的包還需要執行
!pip install -U langchain-google-genai
另外出來使用 API 金鑰設定 GOOGLE_API_KEY 的環境變數
你還可以顯式的使用 google_api_key kwarg 將您的 API 金鑰傳遞給 ChatGoogle 建構函式
llm = ChatGoogleGenerativeAI(model="gemini-pro",google_api_key=api_key)

# 這裡換成自己的代理埠
os.environ['http_proxy'] = 'http://127.0.0.1:10809'
os.environ['https_proxy'] = 'http://127.0.0.1:10809'
GOOGLE_API_KEY="AIzaSyBZFSrCWg7Ow2D3I2rMbf37qXp1SlF9T5k"
os.environ["GOOGLE_API_KEY"] =GOOGLE_API_KEY
llm = ChatGoogleGenerativeAI(model="gemini-pro")

result = llm.invoke("Write a ballad about LangChain")
print(type(result))
print(result.content)

下面的圖片是相應的結果

LLM學習(四)——構建 RAG 應用

4.2 向量庫的搭建

4.2.1 載入向量資料庫

path = '輸入自己的持久化向量庫的位置記得改'
embeddings = GoogleGenerativeAIEmbeddings(model="models/embedding-001",task_type="retrieval_document")
vectordb=Chroma(
    persist_directory=path,
    embedding_function=embeddings
)

print(f"向量庫中儲存的數量：{vectordb._collection.count()}")

question = "什麼是prompt engineering?"
docs = vectordb.similarity_search(question,k=3)
print(f"檢索到的內容數：{len(docs)}")
for i, doc in enumerate(docs):
    print(f"檢索到的第{i}個內容: \n {doc.page_content}", end="\n-----------------------------------------------------\n")

下面是匯入我的向量庫的結果

LLM學習(四)——構建 RAG 應用

4.2.2 構建檢索問答鏈

PromptTemplate函式是 Langchain 官方提供語言模型的提示模板，提示模板由字串模板組成。它接受來自使用者的一組引數，這些引數可用於生成語言模型的提示，而且可以使用 f-strings（預設）或 jinja2 語法對模板進行格式化。下面的Code中的{context}和{input}就是相應的需要輸入的引數。然後使用 from_template 把我們預定義的模板template轉換成PromptTemplate物件。接著我們使用 create_stuff_documents_chain 函式去建立一個鏈用於把輸入的文件列表也就是 prompt 格式化然後傳給 LLM。
這樣我們就構建好了檢索問答鏈 retrieval_chain
接著透過 create_retrieval_chain 去構建檢索器 並且把

template = """
You are a helpful AI assistant.
Answer based on the context provided. 
context: {context}
input: {input}
answer:
"""
retriever=vectordb.as_retriever()
prompt = PromptTemplate.from_template(template)
combine_docs_chain = create_stuff_documents_chain(llm, prompt)
retrieval_chain = create_retrieval_chain(retriever, combine_docs_chain)

相關文章