Update app.py

app.py

@@ -6,79 +6,50 @@ import uvicorn
 import re
 import multiprocessing
 import urllib.parse
-import psutil
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, Query, Form
+from fastapi.responses import JSONResponse
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import hf_hub_download
-import numpy as np
-import bm25s
+from rank_bm25 import BM25Okapi
+import torch
 import aiohttp
 from bs4 import BeautifulSoup
+from transformers import pipeline
 
-# 🔥 COMPREHENSIVE HOTFIX FOR HUGGING FACE SPACES / OPTIMUM
-# The new Transformers library removed TensorFlow checks, but Optimum still looks for them.
-# We inject these functions into the root namespace BEFORE Optimum imports them.
-import transformers
-import transformers.utils
-import transformers.utils.import_utils
-
-def _is_false(): return False
-def _is_true(): return True
-
-# Patch the Root Namespace
-transformers.is_tf_available = _is_false
-transformers.is_torch_available = _is_true
-# Patch the Utils Namespace
-transformers.utils.is_tf_available = _is_false
-transformers.utils.is_torch_available = _is_true
-# Patch the Import Utils Namespace
-transformers.utils.import_utils.is_tf_available = _is_false
-transformers.utils.import_utils.is_torch_available = _is_true
-
-# NOW it is safe to import optimum
-from transformers import AutoTokenizer, pipeline
-from optimum.onnxruntime import ORTModelForCausalLM
-import onnxruntime as ort
-
-# --- 1. SETUP & STRICT CPU ALLOCATION ---
+# --- 1. SETUP & CPU ALLOCATION ---
 TOTAL_CORES = multiprocessing.cpu_count()
 AI_CORES = max(1, int(TOTAL_CORES * 0.4))
-
-# Try to strictly enforce CPU Affinity, fallback safely if Docker/Space restricts permissions
-try:
-    p = psutil.Process()
-    p.cpu_affinity(list(range(AI_CORES)))
-    print(f"🖥️ Total Cores: {TOTAL_CORES} | 🔒 Process strictly bound to {AI_CORES} cores.")
-except Exception as e:
-    print(f"🖥️ Total Cores: {TOTAL_CORES} | Allocating {AI_CORES} cores to AI. (Strict affinity skipped: {e})")
+torch.set_num_threads(AI_CORES)
+print(f"🖥️ Total Cores: {TOTAL_CORES} | Allocating {AI_CORES} cores strictly to AI.")
 
 raw_data = []
 suggestion_list = []
 bm25_index = None
-
-MAX_CONCURRENT_AI = 1
+MAX_CONCURRENT_AI = max(1, AI_CORES)
 ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 ai_generator = None
-ai_tokenizer = None
 
 # --- TRACKING & TRUST SYSTEM ---
-URL_STATS = {} 
+URL_STATS = {} # Format: { "https://example.com": {"user_visits": 0, "ai_visits": 0, "is_trusted": False} }
+
+# --- AI CACHING SYSTEM ---
 AI_RESPONSE_CACHE = {}
 MAX_CACHE_SIZE = 5000
-STOP_WORDS = ["the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are"]
+STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are"}
 HF_TOKEN = os.getenv("HF_TOKEN")
 
+# 🔥 MASSIVE DESI & MODERN ADULT KEYWORDS LIST
 ADULT_KEYWORDS = {
     "porn", "sex", "nude", "nsfw", "xxx", "adult", "erotic", "onlyfans", 
-    "mms", "chudai", "bhabhi", "devar", "xvideos", "pornhub", 
-    "brazzers", "mia khalifa", "bf video", "blue film", "kamasutra",
+    "mms", "chudai", "bhabhi", "devar", "savita bhabhi", "xvideos", "pornhub", 
+    "brazzers", "mia khalifa", "bf video", "blue film", "kamasutra", "nude",
     "boobs", "tits", "ass", "booty", "naked", "desi porn", "hot desi", "xnxx",
     "kacha badam", "ullu", "kooku", "mastram", "gandi baat", "nangi", "nanga",
     "sex tape", "milf", "stepmom", "incest", "playboy", "scandal", "leaked"
 }
 
-# --- 2. ADVANCED NLP & ULTRA-FAST INDEXING ---
+# --- 2. ADVANCED NLP & INDEXING ---
 def tokenize(text):
     if not text: return []
     tokens = re.findall(r'[a-z0-9]+', str(text).lower())
