import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification
import torch
import numpy as np

# Page config
st.set_page_config(
    page_title="Amazon Review Sentiment Analysis",
    page_icon="📊",
    layout="wide"
)

@st.cache_resource
def load_model():
    """Load the model and tokenizer."""
    model_name = "LiYuan/amazon-review-sentiment-analysis"
    tokenizer = AutoTokenizer.from_pretrained(model_name)
    model = AutoModelForSequenceClassification.from_pretrained(model_name)
    return tokenizer, model

def predict_sentiment(text, tokenizer, model):
    """Predict sentiment for given text."""
    inputs = tokenizer(text, return_tensors="pt", truncation=True, max_length=512)
    outputs = model(**inputs)
    predictions = torch.nn.functional.softmax(outputs.logits, dim=-1)
    return predictions.detach().numpy()[0]

def main():
    st.title("📊 Amazon Review Sentiment Analysis")
    st.write("""
    This application analyzes the sentiment of product reviews and predicts ratings (1-5 stars).
    Enter your review text below to get started!
    """)

    # Load model
    with st.spinner("Loading model..."):
        tokenizer, model = load_model()

    # Text input
    text_input = st.text_area("Enter your review text:", height=150)

    if st.button("Analyze Sentiment"):
        if text_input.strip():
            with st.spinner("Analyzing..."):
                # Get prediction
                predictions = predict_sentiment(text_input, tokenizer, model)
                predicted_rating = np.argmax(predictions) + 1  # Add 1 since ratings are 1-5

                # Display results
                col1, col2 = st.columns(2)
                
                with col1:
                    st.subheader("Predicted Rating")
                    st.markdown(f"<h1 style='text-align: center; color: #1f77b4;'>{'⭐' * predicted_rating}</h1>", unsafe_allow_html=True)
                
                with col2:
                    st.subheader("Confidence Scores")
                    for i, score in enumerate(predictions, 1):
                        st.progress(float(score))
                        st.write(f"{i} Stars: {score:.2%}")
        else:
            st.warning("Please enter some text to analyze.")

    # Additional information
    with st.expander("About this Model"):
        st.write("""
        This application uses the LiYuan/amazon-review-sentiment-analysis model from HuggingFace.
        The model is based on DistilBERT and was trained on a large dataset of Amazon product reviews.
        It can predict ratings from 1 to 5 stars based on the review text.
        
        Supported languages:
        - English
        - Dutch
        - German
        - French
        - Spanish
        - Italian
        """)

if __name__ == "__main__":
    main()