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

model = AutoModelForSequenceClassification.from_pretrained("AlCyede/sarcastic-text_prediction")
tokenizer = AutoTokenizer.from_pretrained("distilbert-base-uncased")

def predict(text):
    inputs = tokenizer(text, padding=True, truncation=True, return_tensors="pt")
    outputs = model(**inputs)
    logits = outputs.logits
    predicted_class_id = logits.argmax().item()
    predicted_class_prob = logits.softmax(dim=1)[0][predicted_class_id].item()
    return f"prediction: {model.config.id2label[predicted_class_id]}\nconfidence: {predicted_class_prob * 100:.02f}%"

demo = gr.Interface(fn=predict, inputs="text", outputs="text")
demo.launch(share=True)