Update app.py

app.py

@@ -6,6 +6,7 @@ import uvicorn
 import re
 import multiprocessing
 import urllib.parse
+from collections import Counter
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, Query, Form
 from fastapi.responses import JSONResponse
@@ -17,47 +18,42 @@ import aiohttp
 from bs4 import BeautifulSoup
 from transformers import pipeline
 
-# --- 1. SETUP & CPU ALLOCATION ---
+# --- 1. SETUP & CPU ALLOCATION (MAX SPEED) ---
 TOTAL_CORES = multiprocessing.cpu_count()
-AI_CORES = max(1, int(TOTAL_CORES * 0.4))
+AI_CORES = max(1, int(TOTAL_CORES * 0.7)) # Give 70% CPU to AI for max speed
 torch.set_num_threads(AI_CORES)
-print(f"🖥️ Total Cores: {TOTAL_CORES} | Allocating {AI_CORES} cores strictly to AI.")
+print(f"🖥️ Total Cores: {TOTAL_CORES} | Dedicated AI Cores: {AI_CORES}")
 
 raw_data = []
-suggestion_list = []
+suggestion_tuples = []
 bm25_index = None
-MAX_CONCURRENT_AI = max(1, AI_CORES)
+MAX_CONCURRENT_AI = max(1, AI_CORES // 2)
 ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 ai_generator = None
-
-# --- TRACKING & TRUST SYSTEM ---
-URL_STATS = {} # Format: { "https://example.com": {"user_visits": 0, "ai_visits": 0, "is_trusted": False} }
-
-# --- AI CACHING SYSTEM ---
+URL_STATS = {}
 AI_RESPONSE_CACHE = {}
 MAX_CACHE_SIZE = 5000
-STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are"}
+
+TOKEN_RE = re.compile(r'[a-z0-9]+')
+STOP_WORDS = {"the", "a", "an", "and", "or", "but", "in", "on", "at", "to", "for", "of", "with", "is", "are", "by", "this", "that", "it"}
 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", "savita bhabhi", "xvideos", "pornhub", 
-    "brazzers", "mia khalifa", "bf video", "blue film", "kamasutra", "nude",
+    "brazzers", "mia khalifa", "bf video", "blue film", "kamasutra", 
     "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 & INDEXING ---
 def tokenize(text):
     if not text: return []
-    tokens = re.findall(r'[a-z0-9]+', str(text).lower())
-    return [t for t in tokens if t not in STOP_WORDS]
+    return [t for t in TOKEN_RE.findall(str(text).lower()) if t not in STOP_WORDS]
 
 def load_and_index_dataset():
-    global bm25_index, raw_data, suggestion_list
-    print("📥 Loading Dataset & Building Deep Index...")
+    global bm25_index, raw_data, suggestion_tuples
+    print("📥 Loading Dataset & Building Fast Index...")
     dataset_path = "dataset.jsonl"
     
     if HF_TOKEN:
@@ -68,184 +64,156 @@ def load_and_index_dataset():
                 repo_type="dataset",
                 token=HF_TOKEN 
             )
-        except Exception as e:
-            print(f"⚠️ Hugging Face Download failed: {e}")
+        except Exception:
+            pass
 
     tokenized_corpus = []
-    MAX_RECORDS_TO_LOAD = 150000 # Prevent OOM Crash
-
+    suggestion_counter = Counter()
+    
     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 >= 150000: break 
                 if line.strip():
                     try:
                         doc = json.loads(line)
                         url = doc.get("url", "")
-                        
-                        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())
+                        desc_str = doc.get("description", "")
                         
-                        title_tokens = tokenize(title_str)
-                        desc_tokens = tokenize(doc.get("description", ""))
-                        url_tokens = tokenize(url)
+                        parsed_url = urllib.parse.urlparse(url)
+                        domain = parsed_url.netloc.lower().replace("www.", "")
+                        clean_domain_name = domain.split('.')[0] if '.' in domain else domain
+
+                        doc_text = f"{title_str} {desc_str} {url}".lower()
+                        is_adult = doc.get("tag", "general").lower() in ["adult", "nsfw"] or any(aw in doc_text for aw in ADULT_KEYWORDS)
+
+                        doc["_fast_domain"] = domain
+                        doc["_fast_path"] = parsed_url.path.strip("/")
+                        doc["_fast_title"] = title_str.lower()
+                        doc["_fast_desc"] = desc_str.lower()
+                        doc["_is_adult"] = is_adult
+                        raw_data.append(doc)
                         
