这是 Sudhanshu Sharma 创建的 bge-reranker-v2-m3 模型的 ONNX 版本

license: apache-2.0 language:

• 多语言 pipeline_tag: 文本分类 tags:
• transformers
• sentence-transformers
• text-embeddings-inference

重排序器（Reranker）

更多详情请参考我们的 GitHub：FlagEmbedding。

• 模型列表
• 使用方法
• 微调
• 评估
• 引用

与嵌入模型不同，重排序器将问题和文档作为输入，直接输出相似度而不是嵌入向量。 你可以通过向重排序器输入查询和段落来获取相关性分数。 该分数可以通过 sigmoid 函数映射到 [0,1] 范围内的浮点值。

模型列表

你可以根据场景和资源选择模型。

• 对于多语言场景，使用 BAAI/bge-reranker-v2-m3 和 BAAI/bge-reranker-v2-gemma

• 对于中文或英文，使用 BAAI/bge-reranker-v2-m3 和 BAAI/bge-reranker-v2-minicpm-layerwise。

• 对于效率，使用 BAAI/bge-reranker-v2-m3 和 BAAI/bge-reranker-v2-minicpm-layerwise 的低层。

• 为了更好的性能，推荐 BAAI/bge-reranker-v2-minicpm-layerwise 和 BAAI/bge-reranker-v2-gemma

使用方法

使用 FlagEmbedding

pip install -U FlagEmbedding

对于普通重排序器（bge-reranker-base / bge-reranker-large / bge-reranker-v2-m3）

获取相关性分数（分数越高表示相关性越强）：

from FlagEmbedding import FlagReranker
reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True) # 设置 use_fp16 为 True 可以加速计算，性能略有下降

score = reranker.compute_score(['query', 'passage'])
print(score) # -5.65234375

# 你可以通过设置 "normalize=True" 将分数映射到 0-1 范围，这将对分数应用 sigmoid 函数
score = reranker.compute_score(['query', 'passage'], normalize=True)
print(score) # 0.003497010252573502

scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
print(scores) # [-8.1875, 5.26171875]

# 你可以通过设置 "normalize=True" 将分数映射到 0-1 范围，这将对分数应用 sigmoid 函数
scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']], normalize=True)
print(scores) # [0.00027803096387751553, 0.9948403768236574]

对于基于 LLM 的重排序器

from FlagEmbedding import FlagLLMReranker
reranker = FlagLLMReranker('BAAI/bge-reranker-v2-gemma', use_fp16=True) # 设置 use_fp16 为 True 可以加速计算，性能略有下降
# reranker = FlagLLMReranker('BAAI/bge-reranker-v2-gemma', use_bf16=True) # 你也可以设置 use_bf16=True 来加速计算，性能略有下降

score = reranker.compute_score(['query', 'passage'])
print(score)

scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']])
print(scores)

对于基于 LLM 的逐层重排序器

from FlagEmbedding import LayerWiseFlagLLMReranker
reranker = LayerWiseFlagLLMReranker('BAAI/bge-reranker-v2-minicpm-layerwise', use_fp16=True) # 设置 use_fp16 为 True 可以加速计算，性能略有下降
# reranker = LayerWiseFlagLLMReranker('BAAI/bge-reranker-v2-minicpm-layerwise', use_bf16=True) # 你也可以设置 use_bf16=True 来加速计算，性能略有下降

score = reranker.compute_score(['query', 'passage'], cutoff_layers=[28]) # 调整 'cutoff_layers' 来选择用于计算分数的层。
print(score)

scores = reranker.compute_score([['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']], cutoff_layers=[28])
print(scores)

使用 Huggingface transformers

对于普通重排序器（bge-reranker-base / bge-reranker-large / bge-reranker-v2-m3）

获取相关性分数（分数越高表示相关性越强）：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-m3')
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-v2-m3')
model.eval()

pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = model(**inputs, return_dict=True).logits.view(-1, ).float()
    print(scores)

