import gradio as gr
from optimum.pipelines import pipeline
from transformers import AutoTokenizer

# Load ONNX optimized model
model_name = "jinaai/jina-reranker-v2-base-multilingual"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = pipeline("text-classification", model=model_name, tokenizer=tokenizer)

# Function to rerank documents
def rerank(query, documents):
    documents = documents.split("&&&")
    inputs = [[query, doc] for doc in documents if doc.strip()]
    scores = model(inputs)
    ranked_docs = sorted(zip(documents, [s['score'] for s in scores]), key=lambda x: x[1], reverse=True)
    return [{"document": doc, "score": round(score, 4)} for doc, score in ranked_docs]

# Gradio Interface
iface = gr.Interface(
    fn=rerank,
    inputs=["text", gr.Textbox(label="Documents (Separate with &&&)", placeholder="Doc1 &&& Doc2 &&& Doc3")], 
    outputs="json",
    title="JinaAI v2 Reranker API (Optimized)",
    description="Enter a query and documents (separated by '&&&'). The model will rank them based on relevance.",
)

iface.launch()