Update app.py

app.py

@@ -1,7 +1,7 @@
 # ==============================================================================
 # NEBULIX SEARCH ENGINE v6 FINAL (PRODUCTION READY)
-# Target: 2 vCPU • 16 GB RAM • True Agentic RAG • AI Streaming • SQLite Analytics
-# Includes Hugging Face Dataset Integration & SmolLM2 (Million Params) integration
+# Target: Hugging Face Spaces | 2 vCPU • 16 GB RAM • True Agentic RAG • AI Streaming
+# Includes Hugging Face Dataset Integration & SmolLM2 (Transformers/PyTorch)
 # ==============================================================================
 import json
 import uuid
@@ -13,7 +13,7 @@ import time
 import pickle
 import sqlite3
 import unicodedata
-import multiprocessing
+import threading
 import numpy as np
 import psutil
 from datetime import datetime
@@ -25,22 +25,24 @@ from rank_bm25 import BM25Okapi
 from cachetools import TTLCache, LRUCache
 import hashlib
 
+# ==============================================================================
+# 25. CPU ALLOCATION (60% Search, 40% AI on 2 vCPUs)
+# ==============================================================================
+os.environ["OMP_NUM_THREADS"] = "1"
+os.environ["OPENBLAS_NUM_THREADS"] = "1"
+os.environ["MKL_NUM_THREADS"] = "1"
+
 # --- ML & Data Imports ---
 try:
     from datasets import load_dataset
-    from llama_cpp import Llama
     from sentence_transformers import SentenceTransformer
+    import torch
+    from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextIteratorStreamer
+    torch.set_num_threads(1) # CPU optimization for Hugging Face Spaces
     ML_AVAILABLE = True
 except ImportError:
     ML_AVAILABLE = False
-    print("⚠️ ML packages missing. Run: pip install llama-cpp-python sentence-transformers psutil numpy datasets")
-
-# ==============================================================================
-# 25. CPU ALLOCATION (60% Search, 40% AI on 2 vCPUs)
-# ==============================================================================
-os.environ["OMP_NUM_THREADS"] = "1"
-os.environ["OPENBLAS_NUM_THREADS"] = "1"
-os.environ["MKL_NUM_THREADS"] = "1"
+    print("⚠️ ML packages missing. Check requirements.txt")
 
 # ==============================================================================
 # 22. SQLITE ANALYTICS SETUP
@@ -75,7 +77,8 @@ AI_ANSWER_CACHE = LRUCache(maxsize=2000)
 
 doc_embeddings = None
 text_embedder = None
-llm = None
+tokenizer = None
+pipe = None
 AI_SEMAPHORE = asyncio.Semaphore(1)
 
 CLICK_TRACKER = {}
@@ -108,9 +111,6 @@ def update_dynamic_fields(doc):
     
     doc["special_tags"] = list(tags)
 
-# ==============================================================================
-# TOKENIZER & UTILS
-# ==============================================================================
 def normalize_text(text: str) -> str:
     if not text: return ""
     text = unicodedata.normalize('NFKD', str(text)).encode('ASCII', 'ignore').decode('utf-8')
@@ -134,24 +134,16 @@ def cosine_similarity(vec_a, mat_b):
 def detect_language(text):
     return "English"
 
-# ==============================================================================
-# DATASET LOADING & NEW SCHEMA MAPPING
-# ==============================================================================
 def build_searchable_text(doc):
-    """Combines title, description, tag, and string values from meta/og tags into one text."""
     parts = [doc.get("title", ""), doc.get("description", ""), doc.get("tag", "")]
-    
     for k, v in doc.get("meta_tags", {}).items():
         if isinstance(v, str): parts.append(v)
-            
     for k, v in doc.get("og_tags", {}).items():
         if isinstance(v, str): parts.append(v)
-        
     return " ".join(p for p in parts if p).strip()
 
 def build_runtime_record(doc):
     r = doc.copy()
-    
     r["meta_tags"] = doc.get("meta_tags", {})
     r["og_tags"] = doc.get("og_tags", {})
     r["images"] = doc.get("images", [])
@@ -179,7 +171,7 @@ def build_runtime_record(doc):
     return r
 
 def init_system():
-    global bm25_index, doc_embeddings, text_embedder, llm, raw_data, url_to_data
+    global bm25_index, doc_embeddings, text_embedder, tokenizer, pipe, raw_data, url_to_data
     init_db()
     print("🚀 [1/4] Fetching Dataset & Building Indexes...")
     
@@ -188,7 +180,7 @@ def init_system():
     if not os.path.exists(CACHE_FILE):
         hf_token = os.getenv("HF_TOKEN")
         if not hf_token:
-            raise ValueError("❌ HF_TOKEN environment variable is missing. Cannot fetch dataset.")
+            print("⚠️ HF_TOKEN environment variable is missing. It's needed if the dataset is private.")
         print("☁️ Downloading dataset from Hugging Face...")
         dataset = load_dataset("Nebulixlabs/Ultimate-Search-Engine-Dataset", split="train", token=hf_token)
         dataset.to_json(CACHE_FILE)
@@ -224,17 +216,16 @@ def init_system():
                 with open("embeddings.bin", "wb") as fb: pickle.dump(doc_embeddings, fb)
                 with open("metadata.json", "w") as fm: json.dump({"count": len(doc_embeddings)}, fm)
             
