Update app.py

app.py

@@ -16,11 +16,13 @@ from rank_bm25 import BM25Okapi
 import torch
 import aiohttp
 from bs4 import BeautifulSoup
-from transformers import pipeline
 
-# --- 1. SETUP & CPU ALLOCATION (MAX SPEED) ---
+# Hugging Face imports for Q4 Quantization
+from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, BitsAndBytesConfig
+
+# --- 1. SETUP & CPU ALLOCATION ---
 TOTAL_CORES = multiprocessing.cpu_count()
-AI_CORES = max(1, int(TOTAL_CORES * 0.7)) # Give 70% CPU to AI for max speed
+AI_CORES = max(1, int(TOTAL_CORES * 0.7)) 
 torch.set_num_threads(AI_CORES)
 print(f"🖥️ Total Cores: {TOTAL_CORES} | Dedicated AI Cores: {AI_CORES}")
 
@@ -113,29 +115,89 @@ def load_and_index_dataset():
             bm25_index = BM25Okapi(tokenized_corpus)
             print(f"🚀 Indexing Complete! {len(raw_data)} records ready.")
 
+# --- QWEN 2.5 (0.5B) + Q4 OPTIMIZATION ---
 def load_ai_models():
     global ai_generator
-    print("🤖 Loading LiquidAI/LFM2.5-350M for BLAZING FAST inference...")
+    print("🤖 Loading Qwen/Qwen2.5-0.5B-Instruct...")
     try:
-        # Optimization: use float32 on CPU but pipeline optimized
+        model_id = "Qwen/Qwen2.5-0.5B-Instruct"
+        tokenizer = AutoTokenizer.from_pretrained(model_id)
+
+        # Smart Hardware Check for 4-bit (Q4) Quantization
+        if torch.cuda.is_available():
+            print("⚡ GPU detected! Loading in 4-bit (Q4) using BitsAndBytes for Extreme Speed...")
+            quant_config = BitsAndBytesConfig(
+                load_in_4bit=True,
+                bnb_4bit_compute_dtype=torch.float16,
+                bnb_4bit_quant_type="nf4",
+            )
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                quantization_config=quant_config,
+                device_map="auto"
+            )
+        else:
+            print("🖥️ CPU detected! Loading optimized model (PyTorch execution)...")
+            # For CPU, 4-bit BitsAndBytes isn't natively optimal, so we load in standard float32
+            # Qwen 0.5B is already tiny (~1GB) so it runs perfectly fast on CPU too.
+            model = AutoModelForCausalLM.from_pretrained(
+                model_id,
+                device_map="cpu",
+                torch_dtype=torch.float32
+            )
+
         ai_generator = pipeline(
             "text-generation", 
-            model="LiquidAI/LFM2.5-350M", 
-            device=-1, 
-            torch_dtype=torch.float32,
-            trust_remote_code=True
+            model=model, 
+            tokenizer=tokenizer
         )
-        print("✅ LiquidAI Model Loaded! Superfast mode activated.")
+        print("✅ Qwen2.5-0.5B-Instruct (Superfast Mode) Ready!")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
+# --- AI TEXT GENERATION FIX ---
+def _generate_text(messages, max_tokens: int):
+    tokenizer = ai_generator.tokenizer
+    
+    # 1. Use Native Chat Template for Qwen (Makes answers highly accurate & stops hallucinations)
+    try:
+        prompt = tokenizer.apply_chat_template(
+            messages, 
+            tokenize=False, 
+            add_generation_prompt=True
+        )
+    except Exception:
+        prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant: "
+
+    # 2. Greedy Decoding (do_sample=False) + Return Setup
+    result = ai_generator(
+        prompt, 
+        max_new_tokens=max_tokens, 
+        do_sample=False, 
+        truncation=True,
+        return_full_text=False, # ❌ FIX: No more repeating the query in answers!
+        pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
+    )
+    
+    text_out = result[0]['generated_text'].strip()
+    return text_out
+
+async def run_ai_task(messages, max_tokens: int, query_key: str):
+    if not ai_generator: return "⚠️ AI Model is starting up."
+    if query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
+
+    async with ai_semaphore:
+        response = await asyncio.to_thread(_generate_text, messages, max_tokens)
+        if len(AI_RESPONSE_CACHE) > MAX_CACHE_SIZE: AI_RESPONSE_CACHE.clear()
+        AI_RESPONSE_CACHE[query_key] = response
+        return response
+
 def chunk_text(text, chunk_size=80):
     words = text.split()
     return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
 
 async def scrape_and_find_best_chunk(url, query):
     try:
-        # 🔥 REDUCED TIMEOUT TO 1.5 SECONDS FOR SPEED
         async with aiohttp.ClientSession() as session:
             async with session.get(url, timeout=1.5) as resp:
                 if resp.status == 200:
@@ -204,55 +266,13 @@ async def perform_core_search(raw_query: str):
     
     return [{"match_score": round(s, 2), **doc} for s, doc in scored_results[:100]], prediction, best_score
 
