【限时免费】 生产力升级：将bge-reranker-large模型封装为可随时调用的API服务

生产力升级：将bge-reranker-large模型封装为可随时调用的API服务

【免费下载链接】bge-reranker-large 项目地址: https://gitcode.com/mirrors/BAAI/bge-reranker-large

引言：为什么要将模型API化？

在当今的AI开发中，将本地模型封装成API服务已成为一种常见的实践。这种做法的好处显而易见：

1. 解耦：将模型逻辑与前端或其他调用方解耦，使得模型更新或替换时不影响调用方。
2. 复用：通过API服务，模型可以被多个应用（如网站、App、小程序）共享调用，避免重复开发。
3. 跨语言调用：API服务通常基于HTTP协议，任何支持HTTP的语言都可以调用，解决了语言环境差异的问题。
4. 简化部署：API服务可以集中部署和管理，降低了模型分发的复杂性。

本文将指导开发者如何将开源模型bge-reranker-large封装成一个标准的RESTful API服务，使其能够被灵活调用。

---

技术栈选择

为了实现这一目标，我们推荐使用FastAPI作为Web框架。选择FastAPI的原因如下：

1. 高性能：FastAPI基于Starlette和Pydantic，性能接近Node.js和Go。
2. 自带文档：FastAPI自动生成交互式API文档（Swagger UI），方便开发者调试和测试。
3. 易于使用：FastAPI的语法简洁，学习成本低，适合快速开发。

当然，如果你更熟悉Flask，也可以选择它，但FastAPI在性能和功能上更具优势。

---

核心代码：模型加载与推理函数

首先，我们需要将bge-reranker-large模型的加载和推理逻辑封装成一个独立的Python函数。以下是核心代码：

from transformers import AutoModelForSequenceClassification, AutoTokenizer
import torch

# 加载模型和分词器
model_name = "BAAI/bge-reranker-large"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)
model.eval()

def rerank(query: str, documents: list[str]) -> list[float]:
    """
    对输入的查询和文档列表进行重新排序，返回相关性分数。
    """
    scores = []
    for doc in documents:
        inputs = tokenizer(query, doc, return_tensors="pt", truncation=True, max_length=512)
        with torch.no_grad():
            outputs = model(**inputs)
            score = outputs.logits[0].item()
            scores.append(score)
    return scores

代码说明：

1. 模型加载：使用transformers库加载预训练的bge-reranker-large模型和分词器。
2. 推理函数：rerank函数接收一个查询字符串和一个文档列表，返回每个文档的相关性分数。

---

API接口设计与实现

接下来，我们使用FastAPI将上述函数封装成一个API服务。以下是完整的服务端代码：

from fastapi import FastAPI
from pydantic import BaseModel
from typing import List

app = FastAPI()

class RerankRequest(BaseModel):
    query: str
    documents: List[str]

@app.post("/rerank")
async def rerank_api(request: RerankRequest):
    scores = rerank(request.query, request.documents)
    return {"scores": scores}

代码说明：

1. 请求模型：使用pydantic定义输入数据的结构，确保请求格式正确。
2. API接口：/rerank接口接收POST请求，调用rerank函数并返回JSON格式的结果。

---

测试API服务

完成API开发后，我们可以使用curl或Python的requests库测试服务是否正常工作。

使用curl测试：

curl -X POST "http://127.0.0.1:8000/rerank" \
-H "Content-Type: application/json" \
-d '{"query": "What is AI?", "documents": ["AI is a branch of computer science.", "The sky is blue."]}'

使用Python requests测试：

import requests

response = requests.post(
    "http://127.0.0.1:8000/rerank",
    json={"query": "What is AI?", "documents": ["AI is a branch of computer science.", "The sky is blue."]},
)
print(response.json())

预期输出：

{"scores": [0.95, 0.12]}

---

部署与性能优化考量

部署方案

1. Gunicorn：使用Gunicorn作为WSGI服务器，提升并发能力。

   gunicorn -w 4 -k uvicorn.workers.UvicornWorker main:app
   

2. Docker：将服务打包成Docker镜像，方便跨环境部署。

性能优化

1. 批量推理：支持批量输入文档，减少多次调用的开销。
2. 缓存：对频繁查询的结果进行缓存，提升响应速度。
3. 异步处理：使用FastAPI的异步特性，提高并发性能。

---

结语

【免费下载链接】bge-reranker-large 项目地址: https://gitcode.com/mirrors/BAAI/bge-reranker-large