Update app.py

app.py

@@ -6,6 +6,7 @@ import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
 import pandas as pd
 import re
+import numpy as np
 import json
 import difflib
 from functools import lru_cache
@@ -48,54 +49,75 @@ else:
 # Hardcoded fallback for recent periods if dataset is incomplete
 fallback_returns = {
     (2020, 2022): 8.3,  # Average annual return based on external data
-    (2015, 2024): 12.2
+    (2015, 2024): 12.2,
+    (2020, 2024): 10.5
 }
 
 # Load model and tokenizer at startup
 model_name = "./finetuned_model" if os.path.exists("./finetuned_model") else "distilgpt2"
-tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=False)
-tokenizer.pad_token = tokenizer.eos_token
-model = AutoModelForCausalLM.from_pretrained(model_name).eval()
-model = model.to_bettertransformer()  # Enable BetterTransformer
-sample_input = tokenizer("Average return of S&P 500", return_tensors="pt")["input_ids"]
-traced_model = torch.jit.trace(model, sample_input)
-traced_model.save("distilgpt2_traced.pt")
-model = torch.jit.load("distilgpt2_traced.pt")
+try:
+    logger.info(f"Loading tokenizer for {model_name}")
+    tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=False)
+    tokenizer.pad_token = tokenizer.eos_token
+    logger.info(f"Loading model {model_name}")
+    with torch.inference_mode():
+        if os.path.exists("./finetuned_model/distilgpt2_traced.pt"):
+            model = torch.jit.load("./finetuned_model/distilgpt2_traced.pt")
+        else:
+            model = AutoModelForCausalLM.from_pretrained(model_name, torch_dtype=torch.float16, low_cpu_mem_usage=True).to(device)
+            sample_input = tokenizer("What was the average annual return of the S&P 500 between 2020 and 2022?", return_tensors="pt")["input_ids"].to(device)
+            model = torch.jit.trace(model, sample_input)
+            model.save("./finetuned_model/distilgpt2_traced.pt")
+    logger.info(f"Successfully loaded model: {model_name}")
+except Exception as e:
+    logger.error(f"Error loading model/tokenizer: {e}")
+    raise RuntimeError(f"Failed to load model: {str(e)}")
 
-# Response cache with financial data entries
-response_cache = {
-    "hi": "Hello! I'm FinChat, your financial advisor. How can I help with investing?",
-    "what is the average return rate of the s&p 500 in the past 10 years?": (
-        "The S&P 500’s average annual return rate from 2015 to 2024 was approximately 12.2%, including dividends, based on historical data."
-    ),
-    "what was the average annual return of the s&p 500 between 2020 and 2022?": (
-        "The S&P 500’s average annual return from 2020 to 2022 was approximately 8.3%, including dividends, with significant volatility due to the COVID-19 recovery and 2022 bear market."
-    )
-}
+# Pre-tokenize prompt prefix
+prompt_prefix = (
+    "You are FinChat, a financial advisor with expertise in stock market performance. Provide concise, accurate answers with historical data for S&P 500 queries. "
+    "For period-specific queries, use precise year ranges and calculate average annual returns. For investment return queries, use compound interest calculations "
+    "based on historical averages. Avoid repetition and ensure answers are relevant.\n\n"
+    "Example 1:\n"
+    "Q: What is the S&P 500’s average annual return?\n"
+    "A: The S&P 500’s average annual return is ~10–12% over the long term (1927–2025), including dividends.\n\n"
+    "Example 2:\n"
+    "Q: What will $5,000 be worth in 10 years if invested in the S&P 500?\n"
+    "A: Assuming a 10% average annual return, a $5,000 investment in the S&P 500 would grow to approximately $12,974 in 10 years with annual compounding.\n\n"
+    "Example 3:\n"
+    "Q: What was the average annual return of the S&P 500 between 2020 and 2022?\n"
+    "A: The S&P 500’s average annual return from 2020 to 2022 was approximately 8.3%, including dividends, with significant volatility due to the COVID-19 recovery and 2022 bear market.\n\n"
+    "Q: "
+)
+prefix_tokens = tokenizer(prompt_prefix, return_tensors="pt", truncation=True, max_length=512)["input_ids"].to(device)
 
 # Substring matching for cache with exact year matching
