從零開始編寫一個簡單的 RAG

最近，檢索增強生成 (Retrieval-Augmented Generation, RAG) 已成為人工智慧和大型語言模型 (LLM) 領域一個強大的正規化。RAG 將資訊檢索與文字生成相結合，透過整合外部知識源來增強語言模型的效能。這種方法在各種應用中都顯示出有希望的結果，例如問答、對話系統和內容生成。

在這篇博文中，我們將探討 RAG，並使用 Python 和 ollama 從零開始構建一個簡單的 RAG 系統。這個專案將幫助你理解 RAG 系統的關鍵元件，以及如何使用基本的程式設計概念來實現它們。

什麼是 RAG

首先，讓我們看一個沒有 RAG 的簡單聊天機器人系統

雖然聊天機器人可以根據其訓練資料集回答常見問題，但它可能無法獲取最新或特定領域的知識。

一個現實世界的例子是問 ChatGPT“我媽媽的名字是什麼？”。ChatGPT 無法回答這個問題，因為它無法訪問外部知識，比如你的家庭成員資訊。

為了解決這個限制，我們需要向模型提供外部知識（在這個例子中，是一個家庭成員姓名的列表）。

一個 RAG 系統由兩個關鍵元件組成

• 一個檢索模型，它從外部知識源獲取相關資訊，知識源可以是資料庫、搜尋引擎或任何其他資訊庫。
• 一個語言模型，它根據檢索到的知識生成回應。

實現 RAG 的方法有多種，包括圖 RAG (Graph RAG)、混合 RAG (Hybrid RAG) 和分層 RAG (Hierarchical RAG)，我們將在本文末尾討論這些。

簡單的 RAG

讓我們建立一個簡單的 RAG 系統，它可以從一個預定義的資料集中檢索資訊，並根據檢索到的知識生成回應。該系統將包括以下元件

1. 嵌入模型：一個預訓練的語言模型，將輸入文字轉換為嵌入（embedding）——即捕捉語義的向量表示。這些向量將用於在資料集中搜索相關資訊。
2. 向量資料庫：一個用於儲存知識及其對應嵌入向量的系統。雖然有許多向量資料庫技術，如 Qdrant、Pinecone 和 pgvector，但我們將從零開始實現一個簡單的記憶體資料庫。
3. 聊天機器人：一個根據檢索到的知識生成回應的語言模型。這可以是任何語言模型，如 Llama、Gemma 或 GPT。

索引階段

索引階段是建立 RAG 系統的第一步。它涉及將資料集（或文件）分解成小的資料塊 (chunks)，併為每個資料塊計算一個向量表示，以便在生成過程中能被高效地搜尋。

每個資料塊的大小可以根據資料集和應用而變化。例如，在文件檢索系統中，每個資料塊可以是一個段落或一個句子。在對話系統中，每個資料塊可以是一輪對話。

在索引階段之後，每個資料塊及其對應的嵌入向量將被儲存在向量資料庫中。下面是一個索引後向量資料庫可能的樣子：

資料塊	嵌入向量
義大利和法國生產了世界上超過 40% 的葡萄酒。	\[0.1, 0.04, -0.34, 0.21, ...\]
印度的泰姬陵完全由大理石建成。	\[-0.12, 0.03, 0.9, -0.1, ...\]
世界上 90% 的淡水在南極洲。	\[-0.02, 0.6, -0.54, 0.03, ...\]
...	...

嵌入向量稍後可用於根據給定的查詢檢索相關資訊。可以把它想象成 SQL 的 WHERE 子句，但我們不是透過精確的文字匹配來查詢，而是可以根據它們的向量表示來查詢一組資料塊。

為了比較兩個向量的相似度，我們可以使用餘弦相似度、歐幾里得距離或其他距離度量。在這個例子中，我們將使用餘弦相似度。以下是兩個向量 A 和 B 的餘弦相似度公式：

如果你不熟悉上面的公式，別擔心，我們將在下一節中實現它。

檢索階段

在下圖中，我們將以一個來自使用者的給定輸入查詢為例。然後我們計算查詢向量來表示該查詢，並將其與資料庫中的向量進行比較，以找到最相關的資料塊。

向量資料庫返回的結果將包含與查詢最相關的 N 個數據塊。這些資料塊將被聊天機器人用來生成回應。

讓我們開始編碼

在這個例子中，我們將用 Python 編寫一個簡單的 RAG 實現。