@@ -93,74 +64,65 @@ def load_and_index_dataset():
         try:
             dataset_path = hf_hub_download(
                 repo_id="Nebulixlabs/Ultimate-Search-Engine-Dataset",
-                filename="dataset.jsonl", repo_type="dataset", token=HF_TOKEN 
+                filename="dataset.jsonl",
+                repo_type="dataset",
+                token=HF_TOKEN 
             )
         except Exception as e:
             print(f"⚠️ Hugging Face Download failed: {e}")
 
-    boosted_corpus_texts = []
-    MAX_RECORDS_TO_LOAD = 150000 
+    tokenized_corpus = []
+    MAX_RECORDS_TO_LOAD = 150000 # Prevent OOM Crash
 
     if os.path.exists(dataset_path):
         with open(dataset_path, "r", encoding="utf-8") as f:
             for i, line in enumerate(f):
-                if i >= MAX_RECORDS_TO_LOAD: break 
+                if i >= MAX_RECORDS_TO_LOAD:
+                    break 
+                
                 if line.strip():
                     try:
                         doc = json.loads(line)
                         url = doc.get("url", "")
-                        title_str = doc.get("title", "")
-                        desc_str = doc.get("description", "")
                         
                         if url not in URL_STATS:
                             URL_STATS[url] = {"user_visits": 0, "ai_visits": 0, "is_trusted": False}
                             
                         raw_data.append(doc)
+                        
+                        title_str = doc.get("title", "")
                         if title_str and len(title_str) < 60:
                             suggestion_list.append(title_str.lower())
                         
-                        # TF weighting before feeding to bm25s
-                        boosted_text = f"{title_str} " * 5 + f"{desc_str} " * 3 + url
-                        boosted_corpus_texts.append(boosted_text)
-                    except: pass
+                        title_tokens = tokenize(title_str)
+                        desc_tokens = tokenize(doc.get("description", ""))
+                        url_tokens = tokenize(url)
+                        
+                        boosted_tokens = (title_tokens * 5) + (desc_tokens * 3) + url_tokens
+                        tokenized_corpus.append(boosted_tokens)
+                    except: 
+                        pass
                     
-        if boosted_corpus_texts:
-            print("⚡ Building Ultra-Fast BM25s Index (Rust/C Optimized)...")
+        if tokenized_corpus:
+            print("⚡ Building High-Speed BM25 Okapi Index...")
             suggestion_list = sorted(list(set(suggestion_list)))
-            corpus_tokens = bm25s.tokenize(boosted_corpus_texts, stopwords=STOP_WORDS)
-            bm25_index = bm25s.BM25()
-            bm25_index.index(corpus_tokens)
-            print(f"🚀 Indexing Complete! {len(raw_data)} records ready.")
+            bm25_index = BM25Okapi(tokenized_corpus)
+            print(f"🚀 Indexing Complete! {len(raw_data)} records ready for search.")
     else:
         print("❌ Dataset file not found!")
 
 def load_ai_models():
-    global ai_generator, ai_tokenizer
-    print("🤖 Downloading & Optimizing Ultra-Small 135M Model for ONNX...")
+    global ai_generator
+    print("🤖 Loading Small & Fast 360M Text AI Model...")
     try:
-        model_id = "HuggingFaceTB/SmolLM2-135M-Instruct"
-        
-        # Strictly configure ONNX to use exactly the allowed CPU cores
-        sess_options = ort.SessionOptions()
-        sess_options.intra_op_num_threads = AI_CORES
-        sess_options.inter_op_num_threads = 1
-
-        ai_tokenizer = AutoTokenizer.from_pretrained(model_id)
-        
-        # export=True auto-converts the model to ONNX for massive CPU speedup
-        model = ORTModelForCausalLM.from_pretrained(
-            model_id, 
-            export=True, 
-            session_options=sess_options
-        )
-        
+        # 🔥 Using a small 360M parameter model for strict fast performance
         ai_generator = pipeline(
             "text-generation", 
-            model=model, 
-            tokenizer=ai_tokenizer, 
-            device=-1 # -1 Forces strictly CPU
+            model="HuggingFaceTB/SmolLM2-360M-Instruct", 
+            device=-1, 
+            torch_dtype=torch.float32
         )
-        print("✅ ONNX Accelerated 135M AI Engine Loaded! System ready.")
+        print("✅ Small AI Model Loaded! System ready.")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
@@ -176,13 +138,16 @@ async def scrape_and_find_best_chunk(url, query):
                 if resp.status == 200:
                     html = await resp.text()
                     soup = BeautifulSoup(html, "html.parser")
