Update app.py

app.py

@@ -4,6 +4,7 @@ import time
 import torch
 import gradio as gr
 from transformers import AutoModelForCausalLM, AutoTokenizer
+import difflib
 import json
 
 # Set up logging
@@ -14,85 +15,126 @@ logger = logging.getLogger(__name__)
 device = torch.device("cpu")
 logger.info(f"Using device: {device}")
 
-# Expanded response cache with new entries
+# Response cache with financial data entries
 response_cache = {
-    "hi": "Hello! I'm FinChat, your financial advisor. How can I help with investing today?",
-    "hello": "Hello! I'm FinChat, your financial advisor. How can I help with investing today?",
+    "hi": "Hello! I'm FinChat, your financial advisor. How can I help with investing?",
+    "hello": "Hello! I'm FinChat, your financial advisor. How can I help with investing?",
     "hey": "Hi there! Ready to discuss investment goals with FinChat?",
-    "how can i start investing with $100 a month?": (
-        "Here’s a step-by-step guide to start investing with $100 a month:\n"
-        "1. **Open a brokerage account** with a platform like Fidelity or Robinhood. They offer low fees and no minimums.\n"
-        "2. **Deposit your $100 monthly**. You can set up automatic transfers from your bank.\n"
-        "3. **Choose a low-cost ETF** like VOO, which tracks the S&P 500 for broad market exposure.\n"
-        "4. **Set up automatic investments** to buy shares regularly, reducing the impact of market fluctuations.\n"
-        "5. **Track your progress** every few months to stay on top of your investments.\n"
-        "Consult a financial planner for personalized advice."
+    "what is better individual stocks or etfs?": (
+        "Here’s a comparison of individual stocks vs. ETFs:\n"
+        "1. **Individual Stocks**: High returns possible (e.g., Apple up 80% in 2020) but riskier due to lack of diversification. Require active research.\n"
+        "2. **ETFs**: Diversify risk by tracking indices (e.g., VOO, S&P 500, ~12% avg. return 2015–2025). Lower fees and less research needed.\n"
+        "3. **Recommendation**: Beginners should start with ETFs; experienced investors may add stocks.\n"
+        "Consult a financial planner."
     ),
-    "where can i open a brokerage account?": (
-        "You can open a brokerage account with platforms like Vanguard, Fidelity, Charles Schwab, or Robinhood. "
-        "They are beginner-friendly and offer low fees. Choose one that fits your needs and sign up online."
+    "is $100 per month enough to invest?": (
+        "Yes, $100 per month is enough to start investing. Here’s why and how:\n"
+        "1. **Feasibility**: Brokerages like Fidelity have no minimums, and commission-free trading eliminates fees.\n"
+        "2. **Options**: Buy fractional shares of ETFs (e.g., VOO, ~$500/share) with $100.\n"
+        "3. **Strategy**: Use dollar-cost averaging to invest monthly, reducing market timing risks.\n"
+        "4. **Growth**: At 10% annual return, $100 monthly could grow to ~$41,000 in 20 years.\n"
+        "5. **Tips**: Ensure an emergency fund; diversify.\n"
+        "Consult a financial planner."
     ),
-    "start investing with 100 dollars a month": (
-        "Here’s how to start investing with $100 a month:\n"
-        "1. **Open a brokerage account** with a platform like Fidelity or Robinhood.\n"
-        "2. **Deposit $100 monthly** via automatic transfers.\n"
-        "3. **Invest in a low-cost ETF** like VOO for diversification.\n"
-        "4. **Use dollar-cost averaging** to invest regularly.\n"
-        "5. **Monitor your investments** quarterly.\n"
-        "Consult a financial planner for tailored advice."
+    "can i invest $100 a month?": (
+        "Yes, $100 a month is sufficient. Here’s how:\n"
+        "1. **Brokerage**: Open an account with Fidelity or Vanguard (no minimums).\n"
+        "2. **Investments**: Buy fractional shares of ETFs like VOO ($100 buys ~0.2 shares).\n"
+        "3. **Approach**: Use dollar-cost averaging for steady growth.\n"
+        "4. **Long-Term**: At 10% return, $100 monthly could reach ~$41,000 in 20 years.\n"
+        "5. **Tips**: Prioritize an emergency fund and diversify.\n"
+        "Consult a financial planner."
     ),
-    "best places to open a brokerage account": (
-        "The best places to open a brokerage account include Vanguard, Fidelity, Charles Schwab, and Robinhood. "
-        "They offer low fees, no minimums, and user-friendly platforms for beginners."
+    "hi, give me step-by-step investing advice": (
+        "Here’s a step-by-step guide to start investing:\n"
+        "1. Open a brokerage account (e.g., Fidelity, Vanguard) if 18 or older.\n"
+        "2. Deposit an affordable amount, like $100, after an emergency fund.\n"
+        "3. Research and buy an ETF (e.g., VOO) using Yahoo Finance.\n"
+        "4. Monitor monthly and enable dividend reinvesting.\n"
+        "5. Use dollar-cost averaging ($100 monthly) to reduce risk.\n"
+        "6. Diversify across sectors.\n"
+        "Consult a financial planner."
     ),
-    "what is dollar-cost averaging?": (
-        "Dollar-cost averaging is investing a fixed amount regularly (e.g., $100 monthly) in ETFs, "
-        "reducing risk by spreading purchases over time."
+    "hi, pretend you are a financial advisor. now tell me how can i start investing in stock market?": (
+        "Here’s a guide to start investing:\n"
+        "1. Learn from Investopedia or 'The Intelligent Investor.'\n"
+        "2. Set goals (e.g., retirement) and assess risk.\n"
+        "3. Choose a brokerage (Fidelity, Vanguard).\n"
+        "4. Start with ETFs (e.g., VOO) or mutual funds.\n"
+        "5. Use dollar-cost averaging ($100-$500 monthly).\n"
+        "6. Diversify and monitor.\n"
+        "Consult a financial planner."
+    ),
+    "do you have a list of companies you recommend?": (
+        "I can’t recommend specific companies without data. Try ETFs like VOO (S&P 500, ~12% avg. return 2015–2025) or QQQ (tech). "
+        "Research stocks like Apple (AAPL, ~80% return in 2020) or Johnson & Johnson on Yahoo Finance.\n"
+        "Consult a financial planner."
+    ),
+    "how do i start investing in stocks?": (
+        "Learn from Investopedia. Set goals and assess risk. Open a brokerage account (Fidelity, Vanguard) "
+        "and start with ETFs (e.g., VOO, ~10% avg. return). Consult a financial planner."
+    ),
+    "what's the difference between stocks and bonds?": (
+        "Stocks are company ownership with high risk and growth potential (e.g., S&P 500 ~10% avg. return). Bonds are loans to companies/governments "
+        "with lower risk and steady interest. Diversify for balance."
     ),
     "how much should i invest?": (
         "Invest what you can afford after expenses and an emergency fund. Start with $100-$500 monthly "
-        "in ETFs like VOO using dollar-cost averaging. Consult a financial planner."
+        "in ETFs like VOO (~10% avg. return). Consult a financial planner."
     ),
-    "how to start investing": (
-        "Here’s how to start investing:\n"
-        "1. Educate yourself using resources like Investopedia.\n"
-        "2. Open a brokerage account with a platform like Fidelity.\n"
-        "3. Deposit an initial amount, such as $100, after building an emergency fund.\n"
-        "4. Choose a low-cost ETF like VOO.\n"
-        "5. Invest regularly using dollar-cost averaging.\n"
-        "Consult a financial planner for personalized advice."
+    "what is dollar-cost averaging?": (
+        "Dollar-cost averaging is investing a fixed amount regularly (e.g., $100 monthly) in ETFs, "
+        "reducing risk by spreading purchases over time."
     ),
-    "best brokerage accounts": (
-        "The best brokerage accounts for beginners include Fidelity, Vanguard, Charles Schwab, and Robinhood. "
-        "They offer low fees, no minimums, and user-friendly platforms."
+    "give me few investing idea": (
+        "Here are investing ideas:\n"
+        "1. Open a brokerage account (e.g., Fidelity) if 18 or older.\n"
+        "2. Deposit $100 or what you can afford.\n"
+        "3. Buy a researched ETF (e.g., VOO, ~10% avg. return) or index fund.\n"
+        "4. Check regularly and enable dividend reinvesting.\n"
+        "5. Use dollar-cost averaging (e.g., monthly buys).\n"
+        "Consult a financial planner."
     ),
-    "investing for beginners": (
-        "Here’s a beginner’s guide to investing:\n"
-        "1. Learn the basics from Investopedia or books like 'The Intelligent Investor.'\n"
-        "2. Set clear investment goals and assess your risk tolerance.\n"
-        "3. Open a brokerage account with a platform like Fidelity or Robinhood.\n"
-        "4. Start with low-cost ETFs like VOO or index funds.\n"
-        "5. Invest regularly using dollar-cost averaging.\n"
-        "6. Monitor your investments quarterly.\n"
-        "Consult a financial planner for tailored advice."
+    "give me investing tips": (
+        "Here are investing tips:\n"
+        "1. Educate yourself with Investopedia or books.\n"
+        "2. Open a brokerage account (e.g., Vanguard).\n"
+        "3. Start small with ETFs like VOO (~10% avg. return).\n"
+        "4. Invest regularly using dollar-cost averaging.\n"
+        "5. Diversify to manage risk.\n"
+        "Consult a financial planner."
     ),
-    "steps to start investing": (
-        "Here are the steps to start investing:\n"
-        "1. Educate yourself on investing basics.\n"
-        "2. Open a brokerage account with a beginner-friendly platform.\n"
-        "3. Deposit an initial amount you can afford.\n"
-        "4. Choose a diversified investment like an ETF.\n"
-        "5. Invest consistently over time.\n"
-        "Consult a financial planner for more guidance."
+    "how to start investing": (
+        "Here’s how to start investing:\n"
+        "1. Study basics on Investopedia.\n"
+        "2. Open a brokerage account (e.g., Fidelity).\n"
+        "3. Deposit $100 or more after securing savings.\n"
+        "4. Buy an ETF like VOO (~10% avg. return) after research.\n"
+        "5. Invest monthly with dollar-cost averaging.\n"
+        "Consult a financial planner."
     ),
-    "recommended etfs": (
-        "Recommended ETFs for beginners include VOO (tracks S&P 500), QQQ (tech-focused), and VT (global market exposure). "
-        "They offer diversification and low fees."
+    "investing advice": (
+        "Here’s investing advice:\n"
+        "1. Learn basics from Investopedia.\n"
+        "2. Open a brokerage account (e.g., Vanguard).\n"
+        "3. Start with $100 in an ETF like VOO (~10% avg. return).\n"
+        "4. Use dollar-cost averaging for regular investments.\n"
+        "5. Monitor and diversify your portfolio.\n"
+        "Consult a financial planner."
     ),
-    "how much to invest": (
-        "The amount to invest depends on your financial situation. Start with what you can afford after covering expenses and an emergency fund. "
-        "A common starting point is $100-$500 monthly in low-cost ETFs like VOO. Consult a financial planner for personalized advice."
+    "steps to invest": (
+        "Here are steps to invest:\n"
+        "1. Educate yourself using Investopedia.\n"
+        "2. Open a brokerage account (e.g., Fidelity).\n"
+        "3. Deposit an initial $100 after savings.\n"
+        "4. Buy an ETF like VOO (~10% avg. return) after research.\n"
+        "5. Use dollar-cost averaging monthly.\n"
+        "Consult a financial planner."
     ),
+    "what is the s&p 500 index fund average growth rate?": (
+        "The S&P 500 index fund’s average annual return is approximately 10–12% over the long term (1927–2025), including dividends, based on historical data. "
+        "For example, from 2015 to 2025, it averaged ~12% annually. Returns vary yearly due to market conditions. Consult a financial planner."
+    )
 }
 
 # Load persistent cache
@@ -105,11 +147,12 @@ try:
 except Exception as e:
     logger.warning(f"Failed to load cache.json: {e}")
 
-# Load model and tokenizer
-model_name = "distilgpt2"
+# Load model and tokenizer (use fine-tuned model if available)
+model_name = "./finetuned_model" if os.path.exists("./finetuned_model") else "distilgpt2"
 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():
         model = AutoModelForCausalLM.from_pretrained(
@@ -121,47 +164,58 @@ except Exception as e:
     logger.error(f"Error loading model/tokenizer: {e}")
     raise RuntimeError(f"Failed to load model: {str(e)}")
 
-# Updated prompt prefix with better instructions and examples
+# Pre-tokenize prompt prefix
 prompt_prefix = (
-    "You are FinChat, a financial advisor. Always provide clear, step-by-step answers to the user's exact question. "
-    "Avoid vague or unrelated topics. Use a numbered list format where appropriate and explain each step.\n\n"
+    "You are FinChat, a financial advisor with expertise in stock market performance. Provide detailed, numbered list advice with clear reasoning for investing prompts, "
+    "including historical data when relevant (e.g., S&P 500 returns). Avoid repetition and incomplete answers. Explain why each step or choice is beneficial.\n\n"
     "Example 1:\n"
-    "Q: How can I start investing with $100 a month?\n"
-    "A: Here’s a step-by-step guide:\n"
-    "1. Open a brokerage account with a platform like Fidelity or Robinhood. They offer low fees and no minimums.\n"
-    "2. Deposit your $100 monthly. You can set up automatic transfers.\n"
-    "3. Choose a low-cost ETF like VOO, which tracks the S&P 500.\n"
-    "4. Set up automatic investments to buy shares regularly.\n"
-    "5. Track your progress every few months.\n\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"
+    "1. This reflects historical data adjusted for inflation and dividends.\n"
+    "2. Returns vary yearly (e.g., 16.3% in 2020) due to market conditions.\n"
+    "3. ETFs like VOO track this index for broad market exposure.\n\n"
     "Example 2:\n"
-    "Q: Where can I open a brokerage account?\n"
-    "A: You can open an account with platforms like Vanguard, Fidelity, Charles Schwab, or Robinhood. "
-    "They are beginner-friendly and have low fees.\n\n"
+    "Q: Can I invest $100 a month?\n"
+    "A: Yes, $100 a month is sufficient. Here’s how:\n"
+    "1. Open a brokerage account (e.g., Fidelity): No minimums allow small investments.\n"
+    "2. Buy fractional ETF shares (e.g., VOO, ~12% avg. return 2015–2025): Diversifies risk.\n"
+    "3. Use dollar-cost averaging: Reduces market timing risks.\n\n"
     "Q: "
 )
+prefix_tokens = tokenizer(prompt_prefix, return_tensors="pt", truncation=True, max_length=512).to(device)
+
+# Substring matching for cache
+def get_closest_cache_key(message, cache_keys):
+    message = message.lower().strip()
+    for key in cache_keys:
+        if key in message:
+            return key
+    return None
 
-# Define chat function with substring matching and reduced max_new_tokens
-def chat_with_model(user_input, history=None):
+# Define chat function
+def chat_with_model(user_input, history=None, is_processing=False):
     try:
         start_time = time.time()
         logger.info(f"Processing user input: {user_input}")
-        
-        user_input_lower = user_input.lower().strip()
-        
-        # Substring matching for cache
-        matching_keys = [key for key in response_cache if key in user_input_lower]
-        if matching_keys:
-            longest_key = max(matching_keys, key=len)
-            logger.info(f"Cache hit for: {longest_key}")
-            response = response_cache[longest_key]
+        is_processing = True
+        logger.info("Showing loading animation")
+
+        # 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)
+        if closest_key:
+            logger.info(f"Cache hit for: {closest_key}")
+            response = response_cache[closest_key]
             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
-        
+            return response, history, False, ""
+
+        # 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!"
@@ -171,39 +225,54 @@ def chat_with_model(user_input, history=None):
             history.append({"role": "assistant", "content": response})
             end_time = time.time()
             logger.info(f"Response time: {end_time - start_time:.2f} seconds")
-            return response, history
+            return response, history, False, ""
 
+        # Construct prompt
         full_prompt = prompt_prefix + user_input + "\nA:"
-        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
-        
+        try:
+            inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512).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=50,  # Reduced for faster generation
+                max_new_tokens=50,
                 min_length=20,
-                do_sample=False,  # Greedy decoding for speed
+                do_sample=False,
                 repetition_penalty=1.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 with exact user input as key
-        response_cache[user_input_lower] = response
+
+        # Update cache
+        response_cache[cache_key] = response
         logger.info("Cache miss, added to in-memory cache")
-        
+
+        # Update history
         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"Total response time: {end_time - start_time:.2f} seconds")
-        return response, history
+        logger.info(f"Response time: {end_time - start_time:.2f} seconds")
+        return response, history, False, ""
+
     except Exception as e:
         logger.error(f"Error generating response: {e}")
         response = f"Error: {str(e)}"
@@ -211,38 +280,74 @@ def chat_with_model(user_input, history=None):
         history = history or []
         history.append({"role": "user", "content": user_input})
         history.append({"role": "assistant", "content": response})
-        return response, history
+        end_time = time.time()
+        logger.info(f"Response time: {end_time - start_time:.2f} seconds")
+        return response, history, False, ""
 
-# Create Gradio interface
-with gr.Blocks(
-    title="FinChat: An LLM based on distilgpt2 model",
-    css=".feedback {display: flex; gap: 10px; justify-content: center; margin-top: 10px;}"
-) as interface:
-    gr.Markdown(
-        """
-        # FinChat: An LLM based on distilgpt2 model
-        FinChat provides financial advice using the lightweight distilgpt2 model, optimized for fast, detailed responses.
-        Ask about investing strategies, ETFs, stocks, or budgeting to get started!
+# Save cache on exit
+def save_cache():
+    try:
+        with open(cache_file, 'w') as f:
+            json.dump(response_cache, f, indent=2)
+        logger.info("Saved cache to cache.json")
+    except Exception as e:
+        logger.warning(f"Failed to save cache.json: {e}")
+
+# Create Gradio interface with loading animation
+logger.info("Initializing Gradio interface")
+try:
+    with gr.Blocks(
+        title="FinChat: An LLM based on distilgpt2 model",
+        css="""
+        .loader {
+            border: 5px solid #f3f3f3;
+            border-top: 5px solid #3498db;
+            border-radius: 50%;
+            width: 30px;
+            height: 30px;
+            animation: spin 1s linear infinite;
+            margin: 10px auto;
+            display: block;
+        }
+        @keyframes spin {
+            0% { transform: rotate(0deg); }
+            100% { transform: rotate(360deg); }
+        }
+        .hidden { display: none; }
         """
-    )
-    chatbot = gr.Chatbot(type="messages")
-    msg = gr.Textbox(label="Your message")
-    submit = gr.Button("Send")
-    clear = gr.Button("Clear")
-    
-    def submit_message(user_input, history):
-        response, updated_history = chat_with_model(user_input, history)
-        return "", updated_history  # Clear input, update chatbot
+    ) as interface:
+        gr.Markdown(
+            """
+            # FinChat: An LLM based on distilgpt2 model
+            FinChat provides financial advice using the lightweight distilgpt2 model, optimized for fast, detailed responses.
+            Ask about investing strategies, ETFs, stocks, or budgeting to get started!
+            """
+        )
+        chatbot = gr.Chatbot(type="messages")
+        msg = gr.Textbox(label="Your message")
+        submit = gr.Button("Send")
+        clear = gr.Button("Clear")
+        loading = gr.HTML('<div class="loader hidden"></div>', label="Loading")
+        is_processing = gr.State(value=False)
 
-    submit.click(
-        fn=submit_message,
-        inputs=[msg, chatbot],
-        outputs=[msg, chatbot]
-    )
-    clear.click(
-        fn=lambda: ("", []),  # Clear input and chatbot
-        outputs=[msg, chatbot]
-    )
+        def submit_message(user_input, history, is_processing):
+            response, updated_history, new_processing, clear_input = chat_with_model(user_input, history, is_processing)
+            loader_html = '<div class="loader"></div>' if new_processing else '<div class="loader hidden"></div>'
+            return clear_input, updated_history, loader_html, new_processing
+
+        submit.click(
+            fn=submit_message,
+            inputs=[msg, chatbot, is_processing],
+            outputs=[msg, chatbot, loading, is_processing]
+        )
+        clear.click(
+            fn=lambda: ("", [], '<div class="loader hidden"></div>', False),
+            outputs=[msg, chatbot, loading, is_processing]
+        )
+    logger.info("Gradio interface initialized successfully")
+except Exception as e:
+    logger.error(f"Error initializing Gradio interface: {e}")
+    raise
 
 # Launch interface (conditional for Spaces)
 if __name__ == "__main__" and not os.getenv("HF_SPACE"):
@@ -252,5 +357,9 @@ if __name__ == "__main__" and not os.getenv("HF_SPACE"):
     except Exception as e:
         logger.error(f"Error launching interface: {e}")
         raise
+    finally:
+        save_cache()
 else:
-    logger.info("Running in Hugging Face Spaces, interface defined but not launched")
+    logger.info("Running in Hugging Face Spaces, interface defined but not launched")
+    import atexit
+    atexit.register(save_cache)