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

# Load model and tokenizer
model_name = "tabularisai/robust-sentiment-analysis"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Function to predict sentiment
def predict_sentiment(text):
    inputs = tokenizer(text.lower(), return_tensors="pt", truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    
    probabilities = torch.nn.functional.softmax(outputs.logits, dim=-1)
    predicted_class = torch.argmax(probabilities, dim=-1).item()
    
    sentiment_map = {0: "Very Negative", 1: "Negative", 2: "Neutral", 3: "Positive", 4: "Very Positive"}
    return sentiment_map[predicted_class], {k: float(v) for k, v in zip(sentiment_map.values(), probabilities[0])}

# Function to generate random example
def random_example():
    examples = [
        "I absolutely loved this movie! The acting was superb and the plot was engaging.",
        "The service at this restaurant was terrible. I'll never go back.",
        "The product works as expected. Nothing special, but it gets the job done.",
        "I'm somewhat disappointed with my purchase. It's not as good as I hoped.",
        "This book changed my life! I couldn't put it down and learned so much."
    ]
    return random.choice(examples)

# Gradio interface
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown(
    """
    # 🎭 Sentiment Analysis Wizard
    Discover the emotional tone behind any text with our advanced AI model!
    """
    )
    
    with gr.Row():
        with gr.Column(scale=2):
            text_input = gr.Textbox(label="Enter your text here", placeholder="Type or paste your text...")
            random_btn = gr.Button("Get Random Example")
        
        with gr.Column(scale=1):
            sentiment_output = gr.Textbox(label="Overall Sentiment")
            confidence_output = gr.Label(label="Confidence Scores")
    
    analyze_btn = gr.Button("Analyze Sentiment", variant="primary")
    
    gr.Markdown(
    """
    ### How it works
    This app uses a state-of-the-art language model to analyze the sentiment of your text.
    It classifies the input into one of five categories: Very Negative, Negative, Neutral, Positive, or Very Positive.
    
    Try it out with your own text or click "Get Random Example" for inspiration!
    """
    )
    
    def analyze(text):
        sentiment, confidences = predict_sentiment(text)
        return sentiment, confidences
    
    analyze_btn.click(analyze, inputs=text_input, outputs=[sentiment_output, confidence_output])
    random_btn.click(random_example, outputs=text_input)

demo.launch()