app.py

import gradio as gr
from sentence_transformers import CrossEncoder

# Load the reranker model
model = CrossEncoder("jinaai/jina-reranker-v2-base-multilingual", trust_remote_code=True )

# Function to score query-document pairs
def rerank(query, documents):
    scores = model.predict([[query, doc] for doc in documents])
    ranked_docs = sorted(zip(documents, scores), key=lambda x: x[1], reverse=True)
    return [{"document": doc, "score": round(score, 4)} for doc, score in ranked_docs]

# Define Gradio interface
iface = gr.Interface(
    fn=rerank,
    inputs=["text", gr.List(gr.Textbox())], 
    outputs="json",
    title="JinaAI v2 Reranker API",
    description="Enter a query and a list of documents. The model will rank them based on relevance.",
)

# Launch Gradio app
iface.launch()