reranker.py

from sentence_transformers import CrossEncoder

reranker = CrossEncoder("cross-encoder/ms-marco-MiniLM-L-6-v2")

def rerank(question,chunks):

    pairs=[[question,c["text"]] for c in chunks]

    scores=reranker.predict(pairs)

    ranked=sorted(
        zip(scores,chunks),
        key=lambda x:x[0],
        reverse=True
    )

    return [c for _,c in ranked]