RAG高階技巧-如何實現窗口上下文檢索

2024年2月6日 30点热度 0人点赞

在本文中,我們將介紹一種提高RAG(Retrieval-Augmented Generation)模型檢索效果的高階技巧,即窗口上下文檢索。我們將首先回顧一下基礎RAG的檢索流程和存在的問題,然後介紹窗口上下文檢索的原理和實現方法,最後通過一個實例展示其效果。

本文為一個系列,之前內容沒有看過的小夥伴可以點擊鏈接查看:「鏈接」

基礎RAG存在的問題及解決方案

基礎RAG檢索流程

RAG是一種結合了檢索和生成的AI應用落地的方案,它可以根據給定的問題生成回答,同時利用外部知識庫(例如維基百科)來增強生成的質量和多樣性。RAG的核心思想是將問題和知識庫中的文檔進行匹配,然後將匹配到的文檔作為生成模型的輸入,從而生成更加相關和豐富的回答。

RAG的檢索流程可以分為以下幾個步驟:

  • load:加載文檔,將各種格式的文件加載後轉化為文檔,例如將pdf加載為文本數據,或者將表格轉換為多個鍵值對。
  • split:將文檔拆分為適合向量存儲的較小單元,以便於與向量存儲,以及檢索時的文檔匹配,例如將“我是kxc。我喜歡唱跳,rap,和籃球。”拆分為“我是kxc。”和“我喜歡唱跳,rap,和籃。”兩個數據分塊(一般稱之為chunk)。
  • embedding:將文檔用向量表示,例如使用BERT或TF-IDF等模型進行向量化。
  • store: 將向量化後的數據分塊,存入向量數據庫。
  • retrive:根據問題和文檔的向量,計算它們之間的相似度,然後根據相似度的高低,選擇最相關的文檔作為檢索結果,例如使用餘弦相似度或點積等度量進行排序。
  • query:將檢索到的文檔作為生成模型的輸入,根據問題生成回答,例如使用GPT-3或T5等模型進行生成。

基礎RAG存在的問題

基礎RAG的檢索流程雖然簡單,但是也存在一些問題,主要是在split和retrive兩個步驟中。這些問題會影響RAG的檢索效果,從而導致生成的回答不準確或不完整。

  • split拆分的塊太大,在retrive時,同一塊中非相關的內容就越多,對問題的檢索匹配度影響越大,會導致檢索的不準確。例如,如果我們將維基百科中的一篇文章作為一個文檔,那麼這個文檔可能包含很多不同的主題和細節,與問題的相關性會很低。如果我們將這個文檔作為檢索結果,那麼生成模型可能會從中提取出一些無關或錯誤的信息,從而影響回答的質量。
  • split拆分的塊太小,檢索的匹配度會提高,然而在最後的query環節,提供給llm使用的信息會由於缺少上下文的信息支撐,導致回答不準確。例如,如果我們將維基百科中的一篇文章拆分為多個句子,那麼每個句子可能隻包含一小部分的信息,與問題的相關性會很高。如果我們將這些句子作為檢索結果,那麼生成模型可能會從中提取出一些有用的信息,但是也可能會忽略一些重要的上下文信息,從而影響回答的完整性。

解決方案-窗口上下文檢索

解決這個問題一般采取的方案是,在split拆分時,盡量將文本切分到最小的語義單元。retrive時,不直接使用匹配到doc,而是通過匹配到的doc拓展其上下文內容後整合到一起,在投遞到LLM使用。這樣,既可以提高檢索的精度,又可以保留上下文的完整性,從而提高生成的質量和多樣性。

具體來說,這種方案的實現步驟如下:

  • 在split拆分時,將文本切分為最小的語義單元,例如句子或段落,並給每個單元分配一個唯一的編號,作為其在文本中的位置信息。
  • 在retrive檢索時,根據問題和文檔的向量,計算它們之間的相似度,然後選擇最相關的文檔作為檢索結果,同時記錄下它們的編號。
  • 在query生成時,根據檢索結果的編號,從文本中獲取它們的上下文信息,例如前後若幹個單元,然後將它們拼接成一個完整的文檔,作為生成模型的輸入,根據問題生成回答。

窗口上下文檢索實踐

上下文檢索實現思路

我們從最終要實現的目標著手,也就是在retrive時要能通過匹配到doc拓展出與這個doc內容相關的上下文。要想實現這個目標,我們就必須建立每個doc與其上下文的關聯關系。這個關系的建立其實十分簡單,隻需要按順序給拆分出來的每個doc進行編號,在檢索時通過當前文檔的編號就能匹配到相關上下文doc的編號,進而獲取上下文的內容。

