Update app.py

app.py

@@ -2,7 +2,7 @@ import streamlit as st
 import requests
 from bs4 import BeautifulSoup
 from transformers import pipeline
-from transformers import AutoTokenizer, AutoModelForSequenceClassification
+from transformers import AutoTokenizer, AutoModelForSequenceClassification, AutoModelForTokenClassification
 import time
 
 # ----------- Page Layout & Custom Styling -----------
@@ -29,18 +29,22 @@ st.markdown("""
 """, unsafe_allow_html=True)
 
 # ----------- Model Setup -----------
-model_id = "LinkLinkWu/Boss_Stock_News_Analysis"
-tokenizer = AutoTokenizer.from_pretrained(model_id)
-model = AutoModelForSequenceClassification.from_pretrained(model_id)
-sentiment_pipeline = pipeline("sentiment-analysis", model=model, tokenizer=tokenizer)
+sentiment_model_id = "LinkLinkWu/Boss_Stock_News_Analysis"
+sentiment_tokenizer = AutoTokenizer.from_pretrained(sentiment_model_id)
+sentiment_model = AutoModelForSequenceClassification.from_pretrained(sentiment_model_id)
+sentiment_pipeline = pipeline("sentiment-analysis", model=sentiment_model, tokenizer=sentiment_tokenizer)
 
-# ----------- Function Definitions -----------
+ner_tokenizer = AutoTokenizer.from_pretrained("dslim/bert-base-NER")
+ner_model = AutoModelForTokenClassification.from_pretrained("dslim/bert-base-NER")
+ner_pipeline = pipeline("ner", model=ner_model, tokenizer=ner_tokenizer, grouped_entities=True)
+
+# ----------- Functions -----------
 def fetch_news(ticker):
     try:
         url = f"https://finviz.com/quote.ashx?t={ticker}"
         headers = {
-            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
-            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
+            'User-Agent': 'Mozilla/5.0',
+            'Accept': 'text/html',
             'Accept-Language': 'en-US,en;q=0.5',
             'Referer': 'https://finviz.com/',
             'Connection': 'keep-alive',
@@ -49,20 +53,20 @@ def fetch_news(ticker):
         if response.status_code != 200:
             st.error(f"Failed to fetch page for {ticker}: Status code {response.status_code}")
             return []
-        
+
         soup = BeautifulSoup(response.text, 'html.parser')
         title = soup.title.text if soup.title else ""
         if ticker not in title:
             st.error(f"Page for {ticker} not found or access denied.")
             return []
-        
+
         news_table = soup.find(id='news-table')
         if news_table is None:
             st.error(f"News table not found for {ticker}. The website structure might have changed.")
             return []
-        
+
         news = []
-        for row in news_table.findAll('tr')[:50]:  # Fetch up to 50 articles
+        for row in news_table.findAll('tr')[:50]:
             a_tag = row.find('a')
             if a_tag:
                 title = a_tag.get_text()
@@ -81,28 +85,45 @@ def analyze_sentiment(text):
         st.error(f"Sentiment analysis failed: {e}")
         return "Unknown"
 
-# ----------- Streamlit UI -----------
-st.title("📊 Stock News Sentiment Analysis")
+def extract_org_entities(text):
+    try:
+        entities = ner_pipeline(text)
+        org_entities = []
+        for ent in entities:
+            if ent["entity_group"] == "ORG":
+                clean_word = ent["word"].replace("##", "").strip()
+                if clean_word.upper() not in org_entities:
+                    org_entities.append(clean_word.upper())
+                if len(org_entities) >= 5:
+                    break
+        return org_entities
+    except Exception as e:
+        st.error(f"NER entity extraction failed: {e}")
+        return []
+
+# ----------- UI -----------
+st.title("\U0001F4CA Stock News Sentiment Analysis")
 st.markdown("""
 This tool parses stock tickers and analyzes the sentiment of related news articles.
-
-💡 *Example input:* `META, NVDA, AAPL, NTES, NCTY`
-
+\U0001F4A1 *Example input:* `META, NVDA, AAPL, NTES, NCTY`
 **Note:** If news fetching fails, it might be due to changes in the Finviz website structure or access restrictions. Please verify the website manually or try again later.
 """)
 
-# Input field for stock tickers
-tickers_input = st.text_input("Enter stock tickers separated by commas:", "META, NVDA, AAPL, NTES, NCTY")
+input_mode = st.radio("Choose input method:", ("Text (auto detect)", "Manual tickers"))
 
-# Parse and display cleaned tickers in real-time
-if tickers_input:
+if input_mode == "Manual tickers":
+    tickers_input = st.text_input("Enter stock tickers separated by commas:", "META, NVDA, AAPL")
     tickers = [ticker.strip().upper() for ticker in tickers_input.split(",") if ticker.strip()]
+else:
+    free_text = st.text_area("Enter text mentioning companies:", height=100)
+    tickers = extract_org_entities(free_text)
+
+if tickers:
     cleaned_input = ", ".join(tickers)
-    st.markdown(f"🔎 **Parsed Tickers:** `{cleaned_input}`")
+    st.markdown(f"\U0001F50E **Parsed Tickers:** `{cleaned_input}`")
 else:
     tickers = []
 
-# Button to trigger sentiment analysis
 if st.button("Get News and Sentiment"):
     if not tickers:
         st.warning("Please enter at least one stock ticker.")
@@ -112,36 +133,34 @@ if st.button("Get News and Sentiment"):
         for idx, ticker in enumerate(tickers):
             st.subheader(f"Analyzing {ticker}...")
             news_list = fetch_news(ticker)
-            
+
             if news_list:
-                # Analyze sentiment for all news articles (up to 50)
                 sentiments = []
                 for news in news_list:
                     sentiment = analyze_sentiment(news['title'])
                     sentiments.append(sentiment)
-                
-                # Determine overall sentiment based on majority
+
                 positive_count = sentiments.count("Positive")
                 negative_count = sentiments.count("Negative")
                 total = len(sentiments)
                 positive_ratio = positive_count / total if total else 0
                 negative_ratio = negative_count / total if total else 0
+
                 if positive_ratio >= 0.4:
-                   overall_sentiment = "Positive"
-                else negative_ratio >= 0.6:
-                   overall_sentiment = "Negative"
-                
-                # Display top 3 news articles with sentiment
+                    overall_sentiment = "Positive"
+                elif negative_ratio >= 0.6:
+                    overall_sentiment = "Negative"
+                else:
+                    overall_sentiment = "Neutral"
+
                 st.write(f"**Top 3 News Articles for {ticker}**")
                 for i, news in enumerate(news_list[:3], 1):
                     sentiment = sentiments[i-1]
                     st.markdown(f"{i}. [{news['title']}]({news['link']}) - **{sentiment}**")
-                
-                # Display overall sentiment
+
                 st.write(f"**Overall Sentiment for {ticker}: {overall_sentiment}**")
             else:
                 st.write(f"No news available for {ticker}.")
-            
-            # Update progress bar
+
             progress_bar.progress((idx + 1) / total_stocks)
-            time.sleep(0.1)  # Simulate processing time
+            time.sleep(0.1)