Update app.py

app.py

@@ -40,8 +40,6 @@ try:
     from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, TextIteratorStreamer
     
     # 🔥 STRICT NO-TRAINING LOCK 🔥
-    # Yeh ensure karta hai ki PyTorch gradients calculate na kare. 
-    # AI 100% pre-trained inference mode me chalega. Cost aur RAM dono bachenge.
     torch.set_num_threads(1)
     torch.set_grad_enabled(False) 
     
@@ -140,6 +138,17 @@ def cosine_similarity(vec_a, mat_b):
 def detect_language(text):
     return "English"
 
+def parse_json_field(val):
+    """Safely converts stringified JSONs to dicts if dataset schema changed."""
+    if isinstance(val, dict):
+        return val
+    if isinstance(val, str):
+        try:
+            return json.loads(val)
+        except:
+            return {}
+    return {}
+
 def build_searchable_text(doc):
     parts = [doc.get("title", ""), doc.get("description", ""), doc.get("tag", "")]
     for k, v in doc.get("meta_tags", {}).items():
@@ -150,13 +159,16 @@ def build_searchable_text(doc):
 
 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", [])
-    r["videos"] = doc.get("videos", [])
-    r["documents"] = doc.get("documents", [])
+    
+    # Safe JSON parsing to prevent crashes from schema changes
+    r["meta_tags"] = parse_json_field(doc.get("meta_tags"))
+    r["og_tags"] = parse_json_field(doc.get("og_tags"))
+    r["specific_data"] = parse_json_field(doc.get("specific_data"))
+    
+    r["images"] = doc.get("images") or []
+    r["videos"] = doc.get("videos") or []
+    r["documents"] = doc.get("documents") or []
     r["tag"] = doc.get("tag", "general")
-    r["specific_data"] = doc.get("specific_data", {})
     r["searchable_text"] = build_searchable_text(r)
     
     r["human_opens"] = r.get("human_opens", 0)
@@ -185,11 +197,24 @@ def init_system():
     
     if not os.path.exists(CACHE_FILE):
         hf_token = os.getenv("HF_TOKEN")
-        if not hf_token:
-            print("⚠️ HF_TOKEN environment variable is missing. Required if 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)
+        
+        try:
+            # Attempt 1: Standard load
+            dataset = load_dataset("Nebulixlabs/Ultimate-Search-Engine-Dataset", split="train", token=hf_token)
+            dataset.to_json(CACHE_FILE)
+        except Exception as e:
+            print(f"⚠️ Standard load failed due to schema mismatch ({e}). Trying fallback raw JSON loader...")
+            try:
+                # Attempt 2: Force raw JSON inference (bypasses broken schemas)
+                dataset = load_dataset("json", data_files="hf://datasets/Nebulixlabs/Ultimate-Search-Engine-Dataset/**/*.json*", split="train", token=hf_token)
+                dataset.to_json(CACHE_FILE)
+            except Exception as e2:
+                print(f"⚠️ Raw JSON load failed. Trying Parquet loader... ({e2})")
+                # Attempt 3: Parquet fallback
+                dataset = load_dataset("parquet", data_files="hf://datasets/Nebulixlabs/Ultimate-Search-Engine-Dataset/**/*.parquet", split="train", token=hf_token)
+                dataset.to_json(CACHE_FILE)
+
         print("✅ Dataset successfully cached locally.")
 
     if os.path.exists(CACHE_FILE):
@@ -198,14 +223,18 @@ def init_system():
         with open(CACHE_FILE, "r", encoding="utf-8") as f:
             for line in f:
                 if line.strip():
-                    doc = build_runtime_record(json.loads(line))
-                    raw_data.append(doc)
-                    url_to_data[doc["url"]] = doc
-                    
-                    tokens = tokenize(doc["searchable_text"])
-                    corpus_tokens.append(tokens)
-                    texts_to_embed.append(doc["searchable_text"])
-                    for t in tokens: add_to_trie(t)
+                    try:
+                        doc = build_runtime_record(json.loads(line))
+                        raw_data.append(doc)
+                        url_to_data[doc["url"]] = doc
+                        
+                        tokens = tokenize(doc["searchable_text"])
+                        corpus_tokens.append(tokens)
+                        texts_to_embed.append(doc["searchable_text"])
+                        for t in tokens: add_to_trie(t)
+                    except Exception as parse_e:
+                        # Ignore highly corrupted single lines
+                        pass
 
         bm25_index = BM25Okapi(corpus_tokens)
         
@@ -217,20 +246,19 @@ def init_system():
                 print("💾 [3/4] Loading Precomputed Embeddings...")
                 with open("embeddings.bin", "rb") as fb: doc_embeddings = pickle.load(fb)
             else:
-                # Yeh sirf text ko numbers me convert kar raha hai search ke liye (Training nahi hai)
                 print("���️ [3/4] Converting Text to Vector Embeddings (One-Time Inference)...")
-                doc_embeddings = text_embedder.encode(texts_to_embed, convert_to_numpy=True)
-                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)
+                if texts_to_embed:
+                    doc_embeddings = text_embedder.encode(texts_to_embed, convert_to_numpy=True)
+                    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)
             
             # =============== PRE-TRAINED TRANSFORMERS BASED SMOLLM2 ===============
             print("🔥 [4/4] Loading Pre-trained SmolLM2 (360M Params)...")
             try:
                 model_id = "HuggingFaceTB/SmolLM2-360M-Instruct"
                 tokenizer = AutoTokenizer.from_pretrained(model_id)
-                # Model is loaded strictly for inference
                 model = AutoModelForCausalLM.from_pretrained(model_id, torch_dtype=torch.float32)
-                model.eval() # Puts the model in evaluation (inference) mode
+                model.eval() # Puts the model in strict inference mode
                 pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=-1)
                 print("✅ Model successfully loaded in INFERENCE mode.")
             except Exception as e:
@@ -247,7 +275,7 @@ def get_ranked_results(query: str):
     bm25_scores = bm25_index.get_scores(q_tokens) if bm25_index else np.zeros(len(raw_data))
     semantic_scores = np.zeros(len(raw_data))
     
-    if text_embedder and doc_embeddings is not None:
+    if text_embedder and doc_embeddings is not None and len(doc_embeddings) > 0:
         q_vec = text_embedder.encode([query])[0]
         semantic_scores = cosine_similarity(q_vec, doc_embeddings)