openbmb/MiniCPM-Reranker
MiniCPM-Reranker
MiniCPM-Reranker 是面壁智能与清华大学自然语言处理实验室(THUNLP)、东北大学信息检索小组(NEUIR)共同开发的中英双语言文本重排序模型,有如下特点:
- 出色的中文、英文重排序能力。
- 出色的中英跨语言重排序能力。
MiniCPM-Reranker 基于 MiniCPM-2B-sft-bf16 训练,结构上采取双向注意力。采取多阶段训练方式,共使用包括开源数据、机造数据、闭源数据在内的约 600 万条训练数据。
欢迎关注 RAG 套件系列:
- 检索模型:MiniCPM-Embedding
- 重排模型:MiniCPM-Reranker
- 面向 RAG 场景的 LoRA 插件:MiniCPM3-RAG-LoRA
MiniCPM-Reranker is a bilingual & cross-lingual text re-ranking model developed by ModelBest Inc. , THUNLP and NEUIR , featuring:
- Exceptional Chinese and English re-ranking capabilities.
- Outstanding cross-lingual re-ranking capabilities between Chinese and English.
MiniCPM-Reranker is trained based on MiniCPM-2B-sft-bf16 and incorporates bidirectional attention in its architecture. The model underwent multi-stage training using approximately 6 million training examples, including open-source, synthetic, and proprietary data.
We also invite you to explore the RAG toolkit series:
- Retrieval Model: MiniCPM-Embedding
- Re-ranking Model: MiniCPM-Reranker
- LoRA Plugin for RAG scenarios: MiniCPM3-RAG-LoRA
模型信息 Model Information
- 模型大小:2.4B
- 最大输入token数:1024
- Model Size: 2.4B
- Max Input Tokens: 1024
使用方法 Usage
输入格式 Input Format
本模型支持指令,输入格式如下:
MiniCPM-Reranker supports instructions in the following format:
<s>Instruction: {{ instruction }} Query: {{ query }}</s>{{ document }}例如:
For example:
<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)也可以不提供指令,即采取如下格式:
MiniCPM-Reranker also works in instruction-free mode in the following format:
<s>Query: {{ query }}</s>{{ document }}我们在BEIR与C-MTEB/Retrieval上测试时使用的指令见 instructions.json,其他测试不使用指令。
When running evaluation on BEIR and C-MTEB/Retrieval, we use instructions in instructions.json. For other evaluations, we do not use instructions.
环境要求 Requirements
transformers==4.37.2
flash-attn>2.3.5示例脚本 Demo
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实验结果 Evaluation Results
中文与英文重排序结果 CN/EN Re-ranking Results
中文对bge-large-zh-v1.5检索的top-100进行重排,英文对bge-large-en-v1.5检索的top-100进行重排。
We re-rank top-100 docments from bge-large-zh-v1.5 in C-MTEB/Retrieval and from bge-large-en-v1.5 in BEIR.
中英跨语言重排序结果 CN-EN Cross-lingual Re-ranking Results
对bge-m3(Dense)检索的top100进行重排。
We re-rank top-100 documents from bge-m3 (Dense).
许可证 License
- 本仓库中代码依照 Apache-2.0 协议开源。
- MiniCPM-Reranker 模型权重的使用则需要遵循 MiniCPM 模型协议。
- MiniCPM-Reranker 模型权重对学术研究完全开放。如需将模型用于商业用途,请填写此问卷。
- The code in this repo is released under the Apache-2.0 License.
- The usage of MiniCPM-Reranker model weights must strictly follow MiniCPM Model License.md.
- The models and weights of MiniCPM-Reranker are completely free for academic research. After filling out a "questionnaire" for registration, MiniCPM-Reranker weights are also available for free commercial use.
