import streamlit as st
from transformers import AutoTokenizer, AutoModelForSequenceClassification

# Load pre-trained model and tokenizer
model_name = "mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis"  # Replace with your chosen model
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSequenceClassification.from_pretrained(model_name)

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

def main():
    st.title("Financial Sentiment Analysis")

    # Get user input
    text = st.text_area("Enter financial content:")

    if st.button("Classify Sentiment"):
        if text:
            predicted_sentiment = predict_sentiment(text)
            st.success(f"Predicted sentiment: {predicted_sentiment}")
        else:
            st.warning("Please enter some text.")

if __name__ == "__main__":
    main()