import streamlit as st
import torch
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import matplotlib.pyplot as plt
import seaborn as sns

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

def analyze_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    outputs = model(**inputs)
    probs = torch.nn.functional.softmax(outputs.logits, dim=-1)
    return probs.detach().numpy()[0]

st.title("Sentiment Analysis with Transformer")

prompt_text = "rt NLP trnsf xmpl w PyTrc Hggng Fc, nd Strml ntrfc fr npts tpts, ncld mtpl grph f ncsry. Cd z t ct pst."
st.write(f"**Prompt:** {prompt_text}")

user_input = st.text_area("Enter text for sentiment analysis:", "I love this product!")

if st.button("Analyze Sentiment"):
    sentiment_scores = analyze_sentiment(user_input)
    
    st.write("Sentiment Scores:")
    st.write(f"Negative: {sentiment_scores[0]:.4f}")
    st.write(f"Positive: {sentiment_scores[1]:.4f}")
    
    # Create and display multiple graphs
    fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(12, 4))
    
    # Bar plot
    ax1.bar(['Negative', 'Positive'], sentiment_scores)
    ax1.set_ylabel('Score')
    ax1.set_title('Sentiment Analysis Results (Bar Plot)')
    
    # Pie chart
    ax2.pie(sentiment_scores, labels=['Negative', 'Positive'], autopct='%1.1f%%')
    ax2.set_title('Sentiment Analysis Results (Pie Chart)')
    
    st.pyplot(fig)
    
    # Heatmap
    fig, ax = plt.subplots(figsize=(8, 2))
    sns.heatmap([sentiment_scores], annot=True, cmap="coolwarm", cbar=False, ax=ax)
    ax.set_xticklabels(['Negative', 'Positive'])
    ax.set_yticklabels(['Sentiment'])
    ax.set_title('Sentiment Analysis Results (Heatmap)')
    st.pyplot(fig)

st.write("Note: This example uses a pre-trained model for English sentiment analysis.")