File size: 720 Bytes
3a2a3ef |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
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))
|