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

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

model.eval()

def predict_sentiment(text):
    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
    )

    with torch.no_grad():
        outputs = model(**inputs)
        predicted_class = torch.argmax(outputs.logits, dim=1).item()
    
    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("Made by Ahmad Moiz with Love","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.")