# 1. Import the required packages
import torch
import gradio as gr

from typing import Dict
from transformers import pipeline

# 2. Define function to use our model on given text
def IMDB_sentimental_analysis(text: str) -> Dict[str, float]:
    # Set up text classification pipeline
    IMDB_sentimental_analysis = pipeline(task="text-classification",
                                        # Because our model is on Hugging Face already, we can pass in the model name directly
                                        model="gokulan006/IMDB_sentimental_analysis-distilbert-base-uncased", # link to model on HF Hub
                                        device="cuda" if torch.cuda.is_available() else "cpu",
                                        top_k=None) # return all possible scores (not just top-1)

    # Get outputs from pipeline (as a list of dicts)
    outputs = IMDB_sentimental_analysis(text)[0]

    # Format output for Gradio (e.g. {"label_1": probability_1, "label_2": probability_2})
    output_dict = {}
    for item in outputs:
        output_dict[item["label"]] = item["score"]

    return output_dict

# 3. Create a Gradio interface with details about our app
description = """
A text classifier to determine if a movie review is positive or negative.

Fine-tuned from [DistilBERT](https://huggingface.co/distilbert/distilbert-base-uncased) on a [small dataset of food and not food text](https://github.com/gokulan006/IMDB-Sentiment-Analysis/raw/refs/heads/master/IMDB%20Dataset.csv).

 
"""

demo = gr.Interface(fn=IMDB_sentimental_analysis,
             inputs="text",
             outputs=gr.Label(num_top_classes=2), # show top 2 classes (that's all we have)
             title="📝🎬🎥 IMDB SENTIMENTAL ANALYSIS",
             description=description,
             examples=[["A visually stunning and thought-provoking film that dares to take its time. Denis Villeneuve masterfully crafts a slow-burn neo-noir that expands upon the themes of identity and humanity introduced in the original Blade Runner. While some may find the pacing too slow, those who appreciate deep, atmospheric storytelling will find it incredibly rewarding. The cinematography by Roger Deakins is breathtaking, and the score perfectly complements the film’s futuristic yet melancholic tone. A near-perfect sequel, though it may not be for everyone."],
                        ["This movie is a complete disaster. The jokes are forced, the script makes no sense, and the humor is cringeworthy. I can't believe such a talented cast was wasted in this mess. It’s painfully unfunny, and I regret spending my time on it."]])

# 4. Launch the interface
if __name__ == "__main__":
    demo.launch()