檢索增強型生成 (Retrieval Augmented Generation, RAG) 是一種用於從您自己的文件中檢索信息的技術。它涉及一個兩步驟的過程:

1. 知識庫創建: 文件被分塊成較小的子文件,並為每個塊計算嵌入。這些嵌入和原始塊通常存儲在向量存儲中。

2. 檢索和生成: 當有新的用戶查詢時,會計算查詢的嵌入。通過比較查詢嵌入與存儲的塊嵌入,檢索最相關的塊。然後將檢索到的塊附加到原始查詢上,並輸入到語言模型 (LLM) 中生成最終響應。

這種設置被稱為「檢索增強型生成」管道,因為 LLM 的生成是通過從知識庫檢索到的相關信息增強的。

雖然有像 Langchain 和 LlamaIndex 這樣的框架可以簡化 RAG 管道的實現,但核心組件可以僅使用 Python、嵌入模型和 LLM 來構建。這種方法可以更好地理解底層概念,並允許更多地定制管道。

在構建檢索增強型生成 (RAG) 系統時,一個關鍵步驟是將輸入文件分塊成更小、更易管理的片段。在這個例子中,我們將輸入文件 (一篇維基百科文章) 分塊成段落。

這種方法的理由是,段落通常代表可以有效檢索和使用來增強用戶查詢的信息單元。通過以這種方式分解文件,我們可以更好地識別要包含在最終響應中的最相關部分。

這一步的代碼如下:

# 將文件分塊成段落
chunks = [paragraph.strip() for paragraph in text.split('\n') if paragraph.strip()]

這裡,我們使用一種簡單的方法,將輸入文本按換行符分割,以提取個別段落。然後我們去除每個段落開頭和結尾的空白字符,以確保乾淨的表示。

結果的 chunks 列表包含個別段落,這些段落可以嵌入並存儲在向量存儲中。這使我們能夠根據用戶的查詢有效地檢索最相關的段落,並將它們包含在由語言模型生成的最終響應中。

為了嵌入文件塊和用戶查詢,我們將使用 Sentence Transformers 庫。這個庫提供了預訓練的模型,可以為文本生成高質量的嵌入。

首先,我們加載嵌入模型:

from sentence_transformers import SentenceTransformer
embedding_model = SentenceTransformer('all-mpnet-base-v2')

接下來,我們計算文件塊的嵌入:

chunk_embeddings = embedding_model.encode(chunks, normalize_embeddings=True)

這裡,chunks 是文件中的文本塊列表。normalize_embeddings=True 選項確保嵌入被歸一化,這對後續的相似性計算很重要。

要獲得用戶查詢的嵌入,我們只需運行:

query_embedding = embedding_model.encode([user_query], normalize_embeddings=True)[0]

現在,我們有了文件塊和用戶查詢的嵌入,可以用於檢索步驟。

為了根據用戶查詢檢索最相關的塊,我們首先需要計算查詢嵌入與文件塊嵌入之間的餘弦相似度。以下是我們如何做到這一點:

1. 使用與文件塊相同的嵌入模型計算用戶查詢的嵌入。
2. 計算查詢嵌入與每個文件塊嵌入之間的點積。這給我們提供了餘弦相似度分數。
3. 根據餘弦相似度分數對文件塊進行排序,並選擇前 K 個最相關的塊。
4. 檢索所選塊的文本內容,作為語言模型的上下文。

關鍵步驟是:

1. 計算查詢嵌入
2. 計算餘弦相似度分數
3. 排序並選擇前 K 個塊
4. 檢索塊文本內容

這種簡單的方法允許我們有效地從文件集合中檢索最相關的信息,以增強用戶的查詢,然後將其傳遞給語言模型進行生成。

為了使用檢索到的塊和用戶查詢生成響應,我們首先創建一個包含相關塊和用戶查詢的提示。然後我們使用 OpenAI GPT-4 模型根據這個提示生成一個響應。

以下是代碼:

# 獲取用戶查詢
user_query = "What is the capital of France?"

# 根據用戶查詢檢索最相關的塊
top_chunk_ids = [6, 8, 5]
top_chunks = [chunks[i] for i in top_chunk_ids]

# 創建提示
prompt = "Use the following context to answer the question at the end. If you don't know the answer, say that you don't know and do not try to make up an answer.\n\nContext:\n"
for chunk in top_chunks:
    prompt += chunk + "\n\n"
prompt += f"\nQuestion: {user_query}"

# 使用 OpenAI GPT-4 生成響應
response = openai_client.chat.create(
    model="gpt-4",
    messages=[
        {"role": "user", "content": prompt}
    ],
    max_tokens=1024,
    n=1,
    stop=None,
    temperature=0.7,
).choices[0].message.content

print(response)

在這段代碼中,我們首先根據用戶查詢檢索最相關的三個塊。然後我們創建一個包含這些塊和用戶查詢的提示。最後,我們使用 OpenAI GPT-4 模型根據這個提示生成一個響應。

生成的響應將是一個簡潔而相關的答案,利用了檢索到的塊中的信息。

在這個教程中,我們學習了如何使用僅 10 行 Python 代碼構建一個帶有文檔檢索系統的完整聊天系統。我們介紹了檢索增強型生成 (RAG) 管道的關鍵組件,包括知識庫創建、文檔分塊、嵌入、檢索和使用大型語言模型進行生成。

關鍵要點是:

1. 您可以在不需要複雜框架(如 Langchain 或 LlamaIndex)的情況下實現基本的 RAG 管道。純 Python 和少量庫就足夠了。
2. 根據結構(如段落)對文檔進行分塊是大多數用例的簡單而有效的策略。
3. 嵌入文檔塊和用戶查詢,然後計算相似度分數,可以讓您檢索最相關的信息來增強用戶的查詢。
4. 將檢索到的塊與用戶的查詢集成,並輸入到大型語言模型中,可以生成相關和有信息量的響應。

雖然這個例子提供了一個堅實的基礎,但還有很多機會來構建更強大和先進的 RAG 系統。像 Langchain 和 LlamaIndex 這樣的框架在與各種向量存儲和語言模型集成時可能會很有幫助。但是,從一個純 Python 實現開始可以幫助您更好地理解 RAG 管道的核心概念和組件。

如果您對探索更高級的 RAG 技術感興趣,我建議您查看我的課程「RAG Beyond Basics」,它深入探討了構建複雜、生產就緒的 RAG 系統。