Update app.py

app.py

@@ -7,15 +7,14 @@ import re
 from datasets import load_dataset
 from sentence_transformers import SentenceTransformer
 import faiss
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM
 import spaces  # <--- Import the Hugging Face spaces library for Free GPU
 
 # =========================================================================
 # 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!
+# To get a truly "smart" pitch, we MUST use a modern State-of-the-Art LLM.
+# We are upgrading from the ancient 'flan-t5' to the brilliant 'Qwen2.5' model.
+# By lazy loading, we ensure Hugging Face never crashes during startup!
 # =========================================================================
 
 # Load the dataset globally so the UI Dropdown knows what sectors exist
@@ -50,58 +49,66 @@ def get_ai_models():
             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")
+        print("Lazy-loading State-of-the-Art GenAI Model (Qwen2.5-0.5B-Instruct)...")
+        # UPGRADED to a massively smarter, modern Causal LLM (ChatGPT equivalent for small models)
+        gen_tokenizer_cache = AutoTokenizer.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct")
+        gen_model_cache = AutoModelForCausalLM.from_pretrained("Qwen/Qwen2.5-0.5B-Instruct", torch_dtype="auto")
         gen_model_cache.eval()
         
     return embedding_model_cache, faiss_index_cache, gen_tokenizer_cache, gen_model_cache
 
 
 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"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:"
-    )
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    
+    # Modern ChatML format used by state-of-the-art models like Qwen and Llama
+    messages = [
+        {"role": "system", "content": "You are a brilliant, aggressive Wall Street investment analyst. Your job is to write a single, highly persuasive, creative sentence explaining why a company is a massive investment opportunity."},
+        {"role": "user", "content": f"Company: {company_name}\nIndustry: {theme}\nWhat they do: {description[:300]}\n\nWrite a 1-sentence sales pitch explaining why this company is the ultimate strategic investment for someone focused on '{user_query}'. Do not just summarize what they do. Be creative and aggressive."}
+    ]
+    
+    # Apply the exact chat template the model was trained on
+    prompt = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
 
     try:
-        device = "cuda" if torch.cuda.is_available() else "cpu"
         model.to(device)
-        inputs = tokenizer(prompt, return_tensors="pt", max_length=512, truncation=True).to(device)
+        inputs = tokenizer(prompt, return_tensors="pt").to(device)
 
+        # 100% STRICT DETERMINISM: Forces the exact same brilliant output every time!
         torch.manual_seed(42)
+        if torch.cuda.is_available():
+            torch.cuda.manual_seed_all(42)
+            torch.backends.cudnn.deterministic = True
+            torch.backends.cudnn.benchmark = False
 
         with torch.no_grad():
             outputs = model.generate(
                 **inputs,
                 max_new_tokens=75,
                 do_sample=False,
-                num_beams=4,
-                repetition_penalty=1.5,
-                early_stopping=True
+                repetition_penalty=1.1,
             )
-        pitch = tokenizer.decode(outputs[0], skip_special_tokens=True)
+            
+        # Causal LMs output the prompt + generation. We slice off the prompt.
+        input_length = inputs.input_ids.shape[1]
+        pitch = tokenizer.decode(outputs[0][input_length:], skip_special_tokens=True).strip()
 
-        pitch = pitch.replace("Pitch:", "").replace("Sales Pitch:", "").strip()
+        # Clean up quotes
         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."
+        # Just in case the AI gets chatty, force it to 1-2 sentences max
+        sentences = pitch.split(". ")
+        if len(sentences) > 2:
+            pitch = ". ".join(sentences[:2]) + "."
 
         if pitch:
             pitch = pitch[0].upper() + pitch[1:]
             
         return pitch
 
-    except Exception:
+    except Exception as e:
+        print(f"GenAI Generation Error: {e}")
         return f"By leveraging their advanced {theme} capabilities, {company_name} is perfectly positioned to capture explosive growth and completely dominate the '{user_query}' space."