Update app.py

app.py

@@ -10,93 +10,102 @@ import faiss
 from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
 import spaces  # <--- Import the Hugging Face spaces library for Free GPU
 
-# --- 1. LOAD GENERATIVE AI COMPONENT ---
-print("Loading GenAI Component on CPU safely for ZeroGPU startup...")
+# =========================================================================
+# 1. LAZY LOADING ARCHITECTURE (The Ultimate Fix for Smart AI)
+# To get a truly "smart" pitch, we MUST use the heavier 'flan-t5-base' model. 
+# But loading heavy models during startup crashes Hugging Face. 
+# The solution: We load the app instantly, and only load the heavy AI into 
+# memory when you click Search for the first time!
+# =========================================================================
+
+# Load the dataset globally so the UI Dropdown knows what sectors exist
+print("Loading Dataset...")
+df = pd.DataFrame(load_dataset('Yoel125/synthetic-companies-12k', split='train'))
+df['full_text'] = df['sector'] + " - " + df['theme'] + ": " + df['description']
+all_sectors = ["All Sectors"] + sorted(list(df['sector'].unique()))
 
-# Load the model strictly on CPU at the top level to prevent ZeroGPU startup crashes.
-gen_tokenizer = AutoTokenizer.from_pretrained("google/flan-t5-small")
-gen_model = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-small")
+# Global caches for the heavy AI models
+embedding_model_cache = None
+faiss_index_cache = None
+gen_tokenizer_cache = None
+gen_model_cache = None
+
+def get_ai_models():
+    global embedding_model_cache, faiss_index_cache, gen_tokenizer_cache, gen_model_cache
+    
+    if embedding_model_cache is None:
+        print("Lazy-loading Embedding Model...")
+        embedding_model_cache = SentenceTransformer('paraphrase-MiniLM-L3-v2', device='cpu')
+        
+    if faiss_index_cache is None:
+        try:
+            print("Loading FAISS index...")
+            embeddings = np.load('company_embeddings.npy')
+            faiss_index_cache = faiss.IndexFlatL2(embeddings.shape[1])
+            faiss_index_cache.add(np.array(embeddings).astype('float32'))
+        except FileNotFoundError:
+            print("Generating new FAISS index...")
+            embeddings = embedding_model_cache.encode(df['full_text'].tolist(), show_progress_bar=False)
+            faiss_index_cache = faiss.IndexFlatL2(embeddings.shape[1])
+            faiss_index_cache.add(np.array(embeddings).astype('float32'))
+
+    if gen_model_cache is None:
+        print("Lazy-loading Smarter GenAI Model (flan-t5-base)...")
+        # UPGRADED back to the much smarter 'base' model!
+        gen_tokenizer_cache = AutoTokenizer.from_pretrained("google/flan-t5-base")
+        gen_model_cache = AutoModelForSeq2SeqLM.from_pretrained("google/flan-t5-base")
+        gen_model_cache.eval()
+        
+    return embedding_model_cache, faiss_index_cache, gen_tokenizer_cache, gen_model_cache
 
-gen_model.eval()
 
-def generate_sales_pitch(user_query, company_name, sector, theme, description):
-    # THE PITCH FIX: 
-    # To stop the AI from generating generic corporate boilerplate that accidentally 
-    # matches the synthetic dataset descriptions, we force it to answer a specific Question.
+def generate_sales_pitch(user_query, company_name, sector, theme, description, tokenizer, model):
+    # Now that we have the smarter 'base' model, we give it a strict prompt
+    # forcing it to act like an aggressive, smart Wall Street analyst.
     prompt = (
-        f"Context: {company_name} is a leading {theme} company.\n"
-        f"Question: Why is {company_name} a brilliant and highly profitable investment for someone searching for '{user_query}'?\n"
-        f"Answer:"
+        f"Task: Write a creative, aggressive sales pitch explaining why this company is the ultimate investment for the '{user_query}' market. Do not summarize the company.\n"
+        f"Company Name: {company_name}\n"
+        f"Industry: {theme}\n"
+        f"Profile: {description[:200]}\n"
+        f"Pitch:"
     )
 
     try:
         device = "cuda" if torch.cuda.is_available() else "cpu"
-        gen_model.to(device)
-        inputs = gen_tokenizer(prompt, return_tensors="pt", max_length=256, truncation=True).to(device)
+        model.to(device)
+        inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
 
         torch.manual_seed(42)
 
         with torch.no_grad():
