import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch

tokenizer = AutoTokenizer.from_pretrained("BAAI/bge-reranker-v2-m3")
model = AutoModelForSequenceClassification.from_pretrained("BAAI/bge-reranker-v2-m3")

def rerank(query, docs):
    docs = docs.strip().split('\n')
    pairs = [(query, doc) for doc in docs]
    inputs = tokenizer(pairs, padding=True, truncation=True, return_tensors="pt")
    with torch.no_grad():
        scores = model(**inputs).logits.squeeze(-1)
    results = sorted(zip(docs, scores.tolist()), key=lambda x: x[1], reverse=True)
    return "\n\n".join([f"Score: {score:.4f}\n{doc}" for doc, score in results])

iface = gr.Interface(
    fn=rerank,
    inputs=[
        gr.Textbox(label="Query", lines=1),
        gr.Textbox(label="Documents (one per line)", lines=10)
    ],
    outputs="text",
    title="BGE Reranker v2 M3",
    description="Input a query and a list of documents. Outputs reranked documents with scores."
)

iface.launch()