Update app.py

app.py

@@ -1,137 +1,114 @@
-import os
+from transformers import AutoModelForCausalLM, AutoTokenizer
+import torch
 import gradio as gr
-import hashlib
-import asyncio
-import pickle
-import time
-from openai import AsyncOpenAI
-from functools import lru_cache
 
-# Persistent cache configuration
-CACHE_FILE = "response_cache.pkl"
+# Define device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
-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()
+# Expanded response cache with more common questions
+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."
+    ),
+    "How do I start investing in stocks?": (
+        "To start investing in stocks, educate yourself using resources like Investopedia or books such as 'The Intelligent Investor.' "
+        "Set investment goals and assess your risk tolerance. Open an account with a brokerage like Fidelity or Vanguard, "
+        "and begin with diversified options like ETFs (e.g., VOO) or mutual funds. Consult a financial planner for personalized advice."
+    ),
+    "What are the best stocks to buy right now?": (
+        "I can’t recommend specific stocks without current market data. For a safer approach, consider ETFs like VOO (S&P 500) or QQQ (tech-focused). "
+        "If you prefer stocks, research companies in strong sectors like technology or healthcare using tools like Yahoo Finance. "
+        "Consult a financial planner for up-to-date advice."
+    ),
+    "What's the difference between stocks and bonds?": (
+        "Stocks represent ownership in a company and offer potential growth but with higher risk. Bonds are loans to a company or government, "
+        "providing steady interest with lower risk. Stocks may pay dividends, while bonds pay coupons. Diversifying with both can balance risk and return."
+    )
+}
 
-# Initialize Grok 3 API async client
+# Load model
+model_name = "facebook/opt-350m"
 try:
-    client = AsyncOpenAI(
-        api_key=os.getenv("XAI_API_KEY"),  # Ensure API key is set in environment variables
-        base_url="https://api.x.ai/v1"
-    )
+    tokenizer = AutoTokenizer.from_pretrained(model_name, clean_up_tokenization_spaces=False)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype=torch.float16
+    ).to(device)
 except Exception as e:
-    print(f"Error initializing Grok 3 API client: {e}")
+    print(f"Error loading OPT-350m: {e}")
     exit()
 
-# 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:
-        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
-        )
-        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
+# Define chat function with enhanced prompt
+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
-    
+        if message in response_cache:
+            return response_cache[message]
+        full_prompt = (
+            "You are a financial advisor with expertise in stock market investments. "
+            "Provide accurate, detailed, and actionable advice in a single response. "
+            "Do not rely on prior conversation context. "
+            "If you cannot provide specific recommendations (e.g., individual companies), "
+            "explain why and offer general guidance or alternative suggestions.\n\n"
+            "Examples:\n"
+            "Q: How can I start investing in the stock market?\n"
+            "A: Start by educating yourself with resources like Investopedia or 'The Intelligent Investor.' "
+            "Set clear goals and assess your risk tolerance. Choose a brokerage like Fidelity or Vanguard, "
+            "and consider starting with ETFs or mutual funds.\n\n"
+            "Q: What are some good companies to invest in?\n"
+            "A: I cannot recommend specific companies without current market data. "
+            "Consider ETFs like VOO or QQQ for broad exposure, or research stable sectors like technology or healthcare.\n\n"
+            "Now, answer the following question:\n"
+            "Q: {message}\n"
+            "A:"
+        ).format(message=message)
+        inputs = tokenizer(full_prompt, return_tensors="pt", truncation=True, max_length=512).to(device)
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=30,  # Adjustable for speed vs. quality
+            do_sample=True,
+            top_p=0.9,
+            temperature=0.4,
+            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
     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))
+        return f"Error generating response: {str(e)}"
 
 # Create Gradio interface
 interface = gr.ChatInterface(
-    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)."
-    ),
+    fn=chat_with_model,
+    title="Financial Advisor Chatbot (OPT-350m)",
+    description="Ask for advice on starting to invest in the stock market! Powered by Meta AI's OPT-350m. Provides single, direct answers without conversation history.",
     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?"
+        "You have a list of companies you recommend?",
+        "What's the difference between stocks and bonds?"
     ]
 )