-                    text = " ".join([tag.get_text(strip=True) for tag in soup.find_all(['p', 'h1', 'h2', 'h3'])])
+                    tags_to_extract = ['p', 'h1', 'h2', 'h3']
+                    text = " ".join([tag.get_text(strip=True) for tag in soup.find_all(tags_to_extract)])
+                    
                     if not text: return ""
                     
                     chunks = chunk_text(text)
                     query_tokens = set(tokenize(query))
                     
-                    best_chunk, max_score = "", 0
+                    best_chunk = ""
+                    max_score = 0
                     early_stop_threshold = max(1, len(query_tokens) // 2) 
 
                     for chunk in chunks:
@@ -198,7 +163,7 @@ async def scrape_and_find_best_chunk(url, query):
         return ""
     return ""
 
-# --- 4. CORE SEARCH ALGORITHM (Fast Pre-filtering) ---
+# --- 4. CORE SEARCH ALGORITHM & RANKING ---
 async def perform_core_search(raw_query: str):
     if not raw_data or not bm25_index:
         return [], "Search index not ready.", 0
@@ -207,60 +172,68 @@ async def perform_core_search(raw_query: str):
     query_tokens = tokenize(query)
     user_wants_adult = any(adult_word in query for adult_word in ADULT_KEYWORDS)
     
-    if query_tokens:
-        bm25_query_tokens = bm25s.tokenize([query], stopwords=STOP_WORDS)
-        scores = bm25_index.get_scores(bm25_query_tokens)[0]
-    else:
-        scores = np.ones(len(raw_data))
-
+    scores = bm25_index.get_scores(query_tokens) if query_tokens else [1.0] * len(raw_data)
     scored_results = []
     seen_urls = set() 
-    
-    # Fast Pre-filtering
-    top_n = min(1000, len(scores))
-    top_indices = np.argsort(scores)[::-1][:top_n]
-
-    for index in top_indices:
-        score = float(scores[index])
-        if score < 0.1 and query_tokens: continue
 
-        doc = raw_data[index]
-        doc_url = doc.get("url", "")
-        
-        if doc_url in seen_urls: continue
-        
-        parsed_url = urllib.parse.urlparse(doc_url)
-        domain = parsed_url.netloc.lower().replace("www.", "")
-        path = parsed_url.path.strip("/")
-        
-        doc_title = doc.get("title", "").lower()
-        doc_desc = doc.get("description", "").lower()
-        doc_text = doc_title + " " + doc_desc + " " + doc_url
-        
-        is_adult_doc = doc.get("tag", "general").lower() in ["adult", "nsfw"] or any(aw in doc_text for aw in ADULT_KEYWORDS)
-        if not user_wants_adult and is_adult_doc: continue 
+    for index, score in enumerate(scores):
+        if score > 0.1 or not query_tokens:
+            doc = raw_data[index]
+            doc_url = doc.get("url", "")
             
-        final_score = score
-        query_clean = query.replace(" ", "")
-        domain_clean = domain.replace(" ", "")
-        
-        if query_clean == domain_clean or domain_clean.startswith(query_clean + "."):
-            final_score *= 50.0 if not path else 10.0
-        elif query_clean in domain_clean:
-            final_score *= 25.0 if not path else 5.0
+            if doc_url in seen_urls: continue
+            
+            parsed_url = urllib.parse.urlparse(doc_url)
+            domain = parsed_url.netloc.lower().replace("www.", "")
+            path = parsed_url.path.strip("/")
+            
+            doc_title = doc.get("title", "").lower()
+            doc_desc = doc.get("description", "").lower()
+            doc_text = doc_title + " " + doc_desc + " " + doc_url
+            
+            is_adult_doc = doc.get("tag", "general").lower() in ["adult", "nsfw"] or any(aw in doc_text for aw in ADULT_KEYWORDS)
+            
+            if not user_wants_adult and is_adult_doc:
+                continue 
                 
-        if query in doc_title: final_score *= 5.0  
-        elif all(word in doc_title for word in query_tokens): final_score *= 3.0  
-        if query in doc_desc: final_score *= 2.0  
+            final_score = score
+            query_clean = query.replace(" ", "")
+            domain_clean = domain.replace(" ", "")
             
-        stats = URL_STATS.get(doc_url, {"user_visits": 0, "ai_visits": 0, "is_trusted": False})
-        if stats.get("is_trusted", False): final_score *= 2.0  
-        final_score += (stats.get("user_visits", 0) * 0.05) + (stats.get("ai_visits", 0) * 0.1)
-        
-        seen_urls.add(doc_url)
-        doc_copy = doc.copy()
-        doc_copy["trust_stats"] = stats
-        scored_results.append((final_score, doc_copy))
+            # 🔥 1. SUPER RANKING FOR DOMAINS
+            if query_clean == domain_clean or domain_clean.startswith(query_clean + "."):
+                if not path:
+                    final_score *= 50.0  # EXACT MATCH ROOT DOMAIN -> 50x Boost
+                else:
+                    final_score *= 10.0  # EXACT MATCH BUT SUBPAGE -> 10x Boost
+            elif query_clean in domain_clean:
+                if not path:
+                    final_score *= 25.0  # PARTIAL ROOT DOMAIN -> 25x Boost
+                else:
+                    final_score *= 5.0   # PARTIAL MATCH SUBPAGE -> 5x Boost
+                    
+            # 🔥 2. KEYWORD & TITLE MATCH BOOST
+            if query in doc_title:
+                final_score *= 5.0  # EXACT query present in Title -> 5x
+            elif all(word in doc_title for word in query_tokens):
+                final_score *= 3.0  # All query words present in Title -> 3x
+                
+            if query in doc_desc:
+                final_score *= 2.0  # EXACT query in Description -> 2x
+                
+            # 🔥 3. CONTEXTUAL TRUST BOOST
+            stats = URL_STATS.get(doc_url, {"user_visits": 0, "ai_visits": 0, "is_trusted": False})
+            if stats.get("is_trusted", False):
+                final_score *= 2.0  
+                
+            final_score += (stats.get("user_visits", 0) * 0.05) + (stats.get("ai_visits", 0) * 0.1)
+            
+            seen_urls.add(doc_url)
+            
+            doc_copy = doc.copy()
+            doc_copy["trust_stats"] = stats
+            
+            scored_results.append((final_score, doc_copy))
 
     scored_results.sort(key=lambda x: x[0], reverse=True)
     best_score = scored_results[0][0] if scored_results else 0
@@ -269,29 +242,19 @@ async def perform_core_search(raw_query: str):
     final_results = [{"match_score": round(s, 2), **doc} for s, doc in scored_results[:100]]
     return final_results, prediction, best_score
 
-# --- AI GENERATION HELPER (Non-Blocking) ---
-def _generate_text_sync(messages, max_tokens: int):
-    # Formats messages accurately using the model's chat template
-    prompt = ai_tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)
-    
-    # Generate output
-    result = ai_generator(
-        prompt, 
-        max_new_tokens=max_tokens, 
-        do_sample=True, 
-        temperature=0.3, 
-        top_p=0.85,
-        return_full_text=False # Ensures we only get the generated answer, not the prompt
-    )
-    return result[0]['generated_text'].strip()
+# --- AI GENERATION HELPER ---
+def _generate_text(messages, max_tokens: int):
+    prompt = "\n".join([f"{m['role'].capitalize()}: {m['content']}" for m in messages]) + "\nAssistant:"
+    result = ai_generator(prompt, max_new_tokens=max_tokens, do_sample=True, temperature=0.3, top_p=0.85)
+    text_out = result[0]['generated_text'].split("Assistant:")[-1].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:
-        # Pushes generation to a thread so FastAPI doesn't freeze
-        response = await asyncio.to_thread(_generate_text_sync, messages, max_tokens)
+        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
@@ -304,7 +267,10 @@ async def lifespan(app: FastAPI):
     yield
 
 app = FastAPI(title="Nebulix Advanced Text AI & RAG Engine", lifespan=lifespan)
