Update app.py

app.py

@@ -8,7 +8,6 @@ import torch
 import yfinance as yf
 from functools import lru_cache
 import pandas as pd
-from dateutil import parser  # Added for flexible date parsing
 
 # Define the list of tickers
 tickers = ['TSLA', 'PLTR', 'SOUN', 'MSFT']
@@ -88,51 +87,153 @@ def fetch_stock_data(symbol, start_date, end_date):
         return None
 
 def parse_dates(query):
-    query_lower = query.lower()
-    start_date = None
-    end_date = None
-
     # Handle year ranges like "between 2010 and 2020"
-    range_match = re.search(r'between\s+(\d{4})\s+and\s+(\d{4})', query_lower)
+    range_match = re.search(r'between\s+(\d{4})\s+and\s+(\d{4})', query.lower())
     if range_match:
         start_year = int(range_match.group(1))
         end_year = int(range_match.group(2))
-        start_date = datetime(start_year, 1, 1)
-        end_date = datetime(end_year, 12, 31)
-
-    if not range_match:
-        # Add pattern for "from X to Y"
-        range_match = re.search(r'from\s+(\d{4})\s+to\s+(\d{4})', query_lower)
-        if range_match:
-            start_year = int(range_match.group(1))
-            end_year = int(range_match.group(2))
+        try:
             start_date = datetime(start_year, 1, 1)
             end_date = datetime(end_year, 12, 31)
+            if start_date >= end_date:
+                raise ValueError("Start date must be before end date")
+            return start_date, end_date
+        except ValueError as e:
+            print(f"Date parsing error: {e}")
+            return None, None
+    # Fallback to period parsing for recent periods
+    period_match = re.search(r'(\d+)\s*(year|month|week|day)s?', query.lower())
+    if period_match:
+        num = int(period_match.group(1))
+        unit = period_match.group(2)
+        if unit == 'year':
+            period = timedelta(days=365 * num)
+        elif unit == 'month':
+            period = timedelta(days=30 * num)
+        elif unit == 'week':
+            period = timedelta(weeks=num)
+        elif unit == 'day':
+            period = timedelta(days=num)
+        else:
+            period = timedelta(days=365)
+        end_date = datetime.now()
+        start_date = end_date - period
+        return start_date, end_date
+    # Default to 1 year
+    end_date = datetime.now()
+    start_date = end_date - timedelta(days=365)
+    return start_date, end_date
 
-    if not range_match:
-        # Add pattern for "X-Y" or "over X to Y"
-        range_match = re.search(r'(\d{4})\s*-\s*(\d{4})|over\s+(\d{4})\s+to\s+(\d{4})', query_lower)
-        if range_match:
-            start_year = int(range_match.group(1) or range_match.group(3))
-            end_year = int(range_match.group(2) or range_match.group(4))
-            start_date = datetime(start_year, 1, 1)
-            end_date = datetime(end_year, 12, 31)
+def find_closest_symbol(input_symbol):
+    input_symbol = input_symbol.upper()
+    # Check if input matches a company name
+    for company, ticker in company_to_ticker.items():
+        if company in input_symbol.lower():
+            return ticker
+    # Fallback to ticker matching
+    closest = difflib.get_close_matches(input_symbol, available_symbols, n=1, cutoff=0.6)
+    return closest[0] if closest else None
 
-    # Handle full dates if present (e.g., "2020-01-01 to 2020-12-31")
-    try:
-        dates = re.findall(r'\d{4}-\d{2}-\d{2}', query)
-        if len(dates) >= 2:
-            start_date = parser.parse(dates[0])
-            end_date = parser.parse(dates[1])
-    except:
-        pass
+def calculate_growth_rate(start_date, end_date, symbol):
+    hist = fetch_stock_data(symbol, start_date.strftime('%Y-%m-%d'), end_date.strftime('%Y-%m-%d'))
+    if hist is None or hist.empty or 'Close' not in hist.columns:
+        print(f"No valid data for {symbol} from {start_date} to {end_date}")
+        return None
+    beginning_value = hist.iloc[0]['Close']
+    ending_value = hist.iloc[-1]['Close']
+    years = (end_date - start_date).days / 365.25
+    if years <= 0:
+        return 0
+    total_return = ending_value / beginning_value
+    cagr = total_return ** (1 / years) - 1
+    return cagr * 100
 
