RAG知識庫最佳化之Rerank應用

AiFly發表於2024-04-15

前面幾篇文章介紹了最基本的RAG實現，也瞭解了RAG的基本原理。最基本的RAG 流程為使用者Query問題，RAG應用在向量庫查詢出Context，然後將Query與 Context喂到LLM，LLM生成滿足Query的Response答案。
從這裡可以存在三元組： Query、Context 和 Response，如要想進一步提升RAG的效能還可在Context(上下文)上做文章，目前得到的Context很可能與Query(問題)之間的相關度不高，透過Rerank可以對Context進行評分排序，去掉分數較低相關度較低的資料從而避免其對上下文資訊的汙染。下圖為基本的RAG流程圖；

在未引入Rerank時可以看到從向量庫中得到的上下文中包含了大量的問題相關度較低的資訊，造成不相干資訊對上下文的干擾。下圖為未進行重排時的上下文；

Rerank模型

Rerank模型使用了bge-reranker-base模型和embedding模型不同，Reranker模型使用問題和文件作為輸入，直接輸出相似度而不是embedding。本文透過向Reranker輸入查詢和段落來獲得相關性分數。Reranker是基於交叉熵損失進行最佳化的相關性得分不受特定範圍的限制，bge-reranker-base適用於中英文的雙語場景。如果有更多語言需求可使用最新發布的BGE Re-Ranker v2.0模型。
Rerank模型可以在不犧牲準確性的情況下加速LLM的查詢並可能提高準確率，Rerank透過從上下文中刪除分數較低(相關性較低)的資料，Rerank透過重新排序上下文得以實現。下圖為引入Rerank後的RAG流程圖。

在引入Rerank後可以看到上下文資訊與問題的相關度提高了不少，不相干資訊對上下文的干擾基本解決。下圖為進行重排後的上下文；

Rerank模型使用

embedding_model_name = 'maidalun1020/bce-embedding-base_v1'
embedding_encode_kwargs = {
              'normalize_embeddings': True
            }
embed_model = HuggingFaceEmbeddings(
  model_name="./bge-base-zh-v1.5",
  encode_kwargs=embedding_encode_kwargs
)

reranker_args = {
             'model_name': '/mnt/d/software/dev/gpt/bge-reranker-base', 
             'top_n': 5
                 }
reranker = BgeRerank(model_name='/mnt/d/software/dev/gpt/bge-reranker-base',top_n=3)

# init documents
db =  Chroma(persist_directory='chroma_DB',embedding_function=embed_model)
retriever =db.as_retriever(search_type="similarity",search_kwargs={"k": 10})

compression_retriever = ContextualCompressionRetriever(
    base_compressor=reranker, base_retriever=retriever
)

qa_chain = RetrievalQA.from_chain_type(llm, retriever=compression_retriever
, chain_type_kwargs={"prompt": PROMPT})

引入Rerank模型後的retrieval引擎能夠去除上下文不相關的汙染資料、提供更精準的上下文資訊。重排後(Rerank)精準的上下文不僅可減少了token的使用量進而還可能提高LLM推理速度與準確率。

相關文章