為了執行模型，我們將使用 ollama，這是一個命令列工具，可以讓你執行來自 Hugging Face 的模型。有了 ollama，你不需要訪問伺服器或雲服務來執行模型。你可以直接在你的電腦上執行模型。

對於模型，我們使用以下這些：

• 嵌入模型：hf.co/CompendiumLabs/bge-base-en-v1.5-gguf
• 語言模型：hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF

至於資料集，我們將使用一個關於貓的事實的簡單列表。每個事實在索引階段將被視為一個數據塊。

下載 ollama 和模型

首先，讓我們從專案網站安裝 ollama：ollama.com

安裝後，開啟一個終端並執行以下命令以下載所需的模型

ollama pull hf.co/CompendiumLabs/bge-base-en-v1.5-gguf
ollama pull hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF

如果你看到以下輸出，說明模型已成功下載

pulling manifest
...
verifying sha256 digest
writing manifest
success

在繼續之前，為了在 python 中使用 ollama，我們還需要安裝 ollama 包

pip install ollama

載入資料集

接下來，建立一個 Python 指令碼並將資料集載入到記憶體中。該資料集包含一系列關於貓的事實，這些事實將在索引階段用作資料塊。

你可以從這裡下載示例資料集。下面是載入資料集的示例程式碼

dataset = []
with open('cat-facts.txt', 'r') as file:
  dataset = file.readlines()
  print(f'Loaded {len(dataset)} entries')

實現向量資料庫

現在，讓我們來實現向量資料庫。

我們將使用 ollama 的嵌入模型將每個資料塊轉換為一個嵌入向量，然後將資料塊及其對應的向量儲存在一個列表中。

下面是一個計算給定文字嵌入向量的示例函式：

import ollama

EMBEDDING_MODEL = 'hf.co/CompendiumLabs/bge-base-en-v1.5-gguf'
LANGUAGE_MODEL = 'hf.co/bartowski/Llama-3.2-1B-Instruct-GGUF'

# Each element in the VECTOR_DB will be a tuple (chunk, embedding)
# The embedding is a list of floats, for example: [0.1, 0.04, -0.34, 0.21, ...]
VECTOR_DB = []

def add_chunk_to_database(chunk):
  embedding = ollama.embed(model=EMBEDDING_MODEL, input=chunk)['embeddings'][0]
  VECTOR_DB.append((chunk, embedding))

在這個例子中，為簡單起見，我們將資料集中的每一行視為一個數據塊。

for i, chunk in enumerate(dataset):
  add_chunk_to_database(chunk)
  print(f'Added chunk {i+1}/{len(dataset)} to the database')

實現檢索函式

接下來，讓我們實現檢索函式，它接受一個查詢並根據餘弦相似度返回前 N 個最相關的資料塊。我們可以想象，兩個向量之間的餘弦相似度越高，它們在向量空間中就越“接近”。這意味著它們在意義上更相似。

下面是計算兩個向量之間餘弦相似度的示例函式：

def cosine_similarity(a, b):
  dot_product = sum([x * y for x, y in zip(a, b)])
  norm_a = sum([x ** 2 for x in a]) ** 0.5
  norm_b = sum([x ** 2 for x in b]) ** 0.5
  return dot_product / (norm_a * norm_b)

現在，讓我們實現檢索函式：

def retrieve(query, top_n=3):
  query_embedding = ollama.embed(model=EMBEDDING_MODEL, input=query)['embeddings'][0]
  # temporary list to store (chunk, similarity) pairs
  similarities = []
  for chunk, embedding in VECTOR_DB:
    similarity = cosine_similarity(query_embedding, embedding)
    similarities.append((chunk, similarity))
  # sort by similarity in descending order, because higher similarity means more relevant chunks
  similarities.sort(key=lambda x: x[1], reverse=True)
  # finally, return the top N most relevant chunks
  return similarities[:top_n]

生成階段

在這個階段，聊天機器人將根據上一步檢索到的知識生成一個回應。這隻需將資料塊新增到將作為聊天機器人輸入的提示 (prompt) 中即可完成。

例如，可以像這樣構建一個提示：

input_query = input('Ask me a question: ')
retrieved_knowledge = retrieve(input_query)

print('Retrieved knowledge:')
for chunk, similarity in retrieved_knowledge:
  print(f' - (similarity: {similarity:.2f}) {chunk}')

instruction_prompt = f'''You are a helpful chatbot.
Use only the following pieces of context to answer the question. Don't make up any new information:
{'\n'.join([f' - {chunk}' for chunk, similarity in retrieved_knowledge])}
'''

