|
from fastapi import FastAPI, HTTPException
|
|
from pydantic import BaseModel
|
|
from FlagEmbedding import FlagReranker
|
|
|
|
|
|
app = FastAPI(title="BGE-Reranker API")
|
|
|
|
|
|
reranker = FlagReranker('BAAI/bge-reranker-v2-m3', use_fp16=True)
|
|
|
|
|
|
class RerankRequest(BaseModel):
|
|
query: str
|
|
passage: str
|
|
|
|
|
|
@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))
|
|
|