-def get_closest_cache_key(message, cache_keys):
+@lru_cache(maxsize=100)
+def get_closest_cache_key(message):
     message = message.lower().strip()
     year_match = re.search(r'(\d{4})\s*(?:and|to|-|–)\s*(\d{4})', message)
     if year_match:
         start_year, end_year = year_match.groups()
-        for key in cache_keys:
+        for key in response_cache.keys():
             if f"{start_year} and {end_year}" in key or f"{start_year} to {end_year}" in key or f"{start_year}–{end_year}" in key:
                 return key
-    matches = difflib.get_close_matches(message, cache_keys, n=1, cutoff=0.7)
+    matches = difflib.get_close_matches(message, response_cache.keys(), n=1, cutoff=0.7)
     return matches[0] if matches else None
 
 # Parse period from user input
 def parse_period(query):
+    # Match specific year ranges (e.g., "between 2020 and 2022", "2020–2022")
     match = re.search(r'(?:between|from)\s*(\d{4})\s*(?:and|to|-|–)\s*(\d{4})', query, re.IGNORECASE)
     if match:
         start_year, end_year = map(int, match.groups())
         return start_year, end_year, None
+    # Match duration-based queries (e.g., "1-year from 2020", "3-year growth rate")
     match = re.search(r'(\d+)-year.*from\s*(\d{4})', query, re.IGNORECASE)
     if match:
         duration, start_year = map(int, match.groups())
         end_year = start_year + duration - 1
         return start_year, end_year, duration
+    # Match general duration queries (e.g., "past 5 years", "3-year growth rate")
     match = re.search(r'past\s*(\d+)-year|\b(\d+)-year.*(?:return|growth\s*rate)', query, re.IGNORECASE)
     if match:
         duration = int(match.group(1) or match.group(2))
@@ -109,7 +131,10 @@ def parse_period(query):
 def calculate_growth_rate(start_year, end_year, duration=None):
     if (start_year, end_year) in fallback_returns:
         avg_return = fallback_returns[(start_year, end_year)]
-        response = f"The S&P 500’s average annual return from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
+        if duration:
+            response = f"The S&P 500’s {duration}-year average annual return from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
+        else:
+            response = f"The S&P 500’s average annual return from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
         return avg_return, response
     if df_yearly is None or start_year is None or end_year is None:
         return None, "Data not available or invalid period."
@@ -117,9 +142,33 @@ def calculate_growth_rate(start_year, end_year, duration=None):
     if df_period.empty:
         return None, f"No data available for {start_year} to {end_year}."
     avg_return = df_period['Return'].mean()
-    response = f"The S&P 500’s average annual return from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
+    if duration:
+        response = f"The S&P 500’s {duration}-year average annual return from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
+    else:
+        response = f"The S&P 500’s average annual return from {start_year} to {end_year} was approximately {avg_return:.1f}%, including dividends."
     return avg_return, response
 
+# Parse investment return query
+def parse_investment_query(query):
+    match = re.search(r'\$(\d+).*\s(\d+)\s*years?.*\bs&p\s*500', query, re.IGNORECASE)
+    if match:
+        amount = float(match.group(1))
+        years = int(match.group(2))
+        return amount, years
+    return None, None
+
+# Calculate future value
+def calculate_future_value(amount, years):
+    if df_yearly is None or amount is None or years is None:
+        return None, "Data not available or invalid input."
+    avg_annual_return = 10.0  # Historical S&P 500 average (1927–2025)
+    future_value = amount * (1 + avg_annual_return / 100) ** years
+    return future_value, (
+        f"Assuming a 10% average annual return, a ${amount:,.0f} investment in the S&P 500 would grow to approximately ${future_value:,.0f} "
+        f"in {years} years with annual compounding. This is based on the historical average return of 10–12% (1927–2025). "
+        "Future returns vary and are not guaranteed. Consult a financial planner."
+    )
+
 # Define chat function
 def chat_with_model(user_input, history=None, is_processing=False):
     try:
@@ -130,8 +179,7 @@ def chat_with_model(user_input, history=None, is_processing=False):
 
         # Normalize and check cache
         cache_key = user_input.lower().strip()
-        cache_keys = list(response_cache.keys())
-        closest_key = cache_key if cache_key in response_cache else get_closest_cache_key(cache_key, cache_keys)
+        closest_key = get_closest_cache_key(cache_key)
         if closest_key:
             logger.info(f"Cache hit for: {closest_key}")
             response = response_cache[closest_key]
