Minicpm Reranker

🚀 MiniCPM-Reranker

MiniCPM-Reranker 是面壁智能與清華大學自然語言處理實驗室（THUNLP）、東北大學信息檢索小組（NEUIR）共同開發的中英雙語言文本重排序模型。它具備出色的中文、英文重排序能力，以及出色的中英跨語言重排序能力，為文本重排序任務提供了強大的支持。

✨ 主要特性

• 具備出色的中文、英文重排序能力。
• 擁有出色的中英跨語言重排序能力。

📦 安裝指南

環境要求

transformers==4.37.2
flash-attn>2.3.5

💻 使用示例

輸入格式

本模型支持指令，輸入格式如下：

<s>Instruction: {{ instruction }} Query: {{ query }}</s>{{ document }}

例如：

<s>Instruction: 為這個醫學問題檢索相關回答。Query: 咽喉癌的成因是什麼？</s>（文檔省略）

<s>Instruction: Given a claim about climate change, retrieve documents that support or refute the claim. Query: However the warming trend is slower than most climate models have forecast.</s>(document omitted)

也可以不提供指令，即採取如下格式：

<s>Query: {{ query }}</s>{{ document }}

我們在BEIR與C-MTEB/Retrieval上測試時使用的指令見 instructions.json，其他測試不使用指令。

示例腳本

Huggingface Transformers

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch
import numpy as np

model_name = "openbmb/MiniCPM-Reranker"
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.padding_side = "right"

model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True, torch_dtype=torch.float16).to("cuda")
# You can also use the following code to use flash_attention_2
# model = AutoModelForSequenceClassification.from_pretrained(model_name, trust_remote_code=True,attn_implementation="flash_attention_2", torch_dtype=torch.float16).to("cuda")

model.eval()

@torch.no_grad()
def rerank(input_query, input_docs):
    tokenized_inputs = tokenizer([[input_query, input_doc] for input_doc in input_docs], return_tensors="pt", padding=True, truncation=True, max_length=1024) 

    for k in tokenized_inputs:
      tokenized_inputs [k] = tokenized_inputs[k].to("cuda")

    outputs = model(**tokenized_inputs)
    score = outputs.logits
    return score.float().detach().cpu().numpy()

queries = ["中國的首都是哪裡？"]
passages = [["beijing", "shanghai"]]

INSTRUCTION = "Query: "
queries = [INSTRUCTION + query for query in queries]

scores = []
for i in range(len(queries)):
    print(queries[i])
    scores.append(rerank(queries[i],passages[i]))

print(np.array(scores))  # [[[-4.7460938][-8.8515625]]]

Sentence Transformer

from sentence_transformers import CrossEncoder
import torch

#
model_name = "openbmb/MiniCPM-Reranker"
model = CrossEncoder(model_name,max_length=1024,trust_remote_code=True, automodel_args={"torch_dtype": torch.float16})
# You can also use the following code to use flash_attention_2
#model = CrossEncoder(model_name,max_length=1024,trust_remote_code=True, automodel_args={"attn_implementation":"flash_attention_2","torch_dtype": torch.float16})

model.tokenizer.padding_side = "right"

query = "中國的首都是哪裡？"
passages = [["beijing", "shanghai"]]

INSTRUCTION = "Query: "
query = INSTRUCTION + query

sentence_pairs = [[query, doc] for doc in passages]

scores = model.predict(sentence_pairs, convert_to_tensor=True).tolist()
rankings = model.rank(query, passages, return_documents=True, convert_to_tensor=True)

print(scores) # [0.0087432861328125, 0.00020503997802734375]
for ranking in rankings:
    print(f"Score: {ranking['score']:.4f}, Corpus: {ranking['text']}")
  
# ID: 0, Score: 0.0087, Text: beijing
# ID: 1, Score: 0.0002, Text: shanghai

📚 詳細文檔

模型信息

實驗結果

中文與英文重排序結果

中文對bge-large-zh-v1.5檢索的top-100進行重排，英文對bge-large-en-v1.5檢索的top-100進行重排。

中英跨語言重排序結果

對bge-m3(Dense)檢索的top100進行重排。

🔧 技術細節

MiniCPM-Reranker 基於 MiniCPM-2B-sft-bf16 訓練，結構上採取雙向注意力。採取多階段訓練方式，共使用包括開源數據、機造數據、閉源數據在內的約 600 萬條訓練數據。

📄 許可證

• 本倉庫中代碼依照 Apache-2.0 協議開源。
• MiniCPM-Reranker 模型權重的使用則需要遵循 MiniCPM 模型協議。
• MiniCPM-Reranker 模型權重對學術研究完全開放。如需將模型用於商業用途，請填寫此問卷。

歡迎關注 RAG 套件系列：

• 檢索模型：MiniCPM-Embedding
• 重排模型：MiniCPM-Reranker
• 面向 RAG 場景的 LoRA 插件：MiniCPM3-RAG-LoRA