TextRanking / app.py
IS361Group4's picture
Create app.py
4d32cd6 verified
raw
history blame contribute delete
885 Bytes
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import gradio as gr
# Load model and tokenizer
model_name = 'cross-encoder/ms-marco-MiniLM-L6-v2'
model = AutoModelForSequenceClassification.from_pretrained(model_name)
tokenizer = AutoTokenizer.from_pretrained(model_name)
model.eval()
# Define inference function
def get_similarity(question, answer):
features = tokenizer(question, answer, padding=True, truncation=True, return_tensors="pt")
with torch.no_grad():
score = model(**features).logits
return float(score[0][0]) # Convert tensor to float
# Create Gradio interface
iface = gr.Interface(
fn=get_similarity,
inputs=[
gr.Textbox(label="Question"),
gr.Textbox(label="Answer")
],
outputs=gr.Number(label="Similarity Score"),
title="Cross-Encoder QA Relevance"
)
iface.launch()