app.py

import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import torch.nn.functional as F

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

# Define a function to get score
def rerank_score(query, document):
    # Format input as "query: document"
    input_text = f"{query} [SEP] {document}"
    inputs = tokenizer(input_text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        scores = F.softmax(outputs.logits, dim=1)
    relevance_score = scores[0][1].item()  # Assuming label 1 means "relevant"
    return f"Relevance Score: {relevance_score:.4f}"

# Gradio UI
demo = gr.Interface(
    fn=rerank_score,
    inputs=[gr.Textbox(label="Query"), gr.Textbox(label="Document")],
    outputs=gr.Textbox(label="Relevance Score"),
    title="BAAI bge-reranker-v2-m3",
    description="Enter a query and a document to get a relevance score."
)

if __name__ == "__main__":
    demo.launch()