-            # =============== CHANGED TO SMOLLM2 (360 Million Params) ===============
-            print("🔥 [4/4] Loading Llama.cpp SmolLM2-360M GGUF (Lightweight CPU Model)...")
+            # =============== TRANSFORMERS BASED SMOLLM2 ===============
+            print("🔥 [4/4] Loading SmolLM2 (360 Million Params) via Transformers & PyTorch...")
             try:
-                llm = Llama(
-                    model_path="./models/smollm2-360m-instruct-q8_0.gguf", # Update filename if downloaded differently
-                    n_threads=1, 
-                    n_ctx=3000, 
-                    verbose=False
-                )
-            except:
-                print("⚠️ LLM not found locally in ./models/, mocked mode active.")
+                model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
+                tokenizer = AutoTokenizer.from_pretrained(model_id)
+                model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
+                pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
+                print("✅ Model successfully loaded into memory.")
+            except Exception as e:
+                print(f"⚠️ LLM loading failed: {e}. Mocked mode active.")
     print("✅ Nebulix v6 Ready.")
 
 # ==============================================================================
@@ -283,9 +274,6 @@ def get_ranked_results(query: str):
         })
     return formatted
 
-# ==============================================================================
-# CHUNKING
-# ==============================================================================
 def get_best_chunks(query, full_text):
     if not full_text: return []
     words = full_text.split()
@@ -308,9 +296,12 @@ def get_best_chunks(query, full_text):
 # TRUE AGENTIC RAG PLANNER (Max 10 Hops)
 # ==============================================================================
 async def agent_planner_call(prompt):
-    if not llm: return "ACTION: ANSWER"
-    res = await asyncio.to_thread(llm, prompt, max_tokens=20, temperature=0.1, stop=["\n"])
-    return res["choices"][0]["text"].strip()
+    if not pipe: return "ACTION: ANSWER"
+    res = await asyncio.to_thread(
+        pipe, prompt, max_new_tokens=20, temperature=0.1, do_sample=True, return_full_text=False
+    )
+    generated = res[0]["generated_text"].strip()
+    return generated.split('\n')[0]
 
 async def execute_agentic_rag(query, candidates, debug=False):
     context_memory = ""
@@ -322,7 +313,6 @@ async def execute_agentic_rag(query, candidates, debug=False):
     
     async with AI_SEMAPHORE:
         for hop in range(10):
-            # Prompt simplified slightly for million-parameter models to ensure strict output
             sys_prompt = f"""System: You are a Planner. Output ONLY the action command.
 Query: {query}
 URLs:
@@ -416,11 +406,13 @@ async def api_aimode(q: str = Query(...), debug: bool = False):
     results = get_ranked_results(q)
     final_prompt, sources, debug_log = await execute_agentic_rag(q, results, debug)
     
-    if not llm: return {"Answer": "Mocked", "Sources": sources}
+    if not pipe: return {"Answer": "Mocked", "Sources": sources}
     
     async with AI_SEMAPHORE:
-        res = await asyncio.to_thread(llm, final_prompt, max_tokens=1024, temperature=0.2)
-        answer = res["choices"][0]["text"].strip()
+        res = await asyncio.to_thread(
+            pipe, final_prompt, max_new_tokens=1024, temperature=0.2, do_sample=True, return_full_text=False
+        )
+        answer = res[0]["generated_text"].strip()
         
     resp = {
         "Answer": answer[:150] + "...",
@@ -438,13 +430,30 @@ async def api_aimodechat(q: str = Query(...)):
     final_prompt, sources, _ = await execute_agentic_rag(q, results)
     
     async def true_event_stream():
-        if llm:
+        if pipe and tokenizer:
             async with AI_SEMAPHORE:
-                for chunk in llm(final_prompt, max_tokens=1024, stream=True, temperature=0.2):
-                    text = chunk["choices"][0]["text"]
-                    if text: yield f"data: {text}\n\n"
+                streamer = TextIteratorStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)
+                
+                def generate():
+                    try:
+                        pipe(final_prompt, max_new_tokens=1024, temperature=0.2, do_sample=True, streamer=streamer, return_full_text=False)
+                    except Exception as e:
+                        print(f"Generation error: {e}")
+                
+                # Model ko background thread mein start karo taaki event loop block na ho
+                thread = threading.Thread(target=generate)
+                thread.start()
+                
+                while True:
+                    try:
+                        # asyncio.to_thread ka use iterator se text fetch karne ke liye 
+                        text = await asyncio.to_thread(next, streamer)
+                        if text: yield f"data: {text}\n\n"
+                    except StopIteration:
+                        break
         else:
             yield "data: [Mocked stream due to missing LLM model]\n\n"
+        
         yield f"data: [SOURCES_JSON] {json.dumps(sources)}\n\n"
         
     return StreamingResponse(true_event_stream(), media_type="text/event-stream")
@@ -485,7 +494,7 @@ async def api_suggestions(q: str = Query(...)):
 async def api_stats():
     mem = psutil.virtual_memory()
     return {
-        "engine": "Nebulix v6 HF Production (SmolLM2 360M)",
+        "engine": "Nebulix v6 HF Production (SmolLM2 360M via Transformers)",
         "documents": len(raw_data),
         "cpu_usage_percent": psutil.cpu_percent(),
         "ram_usage_percent": mem.percent,