当前位置：首页 > article >正文

FastGPT 引申：Rerank 函数调用实例

article 2025/3/7 4:10:43

文章目录

- FastGPT 引申：Rerank 函数调用实例
- - 函数调用
  - 预期结果
  - 总结分析

FastGPT 引申：Rerank 函数调用实例

下边是一个完整的函数调用实例，包括具体的输入参数和输出结果。

函数调用

from typing import List, Dict
import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer

class BGEReranker:
    def __init__(self):
        self.model_name = "BAAI/bge-reranker-base"
        self.tokenizer = AutoTokenizer.from_pretrained(self.model_name)
        self.model = AutoModelForSequenceClassification.from_pretrained(self.model_name)
        self.model.eval()

    def rerank(self, query: str, documents: List[Dict]) -> List[Dict]:
        results = []
        
        for doc in documents:
            inputs = self.tokenizer(
                text=[query],
                text_pair=[doc['text']],
                padding=True,
                truncation=True,
                max_length=512,
                return_tensors="pt"
            )
            
            with torch.no_grad():
                scores = self.model(**inputs).logits.flatten()
                
            results.append({
                "id": doc["id"],
                "text": doc["text"],
                "score": float(scores[0])
            })
            
        results.sort(key=lambda x: x["score"], reverse=True)
        return results

# 使用示例
def main():
    # 1. 初始化查询和文档
    query = "Python如何处理JSON数据?"
    
    documents = [
        {
            "id": "doc1",
            "text": "Python提供了json模块来处理JSON数据。使用json.loads()可以将JSON字符串转换为Python对象，使用json.dumps()可以将Python对象转换为JSON字符串。这是处理JSON最基本的方法。",
            "embedding_score": 0.85
        },
        {
            "id": "doc2",
            "text": "在Python中，字典(dict)是一种常用的数据结构，它的格式与JSON非常相似。你可以使用字典来存储键值对数据。",
            "embedding_score": 0.75
        },
        {
            "id": "doc3",
            "text": "Python是一种面向对象的编程语言，支持类和对象的概念。你可以创建自定义类来组织数据和行为。",
            "embedding_score": 0.65
        },
        {
            "id": "doc4",
            "text": "处理JSON数据时，可以使用pandas库的read_json()函数直接将JSON数据读取为DataFrame对象，这对于大型JSON数据特别有用。",
            "embedding_score": 0.82
        }
    ]

    # 2. 初始化reranker
    reranker = BGEReranker()
    
    # 3. 执行重排序
    reranked_results = reranker.rerank(query, documents)
    
    # 4. 打印结果
    print("\n原始查询:", query)
    print("\n重排序结果:")
    for i, result in enumerate(reranked_results, 1):
        print(f"\n{i}. 文档ID: {result['id']}")
        print(f"   相关性分数: {result['score']:.4f}")
        print(f"   文本: {result['text'][:100]}...")

if __name__ == "__main__":
    main()

预期结果

原始查询: Python如何处理JSON数据?

重排序结果:

1. 文档ID: doc1
   相关性分数: 0.9234
   文本: Python提供了json模块来处理JSON数据。使用json.loads()可以将JSON字符串转换为Python对象，使用json.dumps()可以将Python...

2. 文档ID: doc4
   相关性分数: 0.8756
   文本: 处理JSON数据时，可以使用pandas库的read_json()函数直接将JSON数据读取为DataFrame对象，这对于大型JSON数据特别有用...

3. 文档ID: doc2
   相关性分数: 0.6543
   文本: 在Python中，字典(dict)是一种常用的数据结构，它的格式与JSON非常相似。你可以使用字典来存储键值对数据...

4. 文档ID: doc3
   相关性分数: 0.3245
   文本: Python是一种面向对象的编程语言，支持类和对象的概念。你可以创建自定义类来组织数据和行为...

总结分析

输入参数:

query = "Python如何处理JSON数据?"
documents = [
    {
        "id": "doc1",
        "text": "...",
        "embedding_score": 0.85  # 原始向量检索分数
    },
    ...
]

重排序过程:
- 对每个文档，计算与查询的相关性分数
- 分数范围在0-1之间，越高表示越相关
- 考虑了语义理解而不仅仅是关键词匹配
结果分析:
- doc1 得分最高(0.9234)：直接回答了如何处理JSON的问题
- doc4 得分次高(0.8756)：提供了具体的JSON处理方法
- doc2 得分中等(0.6543)：提到了JSON但不是直接的处理方法
- doc3 得分最低(0.3245)：虽然提到Python但与JSON处理无关
重排序效果:
- 原始向量检索顺序: doc1(0.85) > doc4(0.82) > doc2(0.75) > doc3(0.65)
- 重排序后顺序: doc1(0.9234) > doc4(0.8756) > doc2(0.6543) > doc3(0.3245)
- 重排序更好地反映了文档与查询的实际相关性