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

# Load model - xlm-roberta-base'i doğrudan kullanalım
model = AutoModelForSequenceClassification.from_pretrained("xlm-roberta-base", num_labels=3)
tokenizer = AutoTokenizer.from_pretrained("xlm-roberta-base")

def predict(premise, hypothesis):
    inputs = tokenizer(premise, hypothesis, return_tensors="pt", truncation=True)
    outputs = model(**inputs)
    prediction = outputs.logits.softmax(-1)[0]
    
    return {
        "Entailment": float(prediction[0]),
        "Neutral": float(prediction[1]),
        "Contradiction": float(prediction[2])
    }

demo = gr.Interface(
    fn=predict,
    inputs=[
        gr.Textbox(label="Premise"),
        gr.Textbox(label="Hypothesis")
    ],
    outputs=gr.Label(),
    title="Natural Language Inference",
    examples=[
        ["The cat is sleeping.", "The cat is awake."],
        ["It's raining.", "The ground is wet."]
    ]
)

demo.launch()