Update app.py

app.py

@@ -6,14 +6,16 @@ import plotly.express as px
 from dateutil import parser
 import datetime
 import requests
-from transformers import pipeline
+from transformers import BertTokenizer, BertForSequenceClassification, pipeline
 
 st.set_page_config(page_title="Stock News Sentiment Analyzer", layout="wide")
 
-# Initialize FinBERT pipeline
+# Initialize FinBERT (yiyanghkust/finbert-tone) pipeline
 @st.cache_resource
 def load_model():
-    return pipeline("text-classification", model="yiyanghkust/finbert-tone")
+    model = BertForSequenceClassification.from_pretrained('yiyanghkust/finbert-tone', num_labels=3)
+    tokenizer = BertTokenizer.from_pretrained('yiyanghkust/finbert-tone')
+    return pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
 
 finbert = load_model()
 
@@ -78,10 +80,10 @@ def score_news(parsed_news_df):
         label = pred['label']
         score = pred['score']
         
-        # Convert to -1 to 1 scale
-        if label == 'positive':
+        # Map 'LABEL_0' (neutral), 'LABEL_1' (positive), and 'LABEL_2' (negative)
+        if label == 'LABEL_1':
             sentiment_score = score
-        elif label == 'negative':
+        elif label == 'LABEL_2':
             sentiment_score = -score
         else:  # neutral
             sentiment_score = 0
@@ -172,7 +174,7 @@ try:
         The above charts average the sentiment scores of {ticker} stock hourly and daily.
         The table below shows recent headlines with their sentiment scores and classifications.
         The news headlines are obtained from the FinViz website.
-        Sentiments are analyzed using the ProsusAI/finbert model, which is specifically trained for financial text.
+        Sentiments are analyzed using the yiyanghkust/finbert-tone model, which is specifically trained for financial text.
         Links have been verified for validity.
         """
         
@@ -196,4 +198,4 @@ hide_streamlit_style = """
 footer {visibility: hidden;}
 </style>
 """
-st.markdown(hide_streamlit_style, unsafe_allow_html=True)
+st.markdown(hide_streamlit_style, unsafe_allow_html=True)