-app.add_middleware(CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"])
+
+app.add_middleware(
+    CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"]
+)
 
 @app.get("/suggest")
 async def get_suggestions(q: str = Query(..., min_length=1)):
@@ -314,13 +280,15 @@ async def get_suggestions(q: str = Query(..., min_length=1)):
 
 @app.post("/track-click")
 async def track_url_click(url: str = Query(...)):
-    if url not in URL_STATS: URL_STATS[url] = {"user_visits": 0, "ai_visits": 0, "is_trusted": False}
+    if url not in URL_STATS:
+        URL_STATS[url] = {"user_visits": 0, "ai_visits": 0, "is_trusted": False}
     URL_STATS[url]["user_visits"] += 1
     return {"status": "success", "url": url, "visits": URL_STATS[url]["user_visits"]}
 
 @app.get("/url-stats")
 async def get_url_stats(url: str = Query(...)):
-    return {"url": url, "stats": URL_STATS.get(url, {"user_visits": 0, "ai_visits": 0, "is_trusted": False})}
+    stats = URL_STATS.get(url, {"user_visits": 0, "ai_visits": 0, "is_trusted": False})
+    return {"url": url, "stats": stats}
 
 @app.get("/search")
 async def pure_search(q: str = Query(...)):
@@ -337,14 +305,20 @@ async def ai_mode_short(q: str = Query(...)):
 
     context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:3]])
     
