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

model_name = "cross-encoder/multi-nli-xlm-r-100"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def generate_prediction(input_text):
    input_ids = tokenizer.encode(input_text, truncation=True, padding=True, return_tensors='pt')
    outputs = model(input_ids)
    predicted_label = torch.argmax(outputs.logits)
    label_map = {0: "entailment", 1: "neutral", 2: "contradiction"}
    predicted_label_text = label_map[predicted_label.item()]
    return predicted_label_text

input_text = gr.inputs.Textbox(label="Input text")
output_text = gr.outputs.Textbox(label="Output text")

gr.Interface(
    generate_prediction,
    inputs=input_text,
    outputs=output_text,
    title="Text Classifier",
    description="A Hugging Face cross-encoder model for text classification.",
).launch()