import gradio as gr
from transformers import AutoModelForSequenceClassification
import torch

# Load the pre-trained model
model = AutoModelForSequenceClassification.from_pretrained(
    'jinaai/jina-reranker-v2-base-multilingual',
    torch_dtype="auto",
    trust_remote_code=True,
)

device = 'cuda' if torch.cuda.is_available() else 'cpu'
model.to(device)  # Move model to GPU if available, otherwise CPU
model.eval()

def compute_scores(query, documents):
    """
    Compute scores between a query and multiple documents using the loaded model.
    
    Args:
        query (str): The input query string.
        documents (list of str): List of document strings to compare against the query.
        
    Returns:
        list of float: Scores representing the relevance of each document to the query.
    """
    documents_list = documents.split('\n')
    sentence_pairs = [[query, doc] for doc in documents_list]
    scores = model.compute_score(sentence_pairs, max_length=1024)
    return scores

# Define Gradio interface
iface = gr.Interface(
    fn=compute_scores,
    inputs=[
        gr.Textbox(lines=2, placeholder="Enter your query here..."),
        gr.Textbox(lines=8, placeholder="Enter your documents separated by newlines...")
    ],
    outputs="json",
    title="Sentence Pair Scoring with Jina Reranker Model",
    description="This tool computes the relevance scores between a given query and a set of documents using the Jina Reranker model."
)

# Launch the interface
if __name__ == "__main__":
    iface.launch()