import gradio as gr
import numpy as np
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load the tokenizer and model
model_name = "NLPTeamIITGN/finetuned_llama_sentiment_sst2"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Define a label map
label_map = {0: "Negative", 1: "Positive"}

# Function to predict sentiment
def predict_sentiment(review):
    if not review.strip():
        return "No input provided", 0.0
    inputs = tokenizer(review, return_tensors="pt", truncation=True, padding=True, max_length=512)
    outputs = model(**inputs)
    logits = outputs.logits.detach().numpy()
    predicted_label = np.argmax(logits, axis=-1)[0]
    predicted_class = label_map[predicted_label]
    probability = np.max(logits, axis=-1)[0]
    return predicted_class, round(float(probability), 2)

# Gradio Interface
interface = gr.Interface(
    fn=predict_sentiment,
    inputs=gr.Textbox(lines=3, placeholder="Enter a movie review here...", label="Movie Review"),
    outputs=[
        gr.Label(label="Sentiment"),
        gr.Number(label="Confidence Score")
    ],
    title="Sentiment Analysis for Movie Reviews",
    description="Enter a movie review, and the model will classify it as Positive or Negative with a confidence score."
)

# Launch the Gradio app
interface.launch(share=True)