app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from FlagEmbedding import FlagReranker

# Initialize FastAPI app
app = FastAPI(title="BGE-Reranker API")

# Load the reranker model once during startup
reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True)

# Define the request body schema
class RerankRequest(BaseModel):
    query: str
    passage: str

# Define the API endpoint
@app.post("/rerank")
def rerank_score(request: RerankRequest):
    try:
        score = reranker.compute_score([request.query, request.passage], normalize=True)
        return {"score": score}
    except Exception as e:
        raise HTTPException(status_code=500, detail=str(e))