-                        boosted_tokens = (title_tokens * 5) + (desc_tokens * 3) + url_tokens
-                        tokenized_corpus.append(boosted_tokens)
-                    except: 
-                        pass
+                        if not is_adult:
+                            if len(clean_domain_name) > 2 and clean_domain_name not in ADULT_KEYWORDS:
+                                suggestion_counter[clean_domain_name] += 5 
+                            
+                            clean_t_tokens = [w for w in tokenize(title_str) if w not in ADULT_KEYWORDS]
+                            for w in clean_t_tokens:
+                                if len(w) > 2: suggestion_counter[w] += 1
+                            for j in range(len(clean_t_tokens) - 1):
+                                suggestion_counter[f"{clean_t_tokens[j]} {clean_t_tokens[j+1]}"] += 2
+
+                        tokenized_corpus.append((tokenize(title_str) * 5) + (tokenize(desc_str) * 3) + tokenize(url))
+                    except: pass
                     
         if tokenized_corpus:
-            print("⚡ Building High-Speed BM25 Okapi Index...")
-            suggestion_list = sorted(list(set(suggestion_list)))
+            suggestion_tuples = suggestion_counter.most_common(40000)
             bm25_index = BM25Okapi(tokenized_corpus)
-            print(f"🚀 Indexing Complete! {len(raw_data)} records ready for search.")
-    else:
-        print("❌ Dataset file not found!")
+            print(f"🚀 Indexing Complete! {len(raw_data)} records ready.")
 
 def load_ai_models():
     global ai_generator
-    print("🤖 Loading Small & Fast 360M Text AI Model...")
+    print("🤖 Loading LiquidAI/LFM2.5-350M for BLAZING FAST inference...")
     try:
-        # 🔥 Using a small 360M parameter model for strict fast performance
+        # Optimization: use float32 on CPU but pipeline optimized
         ai_generator = pipeline(
             "text-generation", 
-            model="HuggingFaceTB/SmolLM2-360M-Instruct", 
+            model="LiquidAI/LFM2.5-350M", 
             device=-1, 
-            torch_dtype=torch.float32
+            torch_dtype=torch.float32,
+            trust_remote_code=True
         )
-        print("✅ Small AI Model Loaded! System ready.")
+        print("✅ LiquidAI Model Loaded! Superfast mode activated.")
     except Exception as e:
         print(f"❌ Model Loading Failed: {e}")
 
-# --- 3. SMART SCRIPT RAG ---
-def chunk_text(text, chunk_size=120):
+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=3) as resp:
+            async with session.get(url, timeout=1.5) as resp:
                 if resp.status == 200:
                     html = await resp.text()
                     soup = BeautifulSoup(html, "html.parser")
-                    tags_to_extract = ['p', 'h1', 'h2', 'h3']
-                    text = " ".join([tag.get_text(strip=True) for tag in soup.find_all(tags_to_extract)])
-                    
+                    text = " ".join([tag.get_text(strip=True) for tag in soup.find_all(['p', 'h1', 'h2'])])
                     if not text: return ""
                     
                     chunks = chunk_text(text)
                     query_tokens = set(tokenize(query))
                     
