import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSequenceClassification
import torch.nn.functional as F
import torch

# Load the sentiment analysis pipeline, model, and tokenizer from the Hugging Face Hub
model_name = "cordondata/distilbert_sst2_600_150_acc_89"  # Replace with your actual model ID
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Define the function to analyze sentiment
def predict_sentiment(text):
    # Tokenize the input text
    tokens = tokenizer(text, return_tensors="pt")
    
    # Get token IDs
    token_ids = tokens['input_ids'].tolist()
    
    # Get the model outputs (logits)
    outputs = model(**tokens)
    logits = outputs.logits
    
    # Apply softmax to get probabilities
    probs = F.softmax(logits, dim=1).detach().numpy()[0]
    
    # Get the predicted label
    prediction = torch.argmax(logits, dim=1).item()
    sentiment = "POSITIVE" if prediction == 1 else "NEGATIVE"
    
    # Prepare the output
    tokenized_text = tokenizer.convert_ids_to_tokens(token_ids[0])
    tokenized_text_str = " ".join(tokenized_text)
    
    # Return the required outputs (excluding the original sentence)
    return (
        tokenized_text_str,  # Tokenized Sentence
        token_ids[0],  # Token IDs
        sentiment  # Sentiment without probabilities
    )

# Create the Gradio interface
interface = gr.Interface(
    fn=predict_sentiment,  # The function to run sentiment analysis
    inputs="text",  # Input component: text box
    outputs=[
        # Removed the output for the original sentence
        gr.Textbox(label="Tokenized Sentence"),
        gr.Textbox(label="Token IDs"),
        gr.Textbox(label="Sentiment")  # Only the sentiment (POSITIVE or NEGATIVE)
    ],  # Output components
    title="Sentiment Analysis: Interpreting the attitude within a review",
    description="Enter revew to analyze sentiment (positive/negative).",
    examples=[
        "This product exceeded my expectations! It’s easy to use, effective, and a great value. Highly recommend to everyone",
        "I was extremely disappointed. The quality is poor, and it didn’t work as promised. Waste of money",
        "The product is functional and performs as expected. It doesn’t stand out, but it gets the job done. If you're looking for a basic option that works without any frills, this is a decent choice. Not extraordinary, but reliable enough for everyday use. Overall, a satisfactory purchase"
    ]
)

# Launch the interface
interface.launch()