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

# Load the tokenizer and model
tokenizer = AutoTokenizer.from_pretrained("mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis")
model = AutoModelForSequenceClassification.from_pretrained("mrm8488/distilroberta-finetuned-financial-news-sentiment-analysis")

# Set the model to evaluation mode
model.eval()

def predict_sentiment(text):
    # Tokenize the input text and prepare it to be used as model input
    inputs = tokenizer.encode_plus(
        text,
        add_special_tokens=True,
        return_tensors="pt",
        max_length=512,  # You can adjust the max length based on your requirements
        truncation=True
    )

    # Perform the prediction
    with torch.no_grad():
        outputs = model(**inputs)
        predicted_class = torch.argmax(outputs.logits, dim=1).item()
    
    # Determine the sentiment based on the predicted class
    if predicted_class == 0:
        return "Negative sentiment"
    elif predicted_class == 1:
        return "Neutral sentiment"
    elif predicted_class == 2:
        return "Positive sentiment"

# Streamlit interface
st.title("Financial News Sentiment Analysis")

user_input = st.text_area("Enter the financial news text:")
if st.button("Analyze"):
    if user_input:
        # Get the sentiment prediction
        sentiment = predict_sentiment(user_input)
        st.write("Predicted sentiment:", sentiment)
    else:
        st.write("Please enter some text.")