app.py

from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from sentence_transformers import CrossEncoder

app = FastAPI()
model = CrossEncoder('cross-encoder/ms-marco-MiniLM-L6-v2')

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

@app.post("/rerank")
def rerank(request: RerankRequest):
    if not request.documents:
        raise HTTPException(400, "No documents provided")
    pairs = [[request.query, doc] for doc in request.documents]
    scores = model.predict(pairs).tolist()
    ranked = sorted(zip(request.documents, scores), key=lambda x: x[1], reverse=True)
    return {"results": [{"document": doc, "score": score} for doc, score in ranked]}

@app.get("/health")
def health():
    return {"status": "ok"}