Update app.py

app.py

@@ -1,112 +1,125 @@
 import logging
 import os
-from transformers import AutoModelForCausalLM, AutoTokenizer
 import torch
 import gradio as gr
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import difflib
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s")
 logger = logging.getLogger(__name__)
 
-# Define device (force CPU for Spaces compatibility)
+# Define device (force CPU for Spaces free tier)
 device = torch.device("cpu")
 logger.info(f"Using device: {device}")
 
-# Response cache with step-by-step advice
+# Response cache with expanded entries
 response_cache = {
     "hi": "Hello! I'm your financial advisor. How can I help with investing?",
     "hello": "Hello! I'm your financial advisor. How can I help with investing?",
     "hey": "Hi there! Ready to discuss investment goals?",
     "hi, give me step-by-step investing advice": (
         "Here’s a step-by-step guide to start investing:\n"
-        "1. **Open a Brokerage Account**: If you’re 18 or older, sign up with a platform like Fidelity, Vanguard, or Robinhood.\n"
-        "2. **Deposit Funds**: Add an initial amount you can afford, such as $100, after securing an emergency fund.\n"
-        "3. **Research and Buy**: Choose a stock, ETF (e.g., VOO for S&P 500), or index fund based on research from Yahoo Finance or Morningstar.\n"
-        "4. **Monitor Investments**: Check your portfolio regularly and enable dividend reinvesting for compounding returns.\n"
-        "5. **Use Dollar-Cost Averaging**: Invest a fixed amount (e.g., $100 monthly) consistently to reduce market timing risks.\n"
-        "6. **Diversify**: Spread investments across sectors to manage risk.\n"
-        "Consult a certified financial planner for personalized advice."
+        "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."
     ),
     "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 in the stock market:\n"
-        "1. **Learn**: Use Investopedia or 'The Intelligent Investor' by Benjamin Graham.\n"
-        "2. **Goals**: Set objectives (e.g., retirement) and assess risk tolerance.\n"
-        "3. **Brokerage**: Choose Fidelity, Vanguard, or Robinhood.\n"
-        "4. **Investments**: Start with ETFs (e.g., VOO) or mutual funds.\n"
-        "5. **Strategy**: Use dollar-cost averaging ($100-$500 monthly).\n"
-        "6. **Risks**: Diversify and monitor.\n"
-        "Consult a certified financial planner."
+        "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 cannot recommend specific companies without current data. Consider ETFs like VOO (S&P 500) or QQQ (tech). "
-        "Research sectors like technology (e.g., Apple) or healthcare (e.g., Johnson & Johnson) on Yahoo Finance. "
+        "I can’t recommend specific companies without data. Try ETFs like VOO (S&P 500) or QQQ (tech). "
+        "Research technology (e.g., Apple) or healthcare (e.g., Johnson & Johnson) on Yahoo Finance. "
         "Consult a financial planner."
     ),
     "how do i start investing in stocks?": (
-        "Educate yourself with Investopedia or 'The Intelligent Investor.' Set goals and assess risk tolerance. "
-        "Open a brokerage account with Fidelity or Vanguard and start with ETFs (e.g., VOO). Consult a financial planner."
+        "Learn from Investopedia. Set goals and assess risk. Open a brokerage account (Fidelity, Vanguard) "
+        "and start with ETFs (e.g., VOO). Consult a financial planner."
     ),
     "what's the difference between stocks and bonds?": (
-        "Stocks offer ownership in a company with growth potential but higher risk. Bonds are loans to companies/governments, "
-        "offering steady interest with lower risk. Diversify with both for balance."
+        "Stocks are company ownership with high risk and growth potential. 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 (3-6 months’ savings). Start with $100-$500 monthly "
-        "in ETFs like VOO using dollar-cost averaging. Consult a financial planner."
+        "Invest what you can afford after expenses and an emergency fund. Start with $100-$500 monthly "
+        "in ETFs (e.g., VOO) using dollar-cost averaging. Consult a financial planner."
     ),
     "what is dollar-cost averaging?": (
-        "Dollar-cost averaging is investing a fixed amount regularly (e.g., $100 monthly) in assets like ETFs, "
+        "Dollar-cost averaging is investing a fixed amount regularly (e.g., $100 monthly) in ETFs, "
         "reducing risk by spreading purchases over time."
     ),
     "give me few investing idea": (
-        "Here are some investing ideas:\n"
-        "1. Open a brokerage account if you are 18 or older (e.g., Fidelity, Vanguard).\n"
-        "2. Deposit an initial amount you can afford (e.g., $100).\n"
-        "3. Buy a researched stock, ETF (e.g., VOO), or index fund.\n"
-        "4. Check your investments regularly and enable dividend reinvesting if desired.\n"
-        "5. Use dollar-cost averaging to buy the same investment regularly (e.g., monthly).\n"
-        "Consult a financial planner for personalized advice."
+        "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) 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."
+    ),
+    "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.\n"
+        "4. Invest regularly using dollar-cost averaging.\n"
+        "5. Diversify to manage risk.\n"
+        "Consult a financial planner."
     )
 }
 
 # Load model and tokenizer