@@ -143,10 +191,25 @@ def chat_with_model(user_input, history=None, is_processing=False):
             logger.info(f"Response time: {end_time - start_time:.2f} seconds")
             return response, history, False, ""
 
+        # Check for investment return query
+        amount, years = parse_investment_query(user_input)
+        if amount and years:
+            future_value, response = calculate_future_value(amount, years)
+            if future_value is not None:
+                response_cache[cache_key] = response
+                logger.info(f"Investment query: ${amount} for {years} years, added to cache")
+                logger.info(f"Chatbot response: {response}")
+                history = history or []
+                history.append({"role": "user", "content": user_input})
+                history.append({"role": "assistant", "content": response})
+                end_time = time.time()
+                logger.info(f"Response time: {end_time - start_time:.2f} seconds")
+                return response, history, False, ""
+
         # Check for period-specific query
-        start_year, end_year, carica_duration = parse_period(user_input)
+        start_year, end_year, duration = parse_period(user_input)
         if start_year and end_year:
-            avg_return, response = calculate_growth_rate(start_year, end_year, carica_duration)
+            avg_return, response = calculate_growth_rate(start_year, end_year, duration)
             if avg_return is not None:
                 response_cache[cache_key] = response
                 logger.info(f"Dynamic period query: {start_year}–{end_year}, added to cache")
@@ -158,15 +221,50 @@ def chat_with_model(user_input, history=None, is_processing=False):
                 logger.info(f"Response time: {end_time - start_time:.2f} seconds")
                 return response, history, False, ""
 
-        # Model inference
-        inputs = tokenizer(user_input, return_tensors="pt")
-        outputs = model.generate(
-            **inputs,
-            max_new_tokens=30,
-            repetition_penalty=2.5,
-            no_repeat_ngram_size=2
-        )
+        # Skip model for short prompts
+        if len(user_input.strip()) <= 5:
+            logger.info("Short prompt, returning default response")
+            response = "Hello! I'm FinChat, your financial advisor. Ask about investing!"
+            logger.info(f"Chatbot response: {response}")
+            history = history or []
+            history.append({"role": "user", "content": user_input})
+            history.append({"role": "assistant", "content": response})
+            end_time = time.time()
+            logger.info(f"Response time: {end_time - start_time:.2f} seconds")
+            return response, history, False, ""
+
+        # Construct prompt
+        full_prompt = prompt_prefix + user_input + "\nA:"
+        try:
+            inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512)["input_ids"].to(device)
+        except Exception as e:
+            logger.error(f"Error tokenizing input: {e}")
+            response = f"Error: Failed to process input: {str(e)}"
+            logger.info(f"Chatbot response: {response}")
+            history = history or []
+            history.append({"role": "user", "content": user_input})
+            history.append({"role": "assistant", "content": response})
+            end_time = time.time()
+            logger.info(f"Response time: {end_time - start_time:.2f} seconds")
+            return response, history, False, ""
+
+        # Generate response
+        with torch.inference_mode():
+            logger.info("Generating response with model")
+            gen_start_time = time.time()
+            outputs = model.generate(
+                inputs,
+                max_new_tokens=20,
+                min_length=10,
+                do_sample=False,
+                repetition_penalty=3.0,
+                no_repeat_ngram_size=2,
+                pad_token_id=tokenizer.eos_token_id
+            )
+            gen_end_time = time.time()
+            logger.info(f"Generation time: {gen_end_time - gen_start_time:.2f} seconds")
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+        response = response[len(full_prompt):].strip() if response.startswith(full_prompt) else response
         logger.info(f"Chatbot response: {response}")
 
         # Update cache
@@ -177,6 +275,7 @@ def chat_with_model(user_input, history=None, is_processing=False):
         history = history or []
         history.append({"role": "user", "content": user_input})
         history.append({"role": "assistant", "content": response})
+        torch.cuda.empty_cache()
         end_time = time.time()
         logger.info(f"Response time: {end_time - start_time:.2f} seconds")
         return response, history, False, ""
@@ -195,7 +294,7 @@ def chat_with_model(user_input, history=None, is_processing=False):
 # Save cache on exit
 def save_cache():
     try:
-        with open("cache.json", 'w') as f:
+        with open("cache.json", "w") as f:
             json.dump(response_cache, f, indent=2)
         logger.info("Saved cache to cache.json")
     except Exception as e: