import streamlit as st
from transformers import BertTokenizer, BertForSequenceClassification
import torch

model_name = "yiyanghkust/finbert-tone"
tokenizer = BertTokenizer.from_pretrained(model_name)
model = BertForSequenceClassification.from_pretrained(model_name)

def analyze_sentiment(text):
    inputs = tokenizer(text, return_tensors="pt", truncation=True, padding=True, max_length=512)
    
    with torch.no_grad():
        outputs = model(**inputs)
        logits = outputs.logits
    
    sentiment = torch.argmax(logits, dim=1).item()
    
    if sentiment == 0:
        return "Negative"
    elif sentiment == 1:
        return "Neutral"
    else:
        return "Positive"

st.title("FinBERT Sentiment Analysis")
st.write(
    "This app uses FinBERT model to analyze sentiment of financial texts. "
    "Enter text below to get its sentiment classification."
)

text_input = st.text_area("Enter your text here:")

if text_input:
    sentiment = analyze_sentiment(text_input)
    st.write(f"Sentiment: {sentiment}")