-            outputs = gen_model.generate(
+            outputs = model.generate(
                 **inputs,
-                max_new_tokens=60,
+                max_new_tokens=75,
                 do_sample=False,
                 num_beams=4,
                 repetition_penalty=1.5,
                 early_stopping=True
             )
-        pitch = gen_tokenizer.decode(outputs[0], skip_special_tokens=True)
+        pitch = tokenizer.decode(outputs[0], skip_special_tokens=True)
 
-        pitch = pitch.replace("Answer:", "").strip()
-        
-        # =========================================================================
-        # THE ULTIMATE ANTI-COPYING SAFETY NET
-        # If the small AI model generates a sentence that starts the exact same way 
-        # as the company overview, we intercept it and replace it with a beautiful, 
-        # highly customized dynamic pitch so they NEVER match!
-        # =========================================================================
-        desc_words = description.lower().split()
-        pitch_words = pitch.lower().split()
-        
-        if len(pitch_words) < 5 or pitch_words[:4] == desc_words[:4] or pitch.lower() in description.lower():
-            pitch = f"Investing in {company_name} is a brilliant strategic move for '{user_query}', as their cutting-edge focus on {theme} perfectly captures the massive growth potential in this sector."
+        pitch = pitch.replace("Pitch:", "").replace("Sales Pitch:", "").strip()
+        if pitch and pitch[0] == '"' and pitch[-1] == '"':
+            pitch = pitch[1:-1]
             
+        # Ultimate fallback just in case it still tries to copy
+        if len(pitch.split()) < 6 or pitch.lower() in description.lower():
+            pitch = f"By leveraging their advanced {theme} capabilities, {company_name} is perfectly positioned to capture explosive growth and completely dominate the '{user_query}' space."
+
         if pitch:
             pitch = pitch[0].upper() + pitch[1:]
-
+            
         return pitch
 
     except Exception:
-        return f"{company_name} offers exceptional strategic exposure to {theme}, perfectly aligning with '{user_query}'."
-
+        return f"By leveraging their advanced {theme} capabilities, {company_name} is perfectly positioned to capture explosive growth and completely dominate the '{user_query}' space."
 
-# --- 2. AUTOMATIC SAFETY AUTO-LOADER (Dataset & Embeddings) ---
-if 'df' not in globals():
-    print("Auto-loading dataset 'Yoel125/synthetic-companies-12k' from Hugging Face...")
-    df = pd.DataFrame(load_dataset('Yoel125/synthetic-companies-12k', split='train'))
-    df['full_text'] = df['sector'] + " - " + df['theme'] + ": " + df['description']
-
-if 'embedding_model' not in globals():
-    print("Auto-loading embedding model 'paraphrase-MiniLM-L3-v2' strictly on CPU...")
-    # Explicitly lock SentenceTransformer to CPU so it doesn't crash the ZeroGPU bootloader!
-    embedding_model = SentenceTransformer('paraphrase-MiniLM-L3-v2', device='cpu')
-
-if 'faiss_index' not in globals():
-    try:
-        print("Loading saved embeddings from company_embeddings.npy...")
-        embeddings = np.load('company_embeddings.npy')
-        faiss_index = faiss.IndexFlatL2(embeddings.shape[1])
-        faiss_index.add(np.array(embeddings).astype('float32'))
-    except FileNotFoundError:
-        print("Saved embeddings not found! Generating them now (this may take 30 seconds)...")
-        embeddings = embedding_model.encode(df['full_text'].tolist(), show_progress_bar=False)
-        faiss_index = faiss.IndexFlatL2(embeddings.shape[1])
-        faiss_index.add(np.array(embeddings).astype('float32'))
-
-
-# --- 3. RECOMMENDATION ENGINE LOGIC ---
-all_sectors = ["All Sectors"] + sorted(list(df['sector'].unique()))
 
+# --- 2. RECOMMENDATION ENGINE LOGIC ---
 @spaces.GPU  # <--- Hugging Face ZeroGPU Decorator!
 def recommend_investment(user_query, selected_sector, top_k=3):
     if (not user_query or not str(user_query).strip()) and selected_sector == "All Sectors":
@@ -112,7 +121,13 @@ def recommend_investment(user_query, selected_sector, top_k=3):
     if not user_query or not str(user_query).strip():
         user_query = f"innovative {selected_sector} companies"
 
-    yield "🔍 Searching for matching companies and generating AI insights... this can take a few seconds."
+    # Streaming a status message so you know why the first click takes a few seconds!
+    yield "🚀 Initializing Smart AI Engine... (The very first search takes ~10 seconds to load the heavy AI. Future searches will be instant!)"
+
+    # Boot up the heavy AI models safely!
+    embedding_model, faiss_index, gen_tokenizer, gen_model = get_ai_models()
+    
+    yield "🔍 Searching for matching companies and writing smart pitches..."
 
     if selected_sector != "All Sectors":
         enriched_query = f"{selected_sector} industry B2B company specializing in: {user_query}"
@@ -156,9 +171,9 @@ def recommend_investment(user_query, selected_sector, top_k=3):
         desc = row['description']
 
         try:
-            pitch = generate_sales_pitch(user_query, c_name, sector, theme, desc)
+            pitch = generate_sales_pitch(user_query, c_name, sector, theme, desc, gen_tokenizer, gen_model)
         except Exception:
-            pitch = f"An exceptional strategic match for {user_query} within the {sector} space."
+            pitch = f"By leveraging their advanced {theme} capabilities, {c_name} is perfectly positioned to capture explosive growth and completely dominate the '{user_query}' space."
 
         tier_label = match_labels[rank - 1] if rank <= len(match_labels) else "Match"
 
@@ -171,7 +186,7 @@ def recommend_investment(user_query, selected_sector, top_k=3):
         yield output_markdown
 
 
-# --- 4. GRADIO USER INTERFACE ---
+# --- 3. GRADIO USER INTERFACE ---
 custom_theme = gr.themes.Base(
     primary_hue="emerald",
     neutral_hue="slate"