import gradio as gr
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "elo4/TinyBERT-sentiment-model"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=128)
    outputs = model(**inputs)
    sentiment = outputs.logits.argmax().item()
    return f"Predicted Sentiment Class: {sentiment}"

interface = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=2, placeholder="Enter a review here..."),
    outputs="text",
    title="TinyBERT Sentiment Analysis",
    description="Enter a text review and get the sentiment class predicted by TinyBERT."
)

interface.launch()