-    # Handle relative dates (e.g., "last year")
-    if "last year" in query_lower:
-        end_date = datetime.now()
-        start_date = end_date - timedelta(days=365)
+def calculate_investment(principal, years, annual_return=0.07):
+    return principal * (1 + annual_return) ** years
+
+model_name = "HuggingFaceTB/SmolLM-135M-Instruct"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+)
+
+def generate_response(user_query, enable_thinking=False):
+    print(f"Processing query: {user_query}")  # Debugging
+    stock_keywords = ['stock', 'growth', 'investment', 'price', 'return', 'cagr']
+    is_stock_query = any(keyword in user_query.lower() for keyword in stock_keywords)
+    print(f"Is stock query: {is_stock_query}")  # Debugging
+    summary = ""
+    if is_stock_query:
+        # Try to find symbol from company name or ticker
+        symbol = None
+        for company, ticker in company_to_ticker.items():
+            if company in user_query.lower():
+                symbol = ticker
+                break
+        if not symbol:
+            symbol_match = re.search(r'\b([A-Z]{1,5})\b', user_query.upper())
+            symbol = find_closest_symbol(symbol_match.group(1)) if symbol_match else None
+        print(f"Detected symbol: {symbol}")  # Debugging
+        if symbol:
+            start_date, end_date = parse_dates(user_query)
+            print(f"Parsed dates: {start_date} to {end_date}")  # Debugging
+            if start_date is None or end_date is None:
+                summary = "Invalid date range provided. Please specify a valid range, e.g., 'between 2010 and 2020'."
+            else:
+                hist = fetch_stock_data(symbol, start_date.strftime('%Y-%m-%d'), end_date.strftime('%Y-%m-%d'))
+                print(f"Data fetched for {symbol}: { 'Valid' if hist is not None and not hist.empty else 'Empty/None'}")  # Debugging
+                if "average price" in user_query.lower():
+                    if hist is not None and not hist.empty and 'Close' in hist.columns:
+                        avg_price = hist['Close'].mean()
+                        summary = f"The average adjusted closing price for {symbol} from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')} is ${avg_price:.2f}."
+                    else:
+                        summary = f"No data available for {symbol} in the specified period."
+                elif "cagr" in user_query.lower() or "return" in user_query.lower():
+                    growth_rate = calculate_growth_rate(start_date, end_date, symbol)
+                    if growth_rate is not None:
+                        summary = f"The CAGR (geometric average annual return) for {symbol} over the period from {start_date.strftime('%Y-%m-%d')} to {end_date.strftime('%Y-%m-%d')} is {growth_rate:.2f}%."
+                    else:
+                        summary = f"No data available for {symbol} in the specified period."
+                investment_match = re.search(r'\$(\d+)', user_query)
+                if investment_match:
+                    principal = float(investment_match.group(1))
+                    years = (end_date - start_date).days / 365.25
+                    projected = calculate_investment(principal, years)
+                    summary += f" Projecting ${principal} at 7% return over {years:.1f} years: ${projected:.2f}."
+        else:
+            summary = "Could not identify a valid stock symbol in the query. Please specify a company or ticker (e.g., 'Tesla' or 'TSLA')."
+    
+    # Enhanced system prompt for better guidance
+    system_prompt = (
+        "You are FinChat, a knowledgeable financial advisor. Always respond in a friendly, professional manner. "
+        "For greetings like 'Hi' or 'Hello', reply warmly, e.g., 'Hi! I'm FinChat, your financial advisor. What can I help you with today regarding stocks, investments, or advice?' "
+        "Provide accurate, concise advice based on the provided data summary. If no data is available, suggest alternatives politely. "
+        "Use the data summary to inform your response where relevant."
+    )
+    
+    # Integrate thinking if enabled (simple chain-of-thought prompt)
+    thinking_prompt = "Think step by step before responding. " if enable_thinking else ""
+    
+    # Prepare messages for the model
+    messages = [
+        {"role": "system", "content": system_prompt},
+        {"role": "user", "content": f"{thinking_prompt}Query: {user_query}\nData summary: {summary if summary else 'No specific data computed; respond generally based on knowledge.'}"}
+    ]
+    
+    # Generate response using the model
+    input_ids = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
+    outputs = model.generate(input_ids, max_new_tokens=200, do_sample=True, temperature=0.7)
+    response = tokenizer.decode(outputs[0][input_ids.shape[1]:], skip_special_tokens=True)
+    
+    return response
 
-    if start_date and end_date and start_date < end_date:
-        return start_date.strftime('%Y-%m-%d'), end_date.strftime('%Y-%m-%d')
-    else:
-        # Default to last
+# Gradio interface setup
+demo = gr.Interface(
+    fn=generate_response,
+    inputs=[gr.Textbox(lines=2, placeholder="Enter your query (e.g., 'average return for TSLA stock between 2010 and 2020')"), gr.Checkbox(label="Enable Thinking")],
+    outputs="text",
+    title="FinChat",
+    description="Ask about stock performance, CAGR, or investments."
+)
+demo.launch()