import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification

model_name = "mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

# Set the page title
st.title("Financial Sentiment Analysis App")

# Add a text input for the user to input financial news
text_input = st.text_area("Enter Financial News:", "Tesla stock is soaring after record-breaking earnings.")

# Function to perform sentiment analysis
def predict_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt")
    outputs = model(**inputs)
    sentiment_class = outputs.logits.argmax(dim=1).item()
    sentiment_mapping = {0: 'Negative', 1: 'Neutral', 2: 'Positive'}
    predicted_sentiment = sentiment_mapping.get(sentiment_class, 'Unknown')
    return predicted_sentiment

# Button to trigger sentiment analysis
if st.button("Analyze Sentiment"):
    # Check if the input text is not empty
    if text_input:
        # Show loading spinner while processing
        with st.spinner("Analyzing sentiment..."):
            sentiment = predict_sentiment(text_input)
            # Change the view based on the predicted sentiment
            st.success(f"Sentiment: {sentiment}")
            if sentiment == 'Positive':
                st.balloons()  # Celebratory animation for positive sentiment
            # Add additional views for other sentiments as needed
    else:
        st.warning("Please enter some text for sentiment analysis.")

# Optional: Display the raw sentiment scores
if st.checkbox("Show Raw Sentiment Scores"):
    if text_input:
        inputs = tokenizer(text_input, return_tensors="pt")
        outputs = model(**inputs)
        raw_scores = outputs.logits[0].tolist()
        st.info(f"Raw Sentiment Scores: {raw_scores}")

# Optional: Display additional information or analysis
# Add more components as needed for your specific use case

# Add a footer
st.text("Built with Streamlit and Transformers")