-                    best_chunk = ""
-                    max_score = 0
-                    early_stop_threshold = max(1, len(query_tokens) // 2) 
-
+                    best_chunk, max_score = "", 0
                     for chunk in chunks:
-                        chunk_tokens = set(tokenize(chunk))
-                        score = len(query_tokens.intersection(chunk_tokens))
+                        score = len(query_tokens.intersection(tokenize(chunk)))
                         if score > max_score:
-                            max_score = score
-                            best_chunk = chunk
-                        if max_score >= early_stop_threshold and max_score > 0:
-                            break
+                            max_score, best_chunk = score, chunk
+                        if max_score >= max(1, len(query_tokens) // 2): break
                     return best_chunk
     except:
         return ""
     return ""
 
-# --- 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
+    if not raw_data or not bm25_index: return [], "Search index not ready.", 0
 
     query = raw_query.strip().lower()
     query_tokens = tokenize(query)
-    user_wants_adult = any(adult_word in query for adult_word in ADULT_KEYWORDS)
+    user_wants_adult = any(aw in query for aw in ADULT_KEYWORDS)
     
     scores = bm25_index.get_scores(query_tokens) if query_tokens else [1.0] * len(raw_data)
-    scored_results = []
-    seen_urls = set() 
+    scored_results, seen_urls = [], set()
+    query_clean = query.replace(" ", "")
 
     for index, score in enumerate(scores):
         if score > 0.1 or not query_tokens:
             doc = raw_data[index]
-            doc_url = doc.get("url", "")
-            
+            doc_url = doc["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 
+            if not user_wants_adult and doc["_is_adult"]: continue 
                 
-            final_score = score
-            query_clean = query.replace(" ", "")
+            final_score, domain = score, doc["_fast_domain"]
             domain_clean = domain.replace(" ", "")
+            doc_title, doc_desc = doc["_fast_title"], doc["_fast_desc"]
             
-            # 🔥 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
+                final_score *= 50.0 if not doc["_fast_path"] else 10.0
             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
+                final_score *= 25.0 if not doc["_fast_path"] else 5.0
                     
-            # 🔥 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
+            if query in doc_title: final_score *= 5.0 
+            elif all(w in doc_title for w in query_tokens): final_score *= 3.0 
+            if query in doc_desc: final_score *= 2.0 
                 
-            # 🔥 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  
-                
+            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 = {k: v for k, v in doc.items() if not k.startswith("_")}
             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
     prediction = "High confidence" if best_score > 15 else "Medium confidence" if best_score > 5 else "Low confidence"
+    
+    return [{"match_score": round(s, 2), **doc} for s, doc in scored_results[:100]], prediction, best_score
 
-    final_results = [{"match_score": round(s, 2), **doc} for s, doc in scored_results[:100]]
-    return final_results, prediction, best_score
-
-# --- 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)
+    
+    # 🔥 GREEDY DECODING (do_sample=False) FOR 3X FASTER AI RESPONSES
+    result = ai_generator(
+        prompt, 
+        max_new_tokens=max_tokens, 
+        do_sample=False, 
+        truncation=True
+    )
     text_out = result[0]['generated_text'].split("Assistant:")[-1].strip()
     return text_out
 
@@ -259,37 +227,21 @@ async def run_ai_task(messages, max_tokens: int, query_key: str):
         AI_RESPONSE_CACHE[query_key] = response
         return response
 
-# --- 5. FASTAPI APP & ENDPOINTS ---
 @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 Advanced Text AI & RAG Engine", lifespan=lifespan)
-
-app.add_middleware(
-    CORSMiddleware, allow_origins=["*"], allow_credentials=True, allow_methods=["*"], allow_headers=["*"]
-)
+app = FastAPI(title="Nebulix Fast AI Search", lifespan=lifespan)
+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)):
-    q = q.lower()
-    matches = [s for s in suggestion_list if s.startswith(q)][:7]
+    q = q.lower().strip()
+    matches = [word for word, count in suggestion_tuples if word.startswith(q)][:7]
     return {"query": q, "suggestions": matches}
 
-@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}
-    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(...)):
-    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(...)):
     req_id = str(uuid.uuid4())
@@ -300,36 +252,28 @@ async def pure_search(q: str = Query(...)):
 async def ai_mode_short(q: str = Query(...)):
     req_id = str(uuid.uuid4())
     results, _, _ = await perform_core_search(q)
-    if not results: 
-        return {"request_id": req_id, "ai_response": "I couldn't find any results.", "results": []}
+    if not results: return {"request_id": req_id, "ai_response": "I couldn't find any results.", "results": []}
 
-    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)
+    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 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 1 short sentence. Put important keywords in (parentheses)."},
         {"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}")
+    # 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"])
-    
     results, _, _ = await perform_core_search(q)
+    
     rag_context = ""
     used_sources = []
 
     if needs_rag and results:
-        top_urls = results[:10]
-        
+        # 🔥 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)
         
@@ -337,14 +281,6 @@ async def process_rag_chat(q: str):
         for i, chunk in enumerate(scraped_chunks):
             doc = top_urls[i]
             url = doc["url"]
-            
-            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
-                
-            doc["trust_stats"] = URL_STATS[url]
-
             if chunk:
                 context_parts.append(f"Source [{i+1}] ({url}): {chunk}")
                 used_sources.append(doc)
@@ -354,34 +290,21 @@ async def process_rag_chat(q: str):
                 
         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)."
-
+    system_prompt = "You are a Smart AI Search Assistant. Provide a short, direct answer. ENCLOSE IMPORTANT KEYWORDS IN (PARENTHESES). Mention Source URLs if used."
     messages = [
         {"role": "system", "content": system_prompt},
-        {"role": "user", "content": f"Query: {q}\n\nLive Context Data:\n{rag_context}" if needs_rag else f"Query: {q}"}
+        {"role": "user", "content": f"Query: {q}\nContext:\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 {
-        "request_id": req_id, 
-        "ai_response": ai_reply, 
-        "rag_used": needs_rag,
-        "sources": used_sources,
-        "results": results[:15]
-    }
+    # 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]}
 
 @app.get("/aimodechat")
-async def ai_mode_chat_get(q: str = Query(...)):
-    return await process_rag_chat(q)
+async def ai_mode_chat_get(q: str = Query(...)): return await process_rag_chat(q)
 
 @app.post("/aimodechat")
-async def ai_mode_chat_post(q: str = Form(...)):
-    return await process_rag_chat(q)
+async def ai_mode_chat_post(q: str = Form(...)): return await process_rag_chat(q)
 
 if __name__ == "__main__":
     uvicorn.run(app, host="0.0.0.0", port=7860)