-def _generate_text(messages, max_tokens: int):
-    tokenizer = ai_generator.tokenizer
-    
-    # 1. Use the model's Native Chat Template (Prevents random autocomplete/echoing)
-    try:
-        prompt = tokenizer.apply_chat_template(
-            messages, 
-            tokenize=False, 
-            add_generation_prompt=True
-        )
-    except Exception:
-        # Fallback if model doesn't have a strict template
-        prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant: "
-
-    # 2. return_full_text=False ensures the prompt isn't returned in the output!
-    result = ai_generator(
-        prompt, 
-        max_new_tokens=max_tokens, 
-        do_sample=False, 
-        truncation=True,
-        return_full_text=False, # 🔥 THIS FIXES THE REPEATING QUERY ISSUE
-        pad_token_id=tokenizer.eos_token_id if tokenizer.eos_token_id is not None else 0
-    )
-    
-    text_out = result[0]['generated_text'].strip()
-    
-    # 3. Extra safety cleanup for small AI models
-    if text_out.lower().startswith("assistant:"):
-        text_out = text_out[10:].strip()
-        
-    return text_out
-
-async def run_ai_task(messages, max_tokens: int, query_key: str):
-    if not ai_generator: return "⚠️ AI Model is starting up."
-    if query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
-
-    async with ai_semaphore:
-        response = await asyncio.to_thread(_generate_text, messages, max_tokens)
-        if len(AI_RESPONSE_CACHE) > MAX_CACHE_SIZE: AI_RESPONSE_CACHE.clear()
-        AI_RESPONSE_CACHE[query_key] = response
-        return response
-
 @asynccontextmanager
 async def lifespan(app: FastAPI):
     asyncio.create_task(asyncio.to_thread(load_and_index_dataset))
     asyncio.create_task(asyncio.to_thread(load_ai_models))
     yield
 
-app = FastAPI(title="Nebulix Fast AI Search", lifespan=lifespan)
+app = FastAPI(title="Nebulix Fast Qwen AI Search", lifespan=lifespan)
 app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
 
 @app.get("/suggest")
@@ -275,23 +295,21 @@ async def ai_mode_short(q: str = Query(...)):
 
     context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:2]])
     messages = [
-        {"role": "system", "content": "You are a concise AI. Answer in 1 short sentence. Put important keywords in (parentheses)."},
-        {"role": "user", "content": f"Query: {q}\nContext: {context}"}
+        {"role": "system", "content": "You are a concise search AI. Provide a direct, 1-sentence answer based on the context. Highlight keywords in (parentheses)."},
+        {"role": "user", "content": f"Context: {context}\n\nQuestion: {q}"}
     ]
-    # Reduce token generation for lightning speed
     ai_reply = await run_ai_task(messages, 50, query_key=f"short_{q}")
     return {"request_id": req_id, "ai_response": ai_reply, "results": results}
 
 async def process_rag_chat(q: str):
     req_id = str(uuid.uuid4())
-    needs_rag = any(w in q.lower() for w in ["what", "who", "where", "how", "why", "tell me", "explain", "latest", "best"])
+    needs_rag = any(w in q.lower() for w in ["what", "who", "where", "how", "why", "tell me", "explain", "latest", "best", "news"])
     results, _, _ = await perform_core_search(q)
     
     rag_context = ""
     used_sources = []
 
     if needs_rag and results:
-        # 🔥 REDUCED FROM 10 TO TOP 3 URLs FOR FASTER RAG
         top_urls = results[:3]
         tasks = [scrape_and_find_best_chunk(doc["url"], q) for doc in top_urls]
         scraped_chunks = await asyncio.gather(*tasks)
@@ -309,13 +327,12 @@ async def process_rag_chat(q: str):
                 
         rag_context = "\n".join(context_parts)
 
-    system_prompt = "You are a Smart AI Search Assistant. Provide a short, direct answer. ENCLOSE IMPORTANT KEYWORDS IN (PARENTHESES). Mention Source URLs if used."
+    system_prompt = "You are a Smart AI Search Assistant. Provide a short, direct answer. ENCLOSE IMPORTANT KEYWORDS IN (PARENTHESES). If you use information from a Source, cite it."
     messages = [
         {"role": "system", "content": system_prompt},
-        {"role": "user", "content": f"Query: {q}\nContext:\n{rag_context}" if needs_rag else f"Query: {q}"}
+        {"role": "user", "content": f"Live Context Data:\n{rag_context}\n\nQuestion: {q}" if needs_rag else f"Question: {q}"}
     ]
 
-    # Max tokens 120 so it generates answers faster
     ai_reply = await run_ai_task(messages, 120, query_key=f"chat_{q}_{needs_rag}") 
     return {"request_id": req_id, "ai_response": ai_reply, "rag_used": needs_rag, "sources": used_sources, "results": results[:15]}