重排序模型计算两个文本的分数

使用BGE Reranker模型计算文本对相关性：

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer# 加载预训练模型与分词器（使用BAAI官方发布的reranker模型）
model = AutoModelForSequenceClassification.from_pretrained('BAAI/bge-reranker-large')
tokenizer = AutoTokenizer.from_pretrained('BAAI/bge-reranker-large')
model.eval()  # 设置为推理模式def calculate_rerank_score(query, documents):"""计算查询与多个文档的相关性分数:param query: 查询文本，如："熊猫是什么":param documents: 候选文档列表，如：["熊猫是熊科动物", "企鹅生活在南极"]:return: 包含分数和文档的排序列表"""# 构造文本对（格式：[[query, doc1], [query, doc2], ...]）pairs = [[query, doc] for doc in documents]with torch.no_grad():# 批量编码文本对inputs = tokenizer(pairs,padding=True,truncation=True,return_tensors='pt',max_length=512  )# 获取模型输出outputs = model(**inputs)scores = torch.sigmoid(outputs.logits).squeeze().tolist()  # 将logits转换为0-1概率值# 组合结果并按分数降序排序sorted_results = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)return sorted_results# 使用示例
if __name__ == "__main__":query = "What is the capital of France?"documents = ["Paris is the most populous city in France","Lyon is a major city in eastern France","The Eiffel Tower is located in Paris"]results = calculate_rerank_score(query, documents)# 打印结果print("Query:", query)for rank, (doc, score) in enumerate(results, 1):print(f"Rank {rank} (Score: {score:.4f}): {doc}")

• 输出示例：

Query: What is the capital of France?
Rank 1 (Score: 0.9872): Paris is the most populous city in France
Rank 2 (Score: 0.8531): The Eiffel Tower is located in Paris
Rank 3 (Score: 0.1023): Lyon is a major city in eastern France

• 关键实现细节说明：

  • 模型选择：使用BAAI/bge-reranker-large模型，该模型专门针对查询-文档相关性任务训练，支持中英文混合场景

  • 输入构造：将查询与每个文档组成二维列表，形成(query, doc)对，这种交叉编码方式能捕捉细粒度语义交互

  • 分数计算：通过sigmoid函数将logits转换为0-1的概率值，分数越高表示相关性越强，0.5为判定阈值

  • 批处理优化：通过padding=True和return_tensors='pt’实现批量推理，提升计算效率