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

app = FastAPI()

class ChunkSimilarityRequest(BaseModel):
    query: str
    chunks: list[str]
    top_k: int = 5


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


def keyword_similarity(query, candidates):
    pairs = [[query, candidate] for candidate in candidates]
    scores = model.predict(pairs)
    return scores



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


@app.post("/")
async def cross_encoding_similarity(
    payload: ChunkSimilarityRequest
):
    scores = keyword_similarity(payload.query, payload.chunks)

    return [
        chunk for chunk, _ in sorted(zip(payload.chunks, scores), key=lambda x: x[1], reverse=True)[:payload.top_k]
    ]