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

model_path = "my_model"

tokenizer = BertTokenizer.from_pretrained(model_path)
model = BertForSequenceClassification.from_pretrained(model_path)

device = torch.device("cpu")
model.to(device)
model.eval()

def predict(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)

    predicted_class = torch.argmax(probs, dim=1).item()
    confidence = probs[0][predicted_class].item()

    label = "🟢 Positive" if predicted_class == 1 else "🔴 Negative"

    return label, f"{confidence:.2f}"

# 🎨 CUSTOM UI
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    
    gr.Markdown("""
    # 🎬 AI Sentiment Analyzer
    ### Analyze emotions in text using BERT 🤖
    """)
    
    with gr.Row():
        text_input = gr.Textbox(
            placeholder="Type your sentence here...",
            lines=3,
            label="Input Text"
        )

    analyze_btn = gr.Button("Analyze Sentiment 🚀")

    with gr.Row():
        result_label = gr.Textbox(label="Prediction")
        confidence_score = gr.Textbox(label="Confidence")

    analyze_btn.click(
        fn=predict,
        inputs=text_input,
        outputs=[result_label, confidence_score]
    )

demo.launch()