Update app.py

app.py

@@ -3,6 +3,7 @@ import re
 import time
 import json
 import random
+import requests  # <-- For Polygon API calls
 import finnhub
 import torch
 import gradio as gr
@@ -14,10 +15,18 @@ from collections import defaultdict
 from datetime import date, datetime, timedelta
 from transformers import AutoTokenizer, AutoModelForCausalLM, TextStreamer
 
+################################################################################
+# Set up environment variables, tokens, model, tokenizer, and other base config
+################################################################################
 
-access_token = os.environ["HF_TOKEN"]
-finnhub_client = finnhub.Client(api_key=os.environ["FINNHUB_API_KEY"])
+# Make sure these environment variables or your chosen tokens are set properly.
+access_token = os.environ.get("HF_TOKEN", "YOUR_HF_ACCESS_TOKEN")
+finnhub_api_key = os.environ.get("FINNHUB_API_KEY", "YOUR_FINNHUB_API_KEY")
+polygon_api_key = "OeD88ExoL458WlvNpCyiZXnHI0s_h05t"  # <--- Replace "X" with your actual Polygon key
 
+finnhub_client = finnhub.Client(api_key=finnhub_api_key)
+
+# Load base model & LoRA
 base_model = AutoModelForCausalLM.from_pretrained(
     'meta-llama/Llama-2-7b-chat-hf',
     token=access_token,
@@ -40,15 +49,25 @@ tokenizer = AutoTokenizer.from_pretrained(
 
 streamer = TextStreamer(tokenizer)
 
+# Special Llama-format prompt tokens
 B_INST, E_INST = "[INST]", "[/INST]"
 B_SYS, E_SYS = "<<SYS>>\n", "\n<</SYS>>\n\n"
 
-SYSTEM_PROMPT = "You are a seasoned stock market analyst. Your task is to list the positive developments and potential concerns for companies based on relevant news and basic financials from the past weeks, then provide an analysis and prediction for the companies' stock price movement for the upcoming week. " \
-    "Your answer format should be as follows:\n\n[Positive Developments]:\n1. ...\n\n[Potential Concerns]:\n1. ...\n\n[Prediction & Analysis]\nPrediction: ...\nAnalysis: ..."
+SYSTEM_PROMPT = (
+    "You are a seasoned stock market analyst. Your task is to list the positive developments and "
+    "potential concerns for companies based on relevant news and basic financials from the past weeks, "
+    "then provide an analysis and prediction for the companies' stock price movement for the upcoming week. "
+    "Your answer format should be as follows:\n\n[Positive Developments]:\n1. ...\n\n[Potential Concerns]:\n1. ...\n\n"
+    "[Prediction & Analysis]\nPrediction: ...\nAnalysis: ..."
+)
+
 
+###############################################################################
+# Utility functions
+###############################################################################
 
 def print_gpu_utilization():
-    
+    """Helper to print GPU utilization (MB) using NVML."""
     nvmlInit()
     handle = nvmlDeviceGetHandleByIndex(0)
     info = nvmlDeviceGetMemoryInfo(handle)
@@ -56,37 +75,41 @@ def print_gpu_utilization():
 
 
 def get_curday():
-    
+    """Returns today's date in YYYY-MM-DD."""
     return date.today().strftime("%Y-%m-%d")
 
 
 def n_weeks_before(date_string, n):
-    
-    date = datetime.strptime(date_string, "%Y-%m-%d") - timedelta(days=7*n)
-
-    return date.strftime("%Y-%m-%d")
+    """Given 'date_string' in YYYY-MM-DD and an integer n, returns the date that is n weeks prior."""
+    d = datetime.strptime(date_string, "%Y-%m-%d") - timedelta(days=7*n)
+    return d.strftime("%Y-%m-%d")
 
 
 def get_stock_data(stock_symbol, steps):
+    """
+    Downloads stock price data using yfinance for the given date steps.
+    Returns a DataFrame containing Start Date, End Date, Start Price, End Price
+    for each of the time intervals.
+    """
     stock_data = yf.download(stock_symbol, steps[0], steps[-1])
     if len(stock_data) == 0:
         raise gr.Error(f"Failed to download stock price data for symbol {stock_symbol} from yfinance!")
-        
-    print(stock_data)
     
+    print(stock_data)  # For debugging
+
     dates, prices = [], []
     available_dates = stock_data.index.format()
-    
-    for date in steps[:-1]:
+
+    for date_ in steps[:-1]:
         for i in range(len(stock_data)):
-            if available_dates[i] >= date:
-                prices.append(stock_data['Close'].iloc[i])  # Use .iloc here
+            if available_dates[i] >= date_:
+                prices.append(stock_data['Close'].iloc[i])
                 dates.append(datetime.strptime(available_dates[i], "%Y-%m-%d"))
                 break
 
-    # Append the last date and price
+    # Append last date & price
     dates.append(datetime.strptime(available_dates[-1], "%Y-%m-%d"))
-    prices.append(stock_data['Close'].iloc[-1])  # Use .iloc here as well
+    prices.append(stock_data['Close'].iloc[-1])
     
     return pd.DataFrame({
         "Start Date": dates[:-1],
@@ -96,195 +119,385 @@ def get_stock_data(stock_symbol, steps):
     })
 
 
+###############################################################################
+# News retrieval
+###############################################################################
+
+def parse_polygon_news_item(item):
+    """
+    Convert a Polygon news item into a {date, headline, summary} dict similar to the finnhub format used.
+    Published_utc is in ISO8601, e.g. '2021-04-23T12:47:00Z'.
+    """
+    published_str = item.get('published_utc', '')
+    try:
+        # Convert e.g. "2021-04-23T12:47:00Z" to a datetime, then to YYYYmmddHHMMSS string
+        dt = datetime.strptime(published_str, "%Y-%m-%dT%H:%M:%SZ")
+        date_fmt = dt.strftime("%Y%m%d%H%M%S")
+    except:
+        # In case of parsing error, just keep it raw
+        date_fmt = published_str
+
+    headline = item.get("title", "")
+    summary = item.get("description", "")
+    return {
+        "date": date_fmt,
+        "headline": headline,
+        "summary": summary
+    }
+
+
 def get_news(symbol, data):
+    """
+    For each row in data (weekly intervals), fetch both Finnhub and Polygon news.
+    Combine them, sort by date, store the combined list in data['News'].
+    """
+
+    combined_news_list = []
     
-    news_list = []
-    
-    for end_date, row in data.iterrows():
+    for idx, row in data.iterrows():
         start_date = row['Start Date'].strftime('%Y-%m-%d')
         end_date = row['End Date'].strftime('%Y-%m-%d')
-#         print(symbol, ': ', start_date, ' - ', end_date)
-        time.sleep(1) # control qpm
-        weekly_news = finnhub_client.company_news(symbol, _from=start_date, to=end_date)
-        if len(weekly_news) == 0:
-            raise gr.Error(f"No company news found for symbol {symbol} from finnhub!")
-        weekly_news = [
-            {
-                "date": datetime.fromtimestamp(n['datetime']).strftime('%Y%m%d%H%M%S'),
-                "headline": n['headline'],
-                "summary": n['summary'],
-            } for n in weekly_news
-        ]
-        weekly_news.sort(key=lambda x: x['date'])
-        news_list.append(json.dumps(weekly_news))
-    
-    data['News'] = news_list
-    
+
+        # Sleep for QPM control if needed
+        time.sleep(1)
+        
+        #######################################################################
+        # 1) Finnhub News for the weekly period
+        #######################################################################
+        finnhub_weekly_news = []
+        try:
+            finnhub_news = finnhub_client.company_news(symbol, _from=start_date, to=end_date)
+            for n in finnhub_news:
+                dt_str = datetime.fromtimestamp(n['datetime']).strftime('%Y%m%d%H%M%S')
+                finnhub_weekly_news.append({
+                    "date": dt_str,
+                    "headline": n['headline'],
+                    "summary": n['summary']
+                })
+        except Exception as e:
+            # If Finnhub returns an error or no results, we can choose to ignore or handle differently
+            finnhub_weekly_news = []
+        
+        #######################################################################
+        # 2) Polygon News 
+        #    There's no date range param in the sample snippet. We'll fetch
+        #    up to 30 news items, then filter them by start_date <= published <= end_date
+        #######################################################################
+        polygon_weekly_news = []
+        polygon_url = (
+            f"https://api.polygon.io/v2/reference/news"
+            f"?ticker={symbol}"
+            f"&order=asc"
+            f"&limit=30"
+            f"&sort=published_utc"
+            f"&apiKey={polygon_api_key}"
+        )
+        try:
+            resp = requests.get(polygon_url)
+            if resp.status_code == 200:
+                polygon_data = resp.json()
+                results = polygon_data.get('results', [])
+                for item in results:
+                    news_item = parse_polygon_news_item(item)
+                    # Filter by date range
+                    # news_item['date'] is "YYYYmmddHHMMSS"
+                    # We compare it with start_date/end_date as YYYY-MM-DD (the day-based boundary).
+                    # So let's parse it properly to do a valid comparison:
+                    try:
+                        dt_item = datetime.strptime(news_item['date'], "%Y%m%d%H%M%S")
+                        dt_start = datetime.strptime(start_date, "%Y-%m-%d")
+                        dt_end = datetime.strptime(end_date, "%Y-%m-%d") + timedelta(days=1) - timedelta(seconds=1)
+                        if dt_start <= dt_item <= dt_end:
+                            polygon_weekly_news.append(news_item)
+                    except:
+                        # If for any reason the date parse fails, skip
+                        pass
+            else:
+                print(f"Polygon news request returned status code {resp.status_code}")
+        except Exception as e:
+            print(f"Polygon news request failed with exception: {e}")
+            polygon_weekly_news = []
+        
+        # Merge the two news lists
+        weekly_news_combined = finnhub_weekly_news + polygon_weekly_news
+        # Sort by date
+        weekly_news_combined.sort(key=lambda x: x['date'])
+        
+        # If no news from either source, you can decide to raise an error or keep empty
+        # Here we won't raise an error if both are empty, so it continues gracefully:
+        # if len(weekly_news_combined) == 0:
+        #     raise gr.Error(f"No company news found for symbol {symbol} from both Finnhub & Polygon in {start_date}-{end_date}!")
+        
+        combined_news_list.append(json.dumps(weekly_news_combined))
+
+    data['News'] = combined_news_list
     return data
 
 
-def get_company_prompt(symbol):
+###############################################################################
+# Polygon Financials
+###############################################################################
+
+def get_current_basics(symbol, curday):
+    """
+    Fetch the latest (limit=1, order=desc) financials from Polygon and parse out
+    a condensed dictionary of relevant metrics. Replaces the old Finnhub approach.
+    """
+    url = (
+        f"https://api.polygon.io/vX/reference/financials"
+        f"?ticker={symbol}"
+        f"&order=desc"
+        f"&limit=1"
+        f"&sort=filing_date"
+        f"&apiKey={polygon_api_key}"
+    )
+    resp = requests.get(url)
+    if resp.status_code != 200:
+        raise gr.Error(f"Failed to retrieve financial data from Polygon! status={resp.status_code}")
+
+    data = resp.json()
+    if 'results' not in data or len(data['results']) == 0:
+        raise gr.Error(f"No financial results found from Polygon for ticker {symbol}!")
+
+    result = data['results'][0]
+
+    # We can store a small set of relevant metrics
+    filing_date = result.get('filing_date', '')
+    fiscal_year = result.get('fiscal_year', 'N/A')
+    fiscal_period = result.get('fiscal_period', 'N/A')
+    fin = result.get('financials', {})
+
+    # Helper to safely retrieve a nested value
+    def get_nested_value(d, path):
+        # path is a list of keys, e.g. ['income_statement','revenues']
+        temp = d
+        for p in path:
+            if p not in temp:
+                return None
+            temp = temp[p]
+        # Return .get("value") if present
+        return temp.get("value", None)
+
+    # Construct a dictionary for the prompt
+    # We'll reference these in the final user prompt
+    basics = {
+        "period": f"{fiscal_year} {fiscal_period}",
+        "filing_date": filing_date,
+        "Revenue": get_nested_value(fin, ["income_statement", "revenues"]),
+        "NetIncomeLoss": get_nested_value(fin, ["income_statement", "net_income_loss_attributable_to_parent"]),
+        "DilutedEPS": get_nested_value(fin, ["income_statement", "diluted_earnings_per_share"]),
+        "BasicEPS": get_nested_value(fin, ["income_statement", "basic_earnings_per_share"]),
+    }
+
+    return basics
+
+
+###############################################################################
+# Prompt engineering
+###############################################################################
 
+def get_company_prompt(symbol):
+    """
+    Pull a minimal Finnhub company profile to produce a short introduction.
+    (We keep this from Finnhub if desired, as that info isn't replaced by Polygon.)
+    """
     profile = finnhub_client.company_profile2(symbol=symbol)
     if not profile:
         raise gr.Error(f"Failed to find company profile for symbol {symbol} from finnhub!")
         
-    company_template = "[Company Introduction]:\n\n{name} is a leading entity in the {finnhubIndustry} sector. Incorporated and publicly traded since {ipo}, the company has established its reputation as one of the key players in the market. As of today, {name} has a market capitalization of {marketCapitalization:.2f} in {currency}, with {shareOutstanding:.2f} shares outstanding." \
-        "\n\n{name} operates primarily in the {country}, trading under the ticker {ticker} on the {exchange}. As a dominant force in the {finnhubIndustry} space, the company continues to innovate and drive progress within the industry."
+    company_template = (
+        "[Company Introduction]:\n\n{name} is a leading entity in the {finnhubIndustry} sector. "
+        "Incorporated and publicly traded since {ipo}, the company has established its reputation "
+        "as one of the key players in the market. As of today, {name} has a market capitalization "
+        "of {marketCapitalization:.2f} in {currency}, with {shareOutstanding:.2f} shares outstanding.\n\n"
+        "{name} operates primarily in the {country}, trading under the ticker {ticker} on the {exchange}. "
+        "As a dominant force in the {finnhubIndustry} space, the company continues to innovate and drive "
+        "progress within the industry."
+    )
 
     formatted_str = company_template.format(**profile)
-    
     return formatted_str
 
 
 def get_prompt_by_row(symbol, row):
+    """
+    Given a single row with Start Price, End Price, News, return a textual summary
+    plus the news items.
+    """
     end_price = float(row['End Price'])
     start_price = float(row['Start Price'])
     term = 'increased' if end_price > start_price else 'decreased'
     
-    start_date = row['Start Date'] if isinstance(row['Start Date'], str) else row['Start Date'].strftime('%Y-%m-%d')
-    end_date = row['End Date'] if isinstance(row['End Date'], str) else row['End Date'].strftime('%Y-%m-%d')
+    start_date = row['Start Date'].strftime('%Y-%m-%d')
+    end_date = row['End Date'].strftime('%Y-%m-%d')
     head = f"From {start_date} to {end_date}, {symbol}'s stock price {term} from {start_price:.2f} to {end_price:.2f}. Company news during this period are listed below:\n\n"
     
+    # row["News"] is a JSON string, parse it
     news = row["News"] if isinstance(row["News"], list) else json.loads(row["News"])
-    basics = json.loads(row['Basics'])
-    
-    return head, news, basics
-
+    # We do not need 'Basics' from each row, because we fetch the latest basics only at the end
+    return head, news, None
 
 
 def sample_news(news, k=5):
-    
+    """Randomly sample up to k news items from a list."""
+    if len(news) <= k:
+        return news
     return [news[i] for i in sorted(random.sample(range(len(news)), k))]
-    
+
+
 def latest_news(news, k=5):
+    """Get up to k of the most recently dated news from 'news', sorted descending."""
     sorted_news = sorted(news, key=lambda x: x['date'], reverse=True)
     return sorted_news[:k]
 
 
-def get_current_basics(symbol, curday):
-
-    basic_financials = finnhub_client.company_basic_financials(symbol, 'all')
-    if not basic_financials['series']:
-        raise gr.Error(f"Failed to find basic financials for symbol {symbol} from finnhub!")
-        
-    final_basics, basic_list, basic_dict = [], [], defaultdict(dict)
-    
-    for metric, value_list in basic_financials['series']['quarterly'].items():
-        for value in value_list:
-            basic_dict[value['period']].update({metric: value['v']})
-
-    for k, v in basic_dict.items():
-        v.update({'period': k})
-        basic_list.append(v)
-        
-    basic_list.sort(key=lambda x: x['period'])
-    
-    for basic in basic_list[::-1]:
-        if basic['period'] <= curday:
-            break
-            
-    return basic
-    
-
 def get_all_prompts_online(symbol, data, curday, with_basics=True):
+    """
+    Build a final prompt string from:
+      - The company prompt (finnhub profile)
+      - The historical intervals from 'data'
+      - The current basics from Polygon
+    """
     company_prompt = get_company_prompt(symbol)
 
     prev_rows = []
-
-    for row_idx, row in data.iterrows():
+    for _, row in data.iterrows():
         head, news, _ = get_prompt_by_row(symbol, row)
-        prev_rows.append((head, news, None))
-        
-    prompt = ""
+        prev_rows.append((head, news))
+
+    # Build up the text describing previous intervals
+    prompt_body = ""
     for i in range(-len(prev_rows), 0):
-        prompt += "\n" + prev_rows[i][0]
-        latest_news_items = latest_news(
-            prev_rows[i][1], 
-            min(5, len(prev_rows[i][1]))
-        )
-        if latest_news_items:
-            # Ensure each news item is formatted as a string
-            formatted_news_items = ["[Headline]: {}\n[Summary]: {}\n".format(n['headline'], n['summary']) for n in latest_news_items]
-            prompt += "\n".join(formatted_news_items)
+        head, news_list = prev_rows[i]
+        prompt_body += "\n" + head
+        # Show up to 5 latest news for each interval
+        chosen_news = latest_news(news_list, min(5, len(news_list)))
+        if chosen_news:
+            # Format each news item
+            formatted_news_items = [
+                "[Headline]: {}\n[Summary]: {}\n".format(n['headline'], n['summary']) 
+                for n in chosen_news
+            ]
+            prompt_body += "\n".join(formatted_news_items)
         else:
-            prompt += "No relative news reported."
-        
-    period = "{} to {}".format(curday, n_weeks_before(curday, -1))
-    
+            prompt_body += "No relative news reported."
+
+    period = f"{curday} to {n_weeks_before(curday, -1)}"
+
+    # If user wants the latest polygon financials
     if with_basics:
         basics = get_current_basics(symbol, curday)
-        basics = "Some recent basic financials of {}, reported at {}, are presented below:\n\n[Basic Financials]:\n\n".format(
-            symbol, basics['period']) + "\n".join(f"{k}: {v}" for k, v in basics.items() if k != 'period')
+        basics_text = (
+            f"Some recent basic financials of {symbol}, reported at {basics['filing_date']} "
+            f"(period: {basics['period']}), are presented below:\n\n[Basic Financials]:\n\n"
+        )
+        # Append each metric
+        for k, v in basics.items():
+            if k not in ["period", "filing_date"]:
+                basics_text += f"{k}: {v}\n"
     else:
-        basics = "[Basic Financials]:\n\nNo basic financial reported."
+        basics_text = "[Basic Financials]:\n\nNo basic financial reported."
 
-    info = company_prompt + '\n' + prompt + '\n' + basics
-    prompt = info + f"\n\nBased on all the information before {curday}, let's first analyze the positive developments and potential concerns for {symbol}. Come up with 2-4 most important factors respectively and keep them concise. Most factors should be inferred from company related news. " \
-        f"Then make your prediction of the {symbol} stock price movement for next week ({period}). Provide a summary analysis to support your prediction."
-        
-    return info, prompt
+    info_block = company_prompt + '\n' + prompt_body + '\n' + basics_text
+    
+    # This is the final request prompt for the language model
+    final_prompt = (
+        info_block
+        + f"\n\nBased on all the information before {curday}, let's first analyze the positive developments "
+          f"and potential concerns for {symbol}. Come up with 2-4 most important factors respectively and "
+          f"keep them concise. Most factors should be inferred from company-related news. Then make your "
+          f"prediction of the {symbol} stock price movement for next week ({period}). Provide a summary analysis "
+          f"to support your prediction."
+    )
+    
+    return info_block, final_prompt
 
 
+###############################################################################
+# Gradio pipeline
+###############################################################################
 
 def construct_prompt(ticker, curday, n_weeks, use_basics):
-
+    """
+    Build the final prompt to feed into the model by collecting:
+      1) Stock data from yfinance (past n_weeks intervals)
+      2) News from Finnhub & Polygon
+      3) Current basics from Polygon (if requested)
+    """
     try:
         steps = [n_weeks_before(curday, n) for n in range(n_weeks + 1)][::-1]
     except Exception:
         raise gr.Error(f"Invalid date {curday}!")
-        
+    
+    # 1) Stock data
     data = get_stock_data(ticker, steps)
+    # 2) News data (from Finnhub & Polygon)
     data = get_news(ticker, data)
+    # We don't store the basics in each row anymore, so just place empty dict
     data['Basics'] = [json.dumps({})] * len(data)
-    # print(data)
-    
-    info, prompt = get_all_prompts_online(ticker, data, curday, use_basics)
+
+    # 3) Construct final prompt
+    info, prompt = get_all_prompts_online(ticker, data, curday, with_basics)
     
+    # Format with system instructions for Llama
     prompt = B_INST + B_SYS + SYSTEM_PROMPT + E_SYS + prompt + E_INST
-    # print(prompt)
     
     return info, prompt
 
 
-def predict(ticker, date, n_weeks, use_basics):
-
+def predict(ticker, date_, n_weeks, use_basics):
+    """
+    Main function triggered by Gradio. 
+    1) Builds the prompt, 
+    2) Generates the response via the LLaMA-2 + LoRA model,
+    3) Returns the full info block + final answer.
+    """
     print_gpu_utilization()
-
-    info, prompt = construct_prompt(ticker, date, n_weeks, use_basics)
-      
-    inputs = tokenizer(
-        prompt, return_tensors='pt', padding=False
-    )
+    
+    # Construct the full prompt
+    info, prompt = construct_prompt(ticker, date_, n_weeks, use_basics)
+    
+    # Tokenize
+    inputs = tokenizer(prompt, return_tensors='pt', padding=False)
     inputs = {key: value.to(model.device) for key, value in inputs.items()}
 
-    print("Inputs loaded onto devices.")
-        
+    print("Inputs loaded onto device...")
+
+    # Generate
     res = model.generate(
-        **inputs, max_length=4096, do_sample=False,
+        **inputs,
+        max_length=4096,
+        do_sample=False,
         eos_token_id=tokenizer.eos_token_id,
-        use_cache=True, streamer=streamer
+        use_cache=True,
+        streamer=streamer
     )
     output = tokenizer.decode(res[0], skip_special_tokens=True)
+    # Remove everything up to the closing [/INST]
     answer = re.sub(r'.*\[/INST\]\s*', '', output, flags=re.DOTALL)
 
+    # Clean up GPU memory
     torch.cuda.empty_cache()
     
     return info, answer
 
 
+################################################################################
+# Gradio Interface
+################################################################################
+
 demo = gr.Interface(
-    predict,
+    fn=predict,
     inputs=[
         gr.Textbox(
             label="Ticker",
             value="AAPL",
-            info="Companys from Dow-30 are recommended"
+            info="Companies from Dow-30 are recommended"
         ),
         gr.Textbox(
             label="Date",
             value=get_curday,
-            info="Date from which the prediction is made, use format yyyy-mm-dd"
+            info="Date from which the prediction is made, format yyyy-mm-dd"
         ),
         gr.Slider(
             minimum=1,
@@ -292,25 +505,22 @@ demo = gr.Interface(
             value=3,
             step=1,
             label="n_weeks",
-            info="Information of the past n weeks will be utilized, choose between 1 and 4"
+            info="Information of the past n weeks will be utilized. Choose between 1 and 4."
         ),
         gr.Checkbox(
             label="Use Latest Basic Financials",
             value=False,
-            info="If checked, the latest quarterly reported basic financials of the company is taken into account."
+            info="If checked, the latest quarterly reported basic financials from Polygon are taken into account."
         )
     ],
     outputs=[
-        gr.Textbox(
-            label="Information"
-        ),
-        gr.Textbox(
-            label="Response"
-        )
+        gr.Textbox(label="Information"),
+        gr.Textbox(label="Response")
     ],
     title="Pro Capital",
-    description="""Pro Capital implementation.**
-"""
+    description="Pro Capital implementation using yfinance for historical prices, "
+                "Finnhub + Polygon for news, and Polygon for financials."
 )
 
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()