对于基于 LLM 的重排序器

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
    if prompt is None:
        prompt = "给定查询 A 和段落 B，判断段落是否包含查询的答案，通过预测 'Yes' 或 'No' 来回答。"
    sep = "\n"
    prompt_inputs = tokenizer(prompt,
                              return_tensors=None,
                              add_special_tokens=False)['input_ids']
    sep_inputs = tokenizer(sep,
                           return_tensors=None,
                           add_special_tokens=False)['input_ids']
    inputs = []
    for query, passage in pairs:
        query_inputs = tokenizer(f'A: {query}',
                                 return_tensors=None,
                                 add_special_tokens=False,
                                 max_length=max_length * 3 // 4,
                                 truncation=True)
        passage_inputs = tokenizer(f'B: {passage}',
                                   return_tensors=None,
                                   add_special_tokens=False,
                                   max_length=max_length,
                                   truncation=True)
        item = tokenizer.prepare_for_model(
            [tokenizer.bos_token_id] + query_inputs['input_ids'],
            sep_inputs + passage_inputs['input_ids'],
            truncation='only_second',
            max_length=max_length,
            padding=False,
            return_attention_mask=False,
            return_token_type_ids=False,
            add_special_tokens=False
        )
        item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
        item['attention_mask'] = [1] * len(item['input_ids'])
        inputs.append(item)
    return tokenizer.pad(
            inputs,
            padding=True,
            max_length=max_length + len(sep_inputs) + len(prompt_inputs),
            pad_to_multiple_of=8,
            return_tensors='pt',
    )

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-gemma')
model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-gemma')
yes_loc = tokenizer('Yes', add_special_tokens=False)['input_ids'][0]
model.eval()

pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
with torch.no_grad():
    inputs = get_inputs(pairs, tokenizer)
    scores = model(**inputs, return_dict=True).logits[:, -1, yes_loc].view(-1, ).float()
    print(scores)

对于基于 LLM 的逐层重排序器

import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

def get_inputs(pairs, tokenizer, prompt=None, max_length=1024):
    if prompt is None:
        prompt = "给定查询 A 和段落 B，判断段落是否包含查询的答案，通过预测 'Yes' 或 'No' 来回答。"
    sep = "\n"
    prompt_inputs = tokenizer(prompt,
                              return_tensors=None,
                              add_special_tokens=False)['input_ids']
    sep_inputs = tokenizer(sep,
                           return_tensors=None,
                           add_special_tokens=False)['input_ids']
    inputs = []
    for query, passage in pairs:
        query_inputs = tokenizer(f'A: {query}',
                                 return_tensors=None,
                                 add_special_tokens=False,
                                 max_length=max_length * 3 // 4,
                                 truncation=True)
        passage_inputs = tokenizer(f'B: {passage}',
                                   return_tensors=None,
                                   add_special_tokens=False,
                                   max_length=max_length,
                                   truncation=True)
        item = tokenizer.prepare_for_model(
            [tokenizer.bos_token_id] + query_inputs['input_ids'],
            sep_inputs + passage_inputs['input_ids'],
            truncation='only_second',
            max_length=max_length,
            padding=False,
            return_attention_mask=False,
            return_token_type_ids=False,
            add_special_tokens=False
        )
        item['input_ids'] = item['input_ids'] + sep_inputs + prompt_inputs
        item['attention_mask'] = [1] * len(item['input_ids'])
        inputs.append(item)
    return tokenizer.pad(
            inputs,
            padding=True,
            max_length=max_length + len(sep_inputs) + len(prompt_inputs),
            pad_to_multiple_of=8,
            return_tensors='pt',
    )

tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True)
model = AutoModelForCausalLM.from_pretrained('BAAI/bge-reranker-v2-minicpm-layerwise', trust_remote_code=True, torch_dtype=torch.bfloat16)
model = model.to('cuda')
model.eval()

pairs = [['what is panda?', 'hi'], ['what is panda?', 'The giant panda (Ailuropoda melanoleuca), sometimes called a panda bear or simply panda, is a bear species endemic to China.']]
with torch.no_grad():
    inputs = get_inputs(pairs, tokenizer).to(model.device)
    all_scores = model(**inputs, return_dict=True, cutoff_layers=[28])
    all_scores = [scores[:, -1].view(-1, ).float() for scores in all_scores[0]]
    print(all_scores)