基於chroma向量庫的代碼實踐

想要實踐以上思路,需要在split環節基於文檔順序,將文檔編碼寫入元數據。在檢索時,則通過元數據中的順序分塊編碼來查找上下文。具體的代碼如下:

  • 1.split時對分塊編碼並寫入元數據
import bs4,uuid
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import WebBaseLoader
from langchain_community.vectorstores import Chroma
from langchain_openai import  OpenAIEmbeddings
# Load, chunk and index the contents of the blog.
loader = WebBaseLoader(
    web_paths=("https://lilianweng.github.io/posts/2023-06-23-agent/",),
    bs_kwargs=dict(
        parse_only=bs4.SoupStrainer(
            class_=("post-content", "post-title", "post-header")
        )
    ),
)
doc = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=200)
docs = text_splitter.split_documents(doc)
# 這裡給每個docs片段的metadata裡註入file_id
file_id = uuid.uuid4().hex
chunk_id_counter = 0
for doc in docs:
    doc.metadata["file_id"] = file_id
    doc.metadata["chunk_id"] = f'{file_id}_{chunk_id_counter}'  # 添加chunk_id到metadata
    chunk_id_counter  = 1
    for key,value in doc.metadata.items():
        if not isinstance(value, (str, int, float, bool)):
            doc.metadata[key] = str(value)
vectorstore = Chroma.from_documents(documents=splits, embedding=OpenAIEmbeddings())
  • 2.retrive時通過元數據中的順序分塊編碼來查找上下文
def expand_doc(group):
    new_cands = []
    group.sort(key=lambda x: int(x.metadata['chunk_id'].split('_')[-1]))
    id_set = set()
    file_id = group[0].metadata['file_id']
    group_scores_map = {}
    # 先找出該文件所有需要搜索的chunk_id
    cand_chunks = []
    for cand_doc in group:
        current_chunk_id = int(cand_doc.metadata['chunk_id'].split('_')[-1])
        group_scores_map[current_chunk_id] = cand_doc.metadata['score']
        for i in range(current_chunk_id - 200, current_chunk_id   200):
            need_search_id = file_id   '_'   str(i)
            if need_search_id not in cand_chunks:
                cand_chunks.append(need_search_id)
    where = {"chunk_id": {"$in": cand_chunks}}
    ids,group_relative_chunks = get(where)
    group_chunk_map = {int(item.metadata['chunk_id'].split('_')[-1]): item.page_content for item in group_relative_chunks}
    group_file_chunk_num = list(group_chunk_map.keys())
    for cand_doc in group:
        current_chunk_id = int(cand_doc.metadata['chunk_id'].split('_')[-1])
        doc = copy.deepcopy(cand_doc)
        id_set.add(current_chunk_id)
        docs_len = len(doc.page_content)
        for k in range(1, 200):
            break_flag = False
            for expand_index in [current_chunk_id   k, current_chunk_id - k]:
                if expand_index in group_file_chunk_num:
                    merge_content = group_chunk_map[expand_index]
                    if docs_len   len(merge_content) > CHUNK_SIZE:
                        break_flag = True
                        break
                    else:
                        docs_len  = len(merge_content)
                        id_set.add(expand_index)
            if break_flag:
                break
    id_list = sorted(list(id_set))
    id_lists = seperate_list(id_list)
    for id_seq in id_lists:
        for id in id_seq:
            if id == id_seq[0]:
                doc = Document(page_content=group_chunk_map[id],
                                metadata={"score": 0, "file_id": file_id})
            else:
                doc.page_content  = " "   group_chunk_map[id]
        doc_score = min([group_scores_map[id] for id in id_seq if id in group_scores_map])
        doc.metadata["score"] = doc_score
        new_cands.append(doc)
    return new_cands

總結

在本文中,我們介紹了提高RAG模型檢索效果的高階技巧-窗口上下文檢索。我們首先回顧了基礎RAG的檢索流程和存在的問題,然後介紹了窗口上下文檢索的原理和實現方法,最後通過一個實例展示了其效果。

我們希望這篇文章能夠幫助你理解和使用窗口上下文檢索這個高階技巧,從而提高你的RAG模型的檢索效果。如果你對這個技巧有任何疑問或建議,歡迎在評論區留言。謝謝你的閱讀。