+    # 🔥 AI MODE SYSTEM PROMPT (STRICT: Max 80 Words & Use Brackets for Important Terms)
     messages = [
-        {"role": "system", "content": "You are a concise AI. Answer in EXACTLY 1 to 2 sentences (Max 80 words). YOU MUST ENCLOSE ANY IMPORTANT KEYWORD, NAME, OR FACTUAL DETAIL IN PARENTHESES, like (this)."},
+        {
+            "role": "system", 
+            "content": "You are a concise AI. Answer in EXACTLY 1 to 2 sentences (Max 80 words). YOU MUST ENCLOSE ANY IMPORTANT KEYWORD, NAME, OR FACTUAL DETAIL IN PARENTHESES, like (this)."
+        },
         {"role": "user", "content": f"Query: {q}\nContext: {context}"}
     ]
 
+    # Max tokens set to 80 to keep it very short (aprox 20 generated tokens logic based on your 4 word multiplier).
     ai_reply = await run_ai_task(messages, 80, 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"])
@@ -355,6 +329,7 @@ async def process_rag_chat(q: str):
 
     if needs_rag and results:
         top_urls = results[:10]
+        
         tasks = [scrape_and_find_best_chunk(doc["url"], q) for doc in top_urls]
         scraped_chunks = await asyncio.gather(*tasks)
         
@@ -365,14 +340,21 @@ async def process_rag_chat(q: str):
             
             if url not in URL_STATS: URL_STATS[url] = {"user_visits": 0, "ai_visits": 0, "is_trusted": False}
             URL_STATS[url]["ai_visits"] += 1
-            if URL_STATS[url]["ai_visits"] >= 100: URL_STATS[url]["is_trusted"] = True
+            if URL_STATS[url]["ai_visits"] >= 100:
+                URL_STATS[url]["is_trusted"] = True
+                
             doc["trust_stats"] = URL_STATS[url]
 
-            context_parts.append(f"Source [{i+1}] ({url}): {chunk if chunk else doc.get('description', '')}")
-            used_sources.append(doc)
+            if chunk:
+                context_parts.append(f"Source [{i+1}] ({url}): {chunk}")
+                used_sources.append(doc)
+            else:
+                context_parts.append(f"Source [{i+1}] ({url}): {doc.get('description', '')}")
+                used_sources.append(doc)
                 
         rag_context = "\n".join(context_parts)
 
+    # 🔥 AI CHAT SYSTEM PROMPT (Use Brackets for Important Terms)
     system_prompt = "You are a Smart AI Search Assistant. Provide a detailed answer. YOU MUST ENCLOSE EVERY IMPORTANT KEYWORD, ENTITY, OR FACTUAL DETAIL INSIDE PARENTHESES, like (keyword). Always mention the Source URL if you use its information."
     if not needs_rag:
         system_prompt = "You are a helpful AI Assistant. Answer the user's query conversationally, BUT ENCLOSE IMPORTANT WORDS IN PARENTHESES like (word)."
@@ -382,6 +364,7 @@ async def process_rag_chat(q: str):
         {"role": "user", "content": f"Query: {q}\n\nLive Context Data:\n{rag_context}" if needs_rag else f"Query: {q}"}
     ]
 
+    # Max tokens set to 200 for detailed chat mode.
     ai_reply = await run_ai_task(messages, 200, query_key=f"chat_{q}_{needs_rag}") 
 
     return {
@@ -400,6 +383,5 @@ async def ai_mode_chat_get(q: str = Query(...)):
 async def ai_mode_chat_post(q: str = Form(...)):
     return await process_rag_chat(q)
 
-# Using `app` variable directly prevents Spaces routing errors if the file is named differently
 if __name__ == "__main__":
-    uvicorn.run(app, host="0.0.0.0", port=7860, workers=1)
+    uvicorn.run(app, host="0.0.0.0", port=7860)