import gradio as gr
from transformers import BertForSequenceClassification, BertTokenizer
import torch

model_name = "ArunNyp7/sentimentclassifier-finetuned-bert"
model = BertForSequenceClassification.from_pretrained(model_name)
tokenizer = BertTokenizer.from_pretrained(model_name)

def classify_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=256)
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        prediction = torch.argmax(logits, dim=-1).item()

    labels = {0: "Positive", 1: "Neutral", 2: "Negative"}
    return labels[prediction]

with gr.Blocks() as demo:
    gr.Markdown("### Sentiment Classification using Fine-Tuned BERT Model")
    
    text_input = gr.Textbox(label="Enter Text", placeholder="Type here...", lines=2)
    sentiment_output = gr.Label()
    submit_btn = gr.Button("Classify Sentiment")
    submit_btn.click(fn=classify_sentiment, inputs=text_input, outputs=sentiment_output)

demo.launch(share=False)