import gradio as gr
from transformers import pipeline

def classifier(sentence):
    classifier = pipeline(
        "text-classification", 
        model="AirrStorm/DistilBERT-SST2", 
        tokenizer="AirrStorm/DistilBERT-SST2",
        device=0
    )
    
    label_mapping = {"LABEL_0": "negative", "LABEL_1": "positive"}
    result = classifier(sentence) 
    predicted_label = label_mapping[result[0]['label']]
    return predicted_label  # Should print "negative" or "positive"

# Define the Gradio Interface
demo = gr.Interface(
    fn=classifier, 
    inputs=gr.Textbox(
        lines=4, 
        placeholder="Type a sentence to classify the sentiment...", 
        label="Input Text"
    ), 
    outputs=gr.Textbox(
        label="Predicted Sentiment"
    ),
    title="Sentiment Analysis", 
    description="Classify the sentiment of the input text as positive or negative.",
    theme="hugging-face",  # Optional, you can experiment with other themes like 'huggingface'
    allow_flagging="never",  # Disable flagging if not needed
)

# Launch the interface
demo.launch()