RAG-Retrieval

使用統一的方式推理不同的RAG Reranker模型| 微調全鏈路的RAG檢索模型| 實驗結果| License

RAG-Retrieval 提供了全鏈路的RAG檢索微調(train)和推理(infer)以及蒸餾(distill)程式碼。

對於微調，支援微調任意開源的RAG檢索模型，包括向量（embedding、圖a）、遲交互式模型（colbert、圖d）、重新排序模型（cross encoder(bert)、圖c, llm-based reranker）。
對於推理，RAG-Retrieval專注於重新排序(reranker)，開發了一個輕量級的python庫rag-retrieval,提供統一的方式調用任意不同的RAG排序模型。
對於蒸餾，支持將基於LLM的reranker模型蒸餾到基於bert的reranker模型中。

ColBERT

社區交流

加入我們微信群聊

Roadmap

RAG-Retrieval Roadmap

使用統一的方式推理不同的RAG Reranker模型

為什麼要做Reranker模型的推理,甚至開發一個包？

排序模型是任何檢索架構的重要組成部分，也是RAG 的重要組成部分，但目前的現狀是：

開源的排序模型很多，在A場景表現好的模型，在B場景不一定表現好，很難知道該使用哪一個。
另外，新的排序模型不斷的出現，如今年3月BGE才發布的LLM Reranker，使用decoder-only的大模型來對段落重排序，非常有前景。
所有不同的排序模型，都傾向於自己開發一套庫來進行排序，這導致了更高的壁壘，新用戶需要熟悉每種排序模型的輸入和輸出，以及安裝各種不同的依賴。

rag-retrieval的特點

因此，RAG-Retrieval發展了一個輕量級的python函式庫rag-retrieval,提供統一的方式來呼叫任意不同的RAG排序模型，其有以下的特點。

支援多種排序模型：支援常見的開源排序模型(Cross Encoder Reranker,Decoder-Only 的LLM Reranker)
長doc友善：支援兩種不同的對於長doc的處理邏輯(最大長度截斷，切分取最大分值)。
益於擴充：如果有新的排序模型，使用者只需要繼承basereranker，並且實作rank以及comput_score函數即可。

安裝環境

 #为了避免自动安装的torch与本地的cuda不兼容，建议进行下一步之前先手动安装本地cuda版本兼容的torch。
pip install rag-retrieval

支持的Reranker模型

Cross Encoder Reranker

對於Cross Encoder Reranker，只要其使用transformers的AutoModelForSequenceClassification ，那麼就可以使用rag_retrieval的Reranker來進行推理。舉例如下。

bge系列的Cross Encoder模型，例如(BAAI/bge-reranker-base, BAAI/bge-reranker-large, BAAI/bge-reranker-v2-m3)
bce的Cross Encoder模型，例如(maidalun1020/bce-reranker-base_v1)

LLM Reranker

對於LLM Reranker，rag_retrieval的Reranker支援多種強大的LLM排序模型。也支援使用任意的LLM的chat模型來進行zero shot排序。舉例如下。

bge系列的LLM Reranker模型，例如(BAAI/bge-reranker-v2-gemma, BAAI/bge-reranker-v2-minicpm-layerwise, BAAI/bge-reranker-v2-m3 )
也支援使用任意LLM的chat模型來進行zero shot排序

使用

rag-retrieval套件詳細的使用方法和注意事項可以參考Tutorial

我們做了大量的測試與下面原有推理框架對齊，詳見Tests，他們需要不同的模組來執行，rag_retrieval則使用了統一的介面。

如FlagEmbedding的FlagReranker,FlagLLMReranker,LayerWiseFlagLLMReranker。

如BCEmbedding 的RerankerModel

微調全鏈路的RAG檢索模型

安裝環境

conda create -n rag-retrieval python=3.8 && conda activate rag-retrieval
#为了避免自动安装的torch与本地的cuda不兼容，建议进行下一步之前先手动安装本地cuda版本兼容的torch。
pip install -r requirements.txt

向量（embedding）模型

支援微調任意開源的embedding模型（bge,m3e等等）
支援對兩種數據進行微調：
- query和正例（負例採用batch內隨機負例），
- query和正例以及難負例。（負例為對應的難負例，以及batch內隨機負例）

微調embedding模型流程,詳細的流程可參考模型目錄下的Tutorial。

 cd ./rag_retrieval/train/embedding
bash train_embedding.sh

遲交互式（colbert）模型

支援微調開源的bge-m3e模型中的colbert。
支持query和正例以及難負例。（負例為對應的難負例，以及batch內隨機負例）

微調colbert模型流程，詳細的流程可參考模型目錄下的Tutorial。

 cd ./rag_retrieval/train/colbert
bash train_colbert.sh

排序（reranker,cross encoder）模型

支援微調任意開源的reranker模型（例如，bge-rerank、bce-rerank等）。
支援兩種數據進行微調：
- query和doc的相關性為二分類（1代表相關、0代表不相關）
- query和doc的相關性為score。（蒸餾任務）

微調reranker模型流程，詳細的流程可參考模型目錄下的Tutorial。

 cd ./rag_retrieval/train/reranker
bash train_reranker.sh

實驗結果

reranker模型在MTEB Reranking 任務的結果

Model	Model Size(GB)	T2Reranking	MMarcoReranking	CMedQAv1	CMedQAv2	Avg
bge-reranker-base	1.11	67.28	35.46	81.27	84.10	67.03
bce-reranker-base_v1	1.11	70.25	34.13	79.64	81.31	66.33
rag-retrieval-reranker	0.41	67.33	31.57	83.54	86.03	67.12

其中，rag-retrieval-reranker是我們使用RAG-Retrieval程式碼在hfl/chinese-roberta-wwm-ext模型上訓練所得，訓練資料使用bge-rerank模型的訓練資料.

colbert模型在MTEB Reranking 任務的結果

Model	Model Size(GB)	Dim	T2Reranking	MMarcoReranking	CMedQAv1	CMedQAv2	Avg
bge-m3-colbert	2.24	1024	66.82	26.71	75.88	76.83	61.56
rag-retrieval-colbert	0.41	1024	66.85	31.46	81.05	84.22	65.90

其中，rag-retrieval-colbert是我們使用RAG-Retrieval程式碼在hfl/chinese-roberta-wwm-ext模型上訓練所得，訓練資料使用bge-rerank模型的訓練資料.

以領域內資料微調開源的BGE系列模型

Model	T2ranking
bge-v1.5-embedding	66.49
bge-v1.5-embedding finetune	67.15	+0.66
bge-m3-colbert	66.82
bge-m3-colbert finetune	67.22	+0.40
bge-reranker-base	67.28
bge-reranker-base finetune	67.57	+0.29