-model_name = "facebook/opt-350m"
+model_name = "distilgpt2"  # Smaller model for CPU
 try:
     logger.info(f"Loading tokenizer for {model_name}")
     tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=False)
     logger.info(f"Loading model {model_name}")
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float16
+        torch_dtype=torch.float16,
+        low_cpu_mem_usage=True,
+        load_in_4bit=True  # 4-bit quantization
     ).to(device)
 except Exception as e:
     logger.error(f"Error loading model/tokenizer: {e}")
     raise
 
-# Pre-tokenize prompt prefix with few-shot example
+# Pre-tokenize prompt prefix
 prompt_prefix = (
-    "You are a financial advisor. Provide concise, actionable advice in a numbered list for step-by-step or idea prompts. "
-    "Avoid repetition and vague statements. Use varied, specific steps.\n\n"
-    "Example:\n"
-    "Q: Give me step-by-step investing advice\n"
-    "A: 1. Open a brokerage account with Fidelity or Vanguard if 18 or older.\n"
-    "2. Deposit an affordable amount, like $100, after building an emergency fund.\n"
-    "3. Research and buy an ETF like VOO using Yahoo Finance data.\n"
-    "4. Check investments monthly and enable dividend reinvesting.\n"
-    "5. Invest regularly with dollar-cost averaging to reduce risk.\n\n"
+    "You are a financial advisor. Provide concise, numbered list advice for investing prompts. "
+    "Avoid repetition and vague statements.\n\n"
+    "Example: Q: Give investing tips\nA: 1. Open a brokerage account.\n2. Start with ETFs like VOO.\n3. Use dollar-cost averaging.\n\n"
     "Q: "
 )
 prefix_tokens = tokenizer(prompt_prefix, return_tensors="pt", truncation=True, max_length=512).to(device)
 
+# Fuzzy matching for cache
+def get_closest_cache_key(message, cache_keys, threshold=0.9):
+    matches = difflib.get_close_matches(message, cache_keys, n=1, cutoff=threshold)
+    return matches[0] if matches else None
+
 # Define chat function
 def chat_with_model(message, history=None):
     try:
         logger.info(f"Processing message: {message}")
-        # Normalize input for cache
+        # Normalize and check cache
         cache_key = message.lower().strip()
-        if cache_key in response_cache:
-            logger.info("Cache hit")
-            return response_cache[cache_key]
+        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}")
+            return response_cache[closest_key]
         
         # Skip model for short prompts
         if len(message.strip()) <= 5:
@@ -117,20 +130,21 @@ def chat_with_model(message, history=None):
         full_prompt = prompt_prefix + message + "\nA:"
         inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
         
-        # Generate response
-        with torch.no_grad():
+        # Generate response with mixed precision
+        with torch.cpu.amp.autocast(), torch.no_grad():
             outputs = model.generate(
                 **inputs,
-                max_new_tokens=100,
-                min_length=20,
+                max_new_tokens=80,
+                min_length=15,
                 do_sample=True,
                 temperature=0.7,
                 top_p=0.9,
-                no_repeat_ngram_size=2,
+                repetition_penalty=1.2,
                 pad_token_id=tokenizer.eos_token_id
             )
         response = tokenizer.decode(outputs[0], skip_special_tokens=True)
         logger.info("Generated response")
+        torch.cuda.empty_cache()  # Clear memory
         return response[len(full_prompt):].strip() if response.startswith(full_prompt) else response
     except Exception as e:
         logger.error(f"Error generating response: {e}")
@@ -140,15 +154,15 @@ def chat_with_model(message, history=None):
 logger.info("Initializing Gradio interface")
 interface = gr.ChatInterface(
     fn=chat_with_model,
-    title="Financial Advisor Chatbot (OPT-350m)",
-    description="Ask about investing! Powered by Meta AI's OPT-350m. Fast, detailed answers.",
+    title="Financial Advisor Chatbot (DistilGPT2)",
+    description="Ask about investing! Fast, detailed answers on CPU.",
     examples=[
         "Hi",
         "Hi, give me step-by-step investing advice",
         "Give me few investing idea",
+        "Give me investing tips",
         "Do you have a list of companies you recommend?",
-        "What's the difference between stocks and bonds?",
-        "How much should I invest?"
+        "What's the difference between stocks and bonds?"
     ]
 )
 
@@ -156,7 +170,7 @@ interface = gr.ChatInterface(
 if __name__ == "__main__" and not os.getenv("HF_SPACE"):
     logger.info("Launching Gradio interface locally")
     try:
-        interface.launch(share=False, debug=True)
+        interface.launch(share=False, debug=True, queue=False)
     except Exception as e:
         logger.error(f"Error launching interface: {e}")
         raise