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

# -----------------------------
# Configuration Section
# -----------------------------
MODEL_NAME = "quocviethere/imdb-roberta"  # Replace with your actual model ID

# -----------------------------
# Model Loading Section
# -----------------------------
try:
    # Load tokenizer and model from Hugging Face Hub
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
    model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME)
    
    # Initialize the sentiment analysis pipeline
    sentiment_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
    
    # Verify label mapping
    label_mapping = model.config.id2label
    print(f"Model label mapping: {label_mapping}")
except Exception as e:
    print(f"Error loading model: {e}")
    raise

# -----------------------------
# Sentiment Analysis Function
# -----------------------------
def analyze_sentiment(text):
    try:
        # Perform sentiment analysis
        result = sentiment_pipeline(text)[0]
        
        # Extract label and score
        label = result['label']
        score = result['score']
        
        # Map label to sentiment
        if label in label_mapping.values():
            sentiment = "Positive 😊" if label == "POSITIVE" else "Negative 😞"
        else:
            # Handle unexpected labels
            sentiment = label
            print(f"Unexpected label received: {label}")
        
        confidence = f"Confidence: {round(score * 100, 2)}%"
        
        return sentiment, confidence
    except Exception as e:
        print(f"Error during sentiment analysis: {e}")
        return "Error", "Could not process the input."

# -----------------------------
# Gradio Interface Section
# -----------------------------
iface = gr.Interface(
    fn=analyze_sentiment,
    inputs=gr.Textbox(
        lines=5,
        placeholder="Enter a movie review here...",
        label="Movie Review"
    ),
    outputs=[
        gr.Textbox(label="Sentiment"),
        gr.Textbox(label="Confidence")
    ],
    title="IMDb Sentiment Analysis with RoBERTa",
    description="Analyze the sentiment of movie reviews using a fine-tuned RoBERTa model.",
    examples=[
        ["I loved the cinematography and the story was captivating."],
        ["The movie was a complete waste of time. Poor acting and boring plot."]
    ],
    theme="default"
)

# -----------------------------
# Launch the Interface
# -----------------------------
iface.launch()