微调

数据格式

训练数据应该是一个 json 文件，每行是一个如下格式的字典：

{"query": str, "pos": List[str], "neg":List[str], "prompt": str}

query 是查询，pos 是正样本文本列表，neg 是负样本文本列表，prompt 表示查询和文本之间的关系。如果某个查询没有负样本文本，你可以从整个语料库中随机采样一些作为负样本。

参见 toy_finetune_data.jsonl 查看示例数据文件。

训练

你可以使用以下代码微调重排序器：

对于基于 LLM 的重排序器

torchrun --nproc_per_node {number of gpus} \
-m FlagEmbedding.llm_reranker.finetune_for_instruction.run \
--output_dir {path to save model} \
--model_name_or_path google/gemma-2b \
--train_data ./toy_finetune_data.jsonl \
--learning_rate 2e-4 \
--num_train_epochs 1 \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 16 \
--dataloader_drop_last True \
--query_max_len 512 \
--passage_max_len 512 \
--train_group_size 16 \
--logging_steps 1 \
--save_steps 2000 \
--save_total_limit 50 \
--ddp_find_unused_parameters False \
--gradient_checkpointing \
--deepspeed stage1.json \
--warmup_ratio 0.1 \
--bf16 \
--use_lora True \
--lora_rank 32 \
--lora_alpha 64 \
--use_flash_attn True \
--target_modules q_proj k_proj v_proj o_proj

对于基于 LLM 的逐层重排序器

torchrun --nproc_per_node {number of gpus} \
-m FlagEmbedding.llm_reranker.finetune_for_layerwise.run \
--output_dir {path to save model} \
--model_name_or_path openbmb/MiniCPM-2B-dpo-bf16 \
--train_data ./toy_finetune_data.jsonl \
--learning_rate 2e-4 \
--num_train_epochs 1 \
--per_device_train_batch_size 1 \
--gradient_accumulation_steps 16 \
--dataloader_drop_last True \
--query_max_len 512 \
--passage_max_len 512 \
--train_group_size 16 \
--logging_steps 1 \
--save_steps 2000 \
--save_total_limit 50 \
--ddp_find_unused_parameters False \
--gradient_checkpointing \
--deepspeed stage1.json \
--warmup_ratio 0.1 \
--bf16 \
--use_lora True \
--lora_rank 32 \
--lora_alpha 64 \
--use_flash_attn True \
--target_modules q_proj k_proj v_proj o_proj \
--start_layer 8 \
--head_multi True \
--head_type simple \
--lora_extra_parameters linear_head

我们的重排序器初始化自 google/gemma-2b（用于基于 LLM 的重排序器）和 openbmb/MiniCPM-2B-dpo-bf16（用于基于 LLM 的逐层重排序器），我们在混合多语言数据集上进行了训练：

• bge-m3-data
• quora train data
• fever train data

评估

• llama-index.

• BEIR.
  对 bge-en-v1.5 large 的前 100 个结果进行重排序。

对 e5 mistral 7b instruct 的前 100 个结果进行重排序。

• CMTEB-retrieval.
  对 bge-zh-v1.5 large 的前 100 个结果进行重排序。

• miracl (多语言).
  对 bge-m3 的前 100 个结果进行重排序。

引用

如果你发现这个仓库有用，请考虑给一个 star 并引用

@misc{li2023making,
      title={Making Large Language Models A Better Foundation For Dense Retrieval}, 
      author={Chaofan Li and Zheng Liu and Shitao Xiao and Yingxia Shao},
      year={2023},
      eprint={2312.15503},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}
@misc{chen2024bge,
      title={BGE M3-Embedding: Multi-Lingual, Multi-Functionality, Multi-Granularity Text Embeddings Through Self-Knowledge Distillation}, 
      author={Jianlv Chen and Shitao Xiao and Peitian Zhang and Kun Luo and Defu Lian and Zheng Liu},
      year={2024},
      eprint={2402.03216},
      archivePrefix={arXiv},
      primaryClass={cs.CL}
}