然後我們使用 ollama 來生成回應。在這個例子中，我們將使用 instruction_prompt 作為系統訊息。

stream = ollama.chat(
  model=LANGUAGE_MODEL,
  messages=[
    {'role': 'system', 'content': instruction_prompt},
    {'role': 'user', 'content': input_query},
  ],
  stream=True,
)

# print the response from the chatbot in real-time
print('Chatbot response:')
for chunk in stream:
  print(chunk['message']['content'], end='', flush=True)

整合所有部分

你可以在這個檔案中找到最終程式碼。要執行程式碼，請將其儲存為名為 `demo.py` 的檔案並執行以下命令：

python demo.py

你現在可以向聊天機器人提問了，它將根據從資料集中檢索到的知識生成回應。

Ask me a question: tell me about cat speed
Retrieved chunks: ...
Chatbot response:
According to the given context, cats can travel at approximately 31 mph (49 km) over a short distance. This is their top speed.

可改進的空間

到目前為止，我們已經用一個小資料集實現了一個簡單的 RAG 系統。然而，仍然存在許多侷限性：

• 如果問題同時涉及多個主題，系統可能無法提供好的答案。這是因為系統僅根據查詢與資料塊的相似性來檢索資料塊，而沒有考慮查詢的上下文。
  解決方案可以是讓聊天機器人根據使用者的輸入編寫自己的查詢，然後根據生成的查詢檢索知識。我們也可以使用多個查詢來檢索更多相關資訊。
• 前 N 個結果是根據餘弦相似度返回的。這並不總能得到最好的結果，特別是當每個資料塊包含大量資訊時。
  為了解決這個問題，我們可以使用重排模型 (reranking model) 來根據與查詢的相關性對檢索到的資料塊進行重新排序。
• 資料庫儲存在記憶體中，對於大型資料集可能不具備可擴充套件性。我們可以使用更高效的向量資料庫，如 Qdrant、Pinecone、pgvector。
• 我們目前將每個句子視為一個數據塊。對於更復雜的任務，我們可能需要使用更復雜的技術來將資料集分解成更小的資料塊。我們也可以在將每個資料塊新增到資料庫之前對其進行預處理。
• 本例中使用的語言模型是一個只有 1B 引數的簡單模型。對於更復雜的任務，我們可能需要使用更大的語言模型。

其他型別的 RAG

在實踐中，實現 RAG 系統的方法有很多。以下是一些常見的 RAG 系統型別：

• 圖 RAG (Graph RAG)：在這種型別的 RAG 中，知識源被表示為一個圖，其中節點是實體，邊是實體之間的關係。語言模型可以遍歷圖來檢索相關資訊。關於這類 RAG 有許多活躍的研究。這裡是關於圖 RAG 的論文集。
• 混合 RAG (Hybrid RAG)：一種結合了知識圖譜 (KGs) 和向量資料庫技術以改進問答系統的 RAG。要了解更多資訊，你可以閱讀這篇論文。
• 模組化 RAG (Modular RAG)：一種超越了基本的“檢索後生成”過程的 RAG，它採用路由、排程和融合機制來建立一個靈活且可重新配置的框架。這種模組化設計允許各種 RAG 模式（線性、條件、分支和迴圈），從而實現更復雜和適應性強的知識密集型應用。要了解更多資訊，你可以閱讀這篇論文。

關於其他型別的 RAG，你可以參考Rajeev Sharma 的這篇文章。

結論

RAG 代表了使語言模型知識更淵博、更準確方面的一個重大進步。透過從零開始實現一個簡單的 RAG 系統，我們探討了嵌入、檢索和生成的基本概念。雖然我們的實現是基礎的，但它展示了為生產環境中使用的更復雜的 RAG 系統提供動力的核心原則。

擴充套件和改進 RAG 系統的可能性是巨大的，從實現更高效的向量資料庫到探索如圖 RAG 和混合 RAG 等高階架構。隨著該領域的不斷發展，RAG 仍然是利用外部知識增強 AI 系統，同時保持其生成能力的關鍵技術。

參考文獻

• https://arxiv.org/abs/2005.11401
• https://aws.amazon.com/what-is/retrieval-augmented-generation/
• https://github.com/varunvasudeva1/llm-server-docs
• https://github.com/ollama/ollama/blob/main/docs
• https://github.com/ollama/ollama-python
• https://www.pinecone.io/learn/series/rag/rerankers/
• https://arxiv.org/html/2407.21059v1
• https://newsletter.armand.so/p/comprehensive-guide-rag-implementations