RAG 架構概述
RAG 是為解決大型語言模型在上下文過長時容易失去焦點並開始胡言亂語的問題而創建的架構。其基本思想如下:
-
首先,將文件分解成許多小塊。
-
然後,對每個小塊進行嵌入。
-
接著,將所有嵌入存儲在向量資料庫中。
-
當用戶提出問題時,在資料庫中找到與問題語義相似的小塊,並將它們一起發送給大型語言模型。
若想先了解 RAG 的基本原理,可查看作者的視頻,該視頻在 10 分鐘內清晰地解釋了 RAG 的核心概念,視頻描述中會提供連結。
選擇獨特的文本
為防止 AI 偷偷使用自己的知識庫,本次使用一篇它絕對沒見過的傑作:《關於令狐沖轉世為黏液並向世界獻上美麗爆炸》。這篇文章講述了令狐沖轉世到另一個世界成為黏液,憑藉前世記憶成功掌握爆炸法術的故事,是一個扮豬吃老虎的勵志故事。有興趣的話,可在視頻描述中查看連結自行閱讀。
分解文章
初始分解
第一步是將文章分解成許多小塊。這篇文章寫得很整齊,每段之間有兩個換行符,且每段的字數也相當均勻。因此,本次以兩個換行符作為分割依據,在 chunk.py 文件中編寫代碼。作者已預先寫好一個名為 read_data 的函數,用於將上述文章讀取為一個完整的字串。接著,添加一個名為 get_chunk 的分塊函數,其返回值是分塊後的文章,即一個字串列表。按照計劃,使用兩個換行符分割文章,第一個版本的分塊函數完成。
優化分塊
然而,存在一個小問題,這些章節標題也被分割成單獨的段落,且太短,將它們視為單獨的段落不太理想。所以,稍微修改一下邏輯。如果一個段落以井號開頭,就將它與後面的正文合併。作者展示了具體實現,這裡不詳細解釋代碼邏輯,都是相對簡單的字串操作。再次運行,可見標題已與正文連接,目前效果相當不錯,分塊函數基本完成。
此外,除了手動編寫分塊函數,網上還有許多現成的分塊演算法,例如 LangChain 的 RecursiveCharacterTextSplitter 就是一個非常強大的分塊演算法,有興趣的朋友可自行研究。本次只是進行基本演示,不會引入額外的複雜性。
嵌入與存儲
準備工作
完成文章分塊後,第二步是對每個小塊進行嵌入並存儲在向量資料庫中。在 embed.py 文件中編寫此代碼。作者選擇的向量資料庫是 ChromaDB,因為它相對容易使用。對於嵌入模型,選擇了 Google 的嵌入模型。首先安裝依賴項,由於使用 Google 的嵌入模型,還需要一個額外的 API KEY,將其放置在名為 GOOGLE_API_KEY 的環境變數中,作者已事先配置好並打印出來查看。
編寫嵌入函數
接著,編寫一個嵌入函數,命名為 embed。該函數接受一段文本作為參數,並返回該文本對應的嵌入,即一個浮點數數組。調用 Gemini 嵌入模型的介面,'model' 參數是嵌入模型的名稱,這裡使用的是 gemini-embedding-exp-03-07;'contents' 參數是需要嵌入的文本內容。
Google 的嵌入模型比較特殊,將嵌入分為兩類:一類用於存儲,一類用於查詢。例如,有兩段文本“王喜歡吃瓜”和“王愛好吃瓜”,這兩句話意思大致相同,它們的嵌入位置會比較接近。但如果後續問題是“王喜歡吃什麼?”,雖然這句話與前兩句語義相關,但形式差異較大,它們的嵌入距離可能不會那麼近。所以,Google 的做法是,在進行嵌入時,需要明確告訴它這段文本的目的是存儲還是查詢。也就是說,存儲“王喜歡吃瓜”和“王愛好吃瓜”時,需要使用存儲模式;查詢“王喜歡吃什麼?”時,需要使用查詢模式。這樣 Gemini 就能神奇地將這兩句看似遙遠的句子在嵌入向量空間中拉近。具體實現方式 Google 並未透露,我們只需遵循其規則。
需要注意的是,除了 Google,大多數其他嵌入模型沒有存儲和查詢的區分,它們使用相同的介面。因此,在嵌入函數中添加一個參數,以指示此嵌入是用於存儲還是查詢。接著,使用 'config' 參數將此信息傳遞給 Gemini 模型,'task_type' 參數存儲時為 RETRIEVAL_DOCUMENT,查詢時為 RETRIEVAL_QUERY,具體值可參考官方 Gemini 文檔,視頻描述中會提供連結。最後,直接返回嵌入結果。這裡為了演示目的,直接使用 'assert' 檢查返回值,在實際開發中,需要實現更強大的錯誤處理。
測試嵌入函數
然後編寫一個主函數進行測試,get_chunks 函數是剛才編寫的分塊函數,但似乎拼錯了函數名,進行修正。接著,打印第一個小塊的嵌入,看看它的樣子。運行代碼,可見打印出的這一長串數字就是文章第一個小塊的嵌入結果。
存儲到向量資料庫
現在有了文章小塊和嵌入方法,接下來需要創建 ChromaDB 的實例,然後將每個文章小塊的嵌入和原始文本一一對應地存儲到向量資料庫中。首先創建向量資料庫的實例,在這段代碼中,指定了資料庫的存儲位置在當前目錄的 chroma.db 資料夾中。然後創建一個資料表,資料表的名稱可自行命名,這裡使用了令狐沖的拼音。
接著,編寫一個函數,依次對每個小塊進行嵌入。注意,這裡 embed 函數的'store' 參數設置為 True,因為現在準備將結果存儲到資料庫中。然後,使用 ChromaDB 的 upsert 方法將嵌入結果存儲到向量資料庫中。Chroma 要求為每條資料提供一個字串類型的 ID,這個 ID 實際上不是很有用,這裡只是使用小塊的索引作為 ID。'documents' 和 'embeddings' 分別是原始文本及其對應的嵌入。
最後,在主函數中調用 create_db 函數,將令狐沖的異世界之旅存儲到資料庫中。運行代碼,可見 create_db 函數會對每個小塊進行嵌入並存儲到資料庫中。執行後,當前目錄中會出現一個新的資料夾,即 chroma.db,這就是剛才創建的資料庫。
查詢函數
有了資料庫,終於可以編寫最重要的查詢函數,命名為 query_db。首先,函數有一個參數 'question',對應用戶想要問的問題。在函數體中,首先需要對這個問題進行嵌入。注意,由於 'question' 是用於查詢的,這裡的'store' 參數為 False。接著,使用 question_embedding 在向量資料庫中查找最相關的文章小塊,這裡將 n_results 參數設置為返回 5 個最相關的記錄。最後,直接返回檢索到的文本小塊。
查詢函數完成,進行測試。例如,問“令狐沖掌握了什麼樣的魔法?”,這次不需要再次調用 create_db 函數,因為之前已經初始化了資料庫,直接調用 query_db 進行查詢。運行代碼,可見已檢索到與問題最相關的內容小塊,例如有關“爆炸爆發風格”的小塊,以及“爆炸”的原型、“獨孤九劍”的意境、“吸星大法”等,整體效果相當不錯,基本可以看到 RAG 的工作原理。
與大型語言模型交互
最後一步是將問題“令狐沖掌握了什麼魔法?”和剛才從向量資料庫中檢索到的文本小塊一起發送給大型語言模型。首先,將這些信息拼接成一個提示,打印出提示看看它的樣子。接著,將這個整個長提示直接發送給大型語言模型,這裡使用的大型語言模型是 Gemini-Flash-2.5。由於 Gemini 返回的是一個相當複雜的資料結構,直接打印出整個返回值。
最後一次運行代碼,看看最終效果。可見 Gemini 非常認真地回答說,令狐沖掌握的不是魔法,而是一種叫做“爆炸爆發風格”的能量爆發。感謝 Gemini 及時糾正了問題中的錯誤。至此,一個完整的 RAG 架構編寫完成,完整代碼將放在視頻描述中。
作者強烈建議讀者按照示例自行嘗試編寫,因為對於代碼之類的東西,只有親自寫出來才會慢慢變成自己的。不僅會理解得更深入,還會在編寫過程中發現新的細節、新的問題和新的想法。雖然我們都不是令狐沖,不知道獨孤九劍,也不能閉上眼睛就掌握爆炸法術,但作者仍然非常願意先扔出手中的小火球看看,因為誰知道呢,也許有一天它真的會變成爆炸。這是程序員王,下次見。