license: apache-2.0 pipeline_tag: text-classification

BGE-RERANKER-V2的ONNX O4版本

pairs = [['我讨厌吃苹果。','我真的很喜欢吃苹果'],['我真的很喜欢吃苹果', 'Realmente me gusta comer manzana.'], ['我真的很喜欢吃苹果', '我讨厌苹果'],['Las manzanas son geniales.','Realmente me gusta comer manzana.']]

from optimum.onnxruntime import ORTModelForFeatureExtraction,ORTModelForSequenceClassification
from transformers import AutoTokenizer

model_checkpoint = "onnxO4_bge_reranker_v2_m3"

ort_model = ORTModelForSequenceClassification.from_pretrained(model_checkpoint)
tokenizer = AutoTokenizer.from_pretrained(model_checkpoint)

# ONNX结果
import torch


with torch.no_grad():
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors='pt', max_length=512)
    scores = ort_model(**inputs, return_dict=True).logits.view(-1, ).float()
    print(scores)

## tensor([ -9.5081,  -3.9569, -10.8632,   0.3756])

# 原始未量化版本

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()

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)

## tensor([ -9.4973,  -3.9538, -10.8504,   0.3660])