Spaces:
Sleeping
Sleeping
File size: 1,548 Bytes
c732e8c 4d2a98a c732e8c 950be38 c732e8c bd04035 c732e8c 950be38 c732e8c bd04035 c732e8c 950be38 4d2a98a |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
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()
|