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

MODEL_ID = "edaUsha/Fine_Tuning_Bert_For_Sentiment_Anaysis"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSequenceClassification.from_pretrained(MODEL_ID)

id2label = {
    0: "NEGATIVE",   # change to your actual label name
    1: "NEUTRAL",    # change if needed
    2: "POSITIVE"    # change if needed
}

def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True)
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
        probs = torch.softmax(logits, dim=-1)[0]
    pred_id = int(torch.argmax(probs))
    label = id2label[pred_id]
    confidence = float(probs[pred_id])
    return f"{label} ({confidence:.2f})"

demo = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=3, label="Input text"),
    outputs=gr.Textbox(label="Prediction"),
    title="Fine-tuned BERT Sentiment Analysis",
    description="Enter a sentence to see its predicted sentiment."
)

if __name__ == "__main__":
    demo.launch()