import streamlit as st
import torch
from transformers import DistilBertTokenizer, DistilBertForSequenceClassification

# Load pre-trained model and tokenizer
tokenizer = DistilBertTokenizer.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")
model = DistilBertForSequenceClassification.from_pretrained("distilbert-base-uncased-finetuned-sst-2-english")

# Define a function for sentiment analysis
def classify_text(text):
    inputs = tokenizer(text, return_tensors="pt")
    with torch.no_grad():
        logits = model(**inputs).logits
    predicted_class_id = logits.argmax().item()
    return model.config.id2label[predicted_class_id]

# Streamlit user interface
st.title('Sentiment Analysis with DistilBert')

# Provide examples for sentiment analysis
examples = [
    "I love this product, it's amazing!",
    "I hate the traffic, it's so frustrating.",
    "The weather is okay, not too bad.",
    "I am very excited about the new project.",
    "This food is terrible, I can't eat it.",
    "I'm not sure how I feel about this.",
    "The movie was quite boring and too long."
]

# Placeholder for text input
text_input = st.empty()

# Default text box value
default_text = "Hello, my dog is cute"

# If the user has typed something in the text box, preserve that input
custom_text = text_input.text_input("Type a sentence to classify", default_text)

# Handle clickable examples
st.write("Or click an example to fill the input box:")

for example in examples:
    if st.button(example):
        # When an example is clicked, populate the text input box with that example
        custom_text = text_input.text_input("Type a sentence to classify", value=example, key=example)

# Perform sentiment classification based on the current text input
if custom_text:
    prediction = classify_text(custom_text)
    st.write(f'**Sentiment:** {prediction}')