Update app.py

app.py

@@ -1,90 +1,134 @@
-from transformers import AutoModelForCausalLM, AutoTokenizer
-import torch
+import os
 import gradio as gr
+import hashlib
+import asyncio
+import pickle
+import time
+from openai import AsyncOpenAI
+from functools import lru_cache
 
-# Define device
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+# Persistent cache configuration
+CACHE_FILE = "response_cache.pkl"
 
-# Response cache
-response_cache = {
-    "Hi, pretend you are a financial advisor. Now tell me how can I start investing in stock market?": (
-        "As a financial advisor, 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 (low fees), Vanguard (index funds like VTI), or Robinhood (commission-free).\n"
-        "4. **Investments**: Start with ETFs (e.g., VOO for S&P 500) or mutual funds.\n"
-        "5. **Strategy**: Use dollar-cost averaging with $100-$500 monthly.\n"
-        "6. **Risks**: Diversify and monitor.\n"
-        "Consult a certified financial planner."
-    ),
-    "do you have a list of companies you recommend?": (
-        "I cannot recommend specific companies without current market data. Instead, consider ETFs like VOO (S&P 500) or QQQ (tech-focused) for broad exposure. "
-        "For stocks, research sectors like technology (e.g., Apple, Microsoft) or consumer goods (e.g., Procter & Gamble) using Yahoo Finance or Morningstar. "
-        "Consult a certified financial planner."
-    ),
-    "can you provide me a list of companies you recommend?": (
-        "I cannot provide specific company recommendations without up-to-date market analysis. For safer investments, consider ETFs like VOO (S&P 500) or QQQ (tech-focused). "
-        "If interested in stocks, explore stable companies in technology (e.g., Apple, Microsoft) or healthcare (e.g., Johnson & Johnson) using Yahoo Finance. "
-        "Always consult a financial planner for tailored advice."
-    ),
-    "You have a list of companies you recommend?": (
-        "I cannot recommend specific companies without current market data. Instead, consider ETFs like VOO (S&P 500) or QQQ (tech-focused) for broad exposure. "
-        "For stocks, research sectors like technology (e.g., Apple, Microsoft) or consumer goods (e.g., Procter & Gamble) using Yahoo Finance or Morningstar. "
-        "Consult a certified financial planner."
-    )
-}
+def load_cache():
+    try:
+        with open(CACHE_FILE, "rb") as f:
+            return pickle.load(f)
+    except:
+        return {
+            hashlib.md5("Hi, pretend you are a financial advisor. Now tell me how can I start investing in stock market?".lower().encode()).hexdigest(): (
+                "As a financial advisor, 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 (low fees), Vanguard (index funds like VTI), or Robinhood (commission-free).\n"
+                "4. **Investments**: Start with ETFs (e.g., VOO for S&P 500) or mutual funds.\n"
+                "5. **Strategy**: Use dollar-cost averaging with $100-$500 monthly.\n"
+                "6. **Risks**: Diversify and monitor.\n"
+                "Consult a certified financial planner."
+            ),
+            hashlib.md5("do you have a list of companies you recommend?".lower().encode()).hexdigest(): (
+                "I cannot recommend specific companies without current market data. Instead, consider ETFs like VOO (S&P 500) or QQQ (tech-focused) for broad exposure. "
+                "For stocks, research sectors like technology (e.g., Apple, Microsoft) or consumer goods (e.g., Procter & Gamble) using Yahoo Finance or Morningstar. "
+                "Consult a certified financial planner."
+            ),
+            hashlib.md5("can you provide me a list of companies you recommend?".lower().encode()).hexdigest(): (
+                "I cannot provide specific company recommendations without up-to-date market analysis. For safer investments, consider ETFs like VOO (S&P 500) or QQQ (tech-focused). "
+                "If interested in stocks, explore stable companies in technology (e.g., Apple, Microsoft) or healthcare (e.g., Johnson & Johnson) using Yahoo Finance. "
+                "Always consult a financial planner for tailored advice."
+            ),
+            hashlib.md5("You have a list of companies you recommend?".lower().encode()).hexdigest(): (
+                "I cannot recommend specific companies without current market data. Instead, consider ETFs like VOO (S&P 500) or QQQ (tech-focused) for broad exposure. "
+                "For stocks, research sectors like technology (e.g., Apple, Microsoft) or consumer goods (e.g., Procter & Gamble) using Yahoo Finance or Morningstar. "
+                "Consult a certified financial planner."
+            )
+        }
+
+def save_cache(cache):
+    with open(CACHE_FILE, "wb") as f:
+        pickle.dump(cache, f)
+
+# Initialize response cache
+response_cache = load_cache()
 
-# Load model with optimizations
-model_name = "distilgpt2"
+# Initialize Grok 3 API async client
 try:
-    tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=False)
-    tokenizer.pad_token = tokenizer.eos_token  # Ensure pad token is set
-    model = AutoModelForCausalLM.from_pretrained(
-        model_name,
-        device_map="auto",
-        torch_dtype=torch.float16,
-        low_cpu_mem_usage=True  # Optimize memory usage
-    ).to(device)
-    model.eval()  # Set model to evaluation mode for faster inference
+    client = AsyncOpenAI(
+        api_key=os.getenv("XAI_API_KEY"),  # Ensure API key is set in environment variables
+        base_url="https://api.x.ai/v1"
+    )
 except Exception as e:
-    print(f"Error loading distilgpt2: {e}")
+    print(f"Error initializing Grok 3 API client: {e}")
     exit()
 
-# Define chat function
-def chat_with_model(message, history=None):  # Ignore history
+# Cache API responses with increased size
+@lru_cache(maxsize=500)
+async def fetch_grok_response(prompt: str, model: str) -> str:
+    start_time = time.time()
     try:
-        if not isinstance(message, str):
-            return "Error: User input must be a string"
-        # Normalize message for cache lookup (case-insensitive, strip whitespace)
-        message = message.strip().lower()
-        for cached_message, response in response_cache.items():
-            if cached_message.lower() == message:
-                return response
-        # Simplified prompt
-        full_prompt = (
-            "Financial advisor: Answer directly about stock market investments. "
-            "No specific company picks without data; suggest ETFs or general advice. "
-            f"User: {message}\nAssistant:"
+        stream = await client.chat.completions.create(
+            model=model,
+            messages=[
+                {"role": "system", "content": (
+                    "You are a financial advisor. Provide concise stock market investment advice. "
+                    "Avoid specific company recommendations without data; suggest ETFs or strategies. "
+                    "Advise consulting a certified financial planner."
+                )},
+                {"role": "user", "content": prompt}
+            ],
+            temperature=0.2,
+            max_tokens=150,
+            top_p=1.0,
+            stream=True
         )
-        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=256).to(device)
-        with torch.no_grad():  # Disable gradient computation for faster inference
-            outputs = model.generate(
-                **inputs,
-                max_new_tokens=50,  # Increased slightly for better responses
-                do_sample=False,  # Greedy decoding for speed
-                num_beams=1,  # Disable beam search for faster generation
-                pad_token_id=tokenizer.eos_token_id
-            )
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response[len(full_prompt):].strip() if response.startswith(full_prompt) else response
+        response = ""
+        async for chunk in stream:
+            if chunk.choices[0].delta.content:
+                response += chunk.choices[0].delta.content
+        print(f"API call took {time.time() - start_time:.2f} seconds")
+        return response.strip()
     except Exception as e:
+        print(f"API call took {time.time() - start_time:.2f} seconds")
         return f"Error generating response: {str(e)}"
 
+# Define async chat function
+async def chat_with_model(message, history=None):  # Ignore history
+    try:
+        if not isinstance(message, str):
+            return "Error: User input must be a string"
+        
+        # Normalize message and check cache
+        message_normalized = message.strip().lower()
+        message_hash = hashlib.md5(message_normalized.encode()).hexdigest()
+        if message_hash in response_cache:
+            return response_cache[message_hash]
+        
+        # Use smaller model for faster responses
+        model = "grok-3-mini"
+        response = await fetch_grok_response(message, model)
+        
+        # Update and save cache for new responses
+        if not response.startswith("Error"):
+            response_cache[message_hash] = response
+            save_cache(response_cache)
+        
+        return response
+    
+    except Exception as e:
+        return f"Error in chat processing: {str(e)}"
+
+# Wrapper for Gradio compatibility
+def chat_with_model_sync(message, history=None):
+    return asyncio.run(chat_with_model(message, history))
+
 # Create Gradio interface
 interface = gr.ChatInterface(
-    fn=chat_with_model,
-    title="Financial Advisor Chatbot (DistilGPT-2)",
-    description="Ask for advice on starting to invest in the stock market! Powered by DistilGPT-2. Provides single, direct answers without conversation history.",
+    fn=chat_with_model_sync,
+    title="Financial Advisor Chatbot (Grok 3 Mini)",
+    description=(
+        "Ask for advice on starting to invest in the stock market! Powered by xAI's Grok 3 Mini API for faster responses. "
+        "Provides single, direct answers without conversation history. "
+        "Monitor API rate limits in the xAI Developer Console (console.x.ai)."
+    ),
     examples=[
         "Hi, pretend you are a financial advisor. Now tell me how can I start investing in stock market?",
         "You have a list of companies you recommend?"