Update app.py

app.py

@@ -7,7 +7,6 @@ import re
 import math
 import multiprocessing
 import urllib.parse
-import difflib
 import gc
 import sys
 import array
@@ -20,9 +19,13 @@ from fastapi import FastAPI, Query
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import snapshot_download
 
+import torch
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+
 # --- CONFIGURATION & MULTIPROCESSING ---
 TOTAL_CORES = multiprocessing.cpu_count()
-print(f"🖥️ Total Cores Assigned for High-Speed Lexical Search: {TOTAL_CORES}")
+torch.set_num_threads(TOTAL_CORES)
+print(f"🖥️ Total Cores Assigned for High-Speed Search: {TOTAL_CORES}")
 
 # Apka HF Token (Environment variable se ya direct yaha daal sakte hain)
 HF_TOKEN = os.getenv("HF_TOKEN", "YOUR_HF_TOKEN_HERE")
@@ -31,6 +34,10 @@ HF_TOKEN = os.getenv("HF_TOKEN", "YOUR_HF_TOKEN_HERE")
 QUERY_CACHE = {}
 MAX_CACHE_SIZE = 50000
 
+# --- LIGHTWEIGHT AI SPELL CHECKER (60M Params) ---
+spell_tokenizer = None
+spell_model = None
+
 # --- 1. DETERMINISTIC PREPROCESSING & CONSTANTS ---
 STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are", "by", "this", "that", "it", "what", "who", "where", "when", "why", "how"}
 ADULT_KEYWORDS = {
@@ -52,17 +59,16 @@ EXPANSION_DICT = {
     "js": ["javascript"]
 }
 
-# 🔥 Tuned Weights for Best Ranking WITHOUT AI Dense Models
 WEIGHTS = {
-    "W_BM25": 15.0,        # Increased priority
+    "W_BM25": 15.0, 
     "W_TITLE": 10.0, 
     "W_HEADINGS": 5.0, 
     "W_DESCRIPTION": 3.0, 
     "W_BODY": 1.5,
     "W_URL": 2.5, 
     "W_DOMAIN": 3.5, 
-    "W_PHRASE": 8.0,       # Exact phrase matching heavily rewarded
-    "W_EXACT": 10.0,       # Exact queries boosted
+    "W_PHRASE": 8.0, 
+    "W_EXACT": 10.0, 
     "W_FRESHNESS": 2.0,
     "W_QUALITY": 3.0, 
     "W_PROXIMITY": 6.0, 
@@ -107,8 +113,6 @@ class HybridSearchIndex:
         self.docs = []
         self.total_docs = 0
         self.total_len = 0
-        self.vocab = set()
-        self.high_freq_vocab = []
         self.df = defaultdict(int)
         self.avg_doc_len = 0.0
         
@@ -141,7 +145,6 @@ class HybridSearchIndex:
         parsed_url = urllib.parse.urlparse(url)
         domain_str = sys.intern(parsed_url.netloc.lower().replace("www.", ""))
         
-        # 🔥 Extracting Image URL
         image_url = None
         open_graph = metadata.get("open_graph") or {}
         if isinstance(open_graph, dict):
@@ -158,7 +161,7 @@ class HybridSearchIndex:
             "_fast_snippet": snippet_str[:300],
             "_fast_content": main_text,
             "_fast_headings": (h1_str + " " + h2_str)[:200],
-            "_fast_image": image_url # Hidden field for images output
+            "_fast_image": image_url 
         }
         
         title_tokens = tokenize(clean_doc["_fast_title"])
@@ -177,11 +180,9 @@ class HybridSearchIndex:
         
         self.docs.append(clean_doc)
         
-        distinct_tokens = set(all_tokens)
-        for term in distinct_tokens:
+        for term in set(all_tokens):
             self.inverted_index[term].append(doc_id)
             self.df[term] += 1
-            self.vocab.add(term)
             
         if not clean_doc["_is_adult"]:
             clean_title = [w for w in title_tokens if w not in ADULT_KEYWORDS and len(w) > 2]
@@ -192,7 +193,6 @@ class HybridSearchIndex:
         self.total_docs = len(self.docs)
         self.avg_doc_len = self.total_len / max(1, self.total_docs)
         self.suggestions = self.suggestion_counter.most_common(20000)
-        self.high_freq_vocab = [word for word, count in Counter(self.df).most_common(15000)]
         self.suggestion_counter.clear()
         gc.collect()
         print(f"🚀 Offline Search Engine Indexing Complete! Total Docs: {self.total_docs}")
@@ -205,13 +205,8 @@ class HybridSearchIndex:
     def _get_candidates(self, query_tokens: List[str]) -> Set[int]:
         candidates = set()
         for term in query_tokens:
-            if term in self.inverted_index: candidates.update(self.inverted_index[term])
-        
-        if len(candidates) < 10 and self.high_freq_vocab:
-            for term in query_tokens:
-                if term not in self.inverted_index:
-                    matches = difflib.get_close_matches(term, self.high_freq_vocab, n=2, cutoff=0.8)
-                    for m in matches: candidates.update(self.inverted_index[m])
+            if term in self.inverted_index: 
+                candidates.update(self.inverted_index[term])
         return candidates
 
     def calculate_features(self, doc: dict, query_tokens_sorted: List[str], query_str: str) -> Dict[str, float]:
@@ -304,42 +299,51 @@ class HybridSearchIndex:
         if self.total_docs == 0:
             return [], [], {"intent": "GENERAL", "signals": []}, None
 
-        query_norm = normalize_text(query)
-        raw_tokens = query_norm.split()
-        
-        corrected_tokens = []
-        is_corrected = False
+        # ==========================================
+        # 🤖 1. TINY AI SPELL CORRECTION (Smart Logic)
+        # ==========================================
+        corrected_query_str = None
+        search_query = query
         
-        if self.high_freq_vocab:
-            for t in raw_tokens:
-                if t in self.vocab or t in EXPANSION_DICT or t in ADULT_KEYWORDS:
-                    corrected_tokens.append(t)
-                else:
-                    matches = difflib.get_close_matches(t, self.high_freq_vocab, n=1, cutoff=0.85)
-                    if matches:
-                        corrected_tokens.append(matches[0])
-                        is_corrected = True
-                    else:
-                        corrected_tokens.append(t)
-        else:
-            corrected_tokens = raw_tokens
-            
-        corrected_query_str = " ".join(corrected_tokens) if is_corrected else None
-        search_tokens = [t for t in corrected_tokens if t not in STOP_WORDS]
+        if spell_model is not None and spell_tokenizer is not None:
+            try:
+                # "correct grammar: " is required by this specific lightweight model
+                input_text = "correct grammar: " + query
+                input_ids = spell_tokenizer(input_text, return_tensors="pt").input_ids
+                
+                # CPU-Optimized Fast Inference
+                with torch.inference_mode():
+                    outputs = spell_model.generate(input_ids, max_length=20, num_beams=2, early_stopping=True)
+                
+                ai_corrected = spell_tokenizer.decode(outputs[0], skip_special_tokens=True).strip()
+                
+                # Validation: Punctuation hata kar match karte hain (ki meaning change na ho)
+                q_clean = normalize_text(query)
+                ai_clean = normalize_text(ai_corrected)
+                
+                # Agar user ne galat type kiya hai tabhi change hoga, warna "ai for coding" same hi rahega
+                if q_clean != ai_clean and len(ai_clean) > 0:
+                    corrected_query_str = ai_corrected
+                    search_query = ai_corrected
+            except Exception as e:
+                pass # Fail safe mechanism
+
+        query_norm = normalize_text(search_query)
+        search_tokens = [t for t in query_norm.split() if t not in STOP_WORDS]
         
         expanded_tokens = list(search_tokens)
         for t in search_tokens:
             if t in EXPANSION_DICT: expanded_tokens.extend(EXPANSION_DICT[t])
         
         sorted_tokens = list(dict.fromkeys(expanded_tokens)) 
-        if not sorted_tokens: sorted_tokens = corrected_tokens 
+        if not sorted_tokens: sorted_tokens = query_norm.split()
             
         sorted_tokens.sort(key=lambda t: self.idf(t), reverse=True)
         intent_info = detect_query_intent(sorted_tokens)
         intent = intent_info["intent"]
         user_wants_adult = any(aw in query_norm for aw in ADULT_KEYWORDS)
         
-        # 1. Pure Fast Lexical Retrieval
+        # 2. Pure Fast Lexical Retrieval
         candidates_ids = self._get_candidates(sorted_tokens)
         
         scored_results = []
@@ -386,7 +390,7 @@ class HybridSearchIndex:
                     "domain": doc.get("_fast_domain", ""),
                     "match_score": round(final_score, 4),
                     "_internal_content": doc.get("_fast_content", ""),
-                    "_fast_image": doc.get("_fast_image", None) # Safe keeping for later
+                    "_fast_image": doc.get("_fast_image", None)
                 }
                 if debug:
                     doc_out["_debug_features"] = features
@@ -395,13 +399,13 @@ class HybridSearchIndex:
 
         scored_results.sort(key=lambda x: x[0], reverse=True)
         
-        # 🔥 Extract Top Images globally first (before applying strict domain filters)
+        # 🔥 Extract Top Images globally
         global_images = []
-        for score, doc_info, b_url, domain in scored_results[:20]: # Check top 20 best results for images
+        for score, doc_info, b_url, domain in scored_results[:20]:
             img_url = doc_info.get("_fast_image")
             if img_url and img_url not in global_images:
                 global_images.append(img_url)
-                if len(global_images) >= 10: break # Collect Max 10 relevant images
+                if len(global_images) >= 10: break 
         
         diversified_results = []
         for score, doc_info, b_url, domain in scored_results:
@@ -409,29 +413,36 @@ class HybridSearchIndex:
             seen_domains[domain] += 1
             seen_urls.add(b_url)
             
-            doc_info.pop("_fast_image", None) # Remove it from URL output 
+            doc_info.pop("_fast_image", None)
             diversified_results.append(doc_info)
-            if len(diversified_results) >= STRICT_RESULT_LIMIT: break # 🔥 Top 5 URLs limit
+            if len(diversified_results) >= STRICT_RESULT_LIMIT: break 
                 
         for idx, item in enumerate(diversified_results): item["rank"] = idx + 1
         return diversified_results, global_images, intent_info, corrected_query_str
 
 search_index = HybridSearchIndex()
 
-# --- DATA LOADER (OFFLINE + SYNC) ---
+# --- DATA LOADER (OFFLINE + SYNC + AI CHECKER) ---
+def load_tiny_ai():
+    global spell_tokenizer, spell_model
+    print("🧠 Loading Ultra-Lightweight T5 Spell Checker (~60M params) on CPU...")
+    try:
+        model_name = "AventIQ-AI/t5-small-grammar-correction"
+        spell_tokenizer = AutoTokenizer.from_pretrained(model_name)
+        spell_model = AutoModelForSeq2SeqLM.from_pretrained(model_name).eval() # Direct Eval mode for Speed
+        print("✅ Tiny AI Spell Checker Loaded Successfully!")
+    except Exception as e:
+        print(f"⚠️ Could not load AI Spell Checker: {e}")
+
 def sync_and_load_datasets():
     print("📥 Syncing Required Datasets securely...")
     try:
-        # Sync Main Lexical dataset to memory
         dataset_path = snapshot_download(
             repo_id="Nebulixlabs/web-search-dataset", 
             repo_type="dataset", 
             allow_patterns="data/*.jsonl", 
             token=HF_TOKEN
         )
-        
-        # Simply downloading the Private Vector repo as per your request for offline storage
-        # (Vectors won't be queried live since you strictly requested NO AI MODEL for search engine)
         snapshot_download(
             repo_id="Nebulixlabs/web-search-faiss-index", 
             repo_type="dataset", 
@@ -440,7 +451,7 @@ def sync_and_load_datasets():
         print("✅ Private Vector & Public Datasets Synced Offline!")
     except Exception as e:
         dataset_path = "."
-        print(f"⚠️ Offline Fallback mode (No internet or token issue): {e}")
+        print(f"⚠️ Offline Fallback mode: {e}")
         
     data_folder = os.path.join(dataset_path, "data")
     if not os.path.exists(data_folder): data_folder = "data" 
@@ -464,6 +475,8 @@ def sync_and_load_datasets():
 # --- FASTAPI APP ENDPOINTS ---
 @asynccontextmanager
 async def lifespan(app: FastAPI):
+    # Dono kaam ek sath chalenge bina time waste kiye
+    asyncio.create_task(asyncio.to_thread(load_tiny_ai))
     asyncio.create_task(asyncio.to_thread(sync_and_load_datasets))
     yield
 
@@ -499,7 +512,7 @@ async def pure_search(q: str = Query(...), content: bool = Query(False), debug:
         "request_id": req_id, 
         "query": q, 
         "corrected_query": corrected_query,
-        "images": global_images,      # 🖼️ Images output globally list ki tarah yaha set hai!
+        "images": global_images,      
         "results_count": len(top_5), 
         "results": top_5
     }