Update app.py

app.py

@@ -10,14 +10,16 @@ import urllib.parse
 from datetime import datetime
 from contextlib import asynccontextmanager
 from fastapi import FastAPI, Query, Request, UploadFile, File
-from fastapi.responses import JSONResponse
+from fastapi.responses import JSONResponse, HTMLResponse
 from fastapi.middleware.cors import CORSMiddleware
 from huggingface_hub import hf_hub_download
 from rank_bm25 import BM25Okapi
 import torch
 from PIL import Image
 import io
-from transformers import pipeline
+
+# DIRECT IMPORTS TO FIX THE PIPELINE BUG
+from transformers import pipeline, BlipProcessor, BlipForConditionalGeneration
 
 # --- 1. SETUP & CPU ALLOCATION (40% for AI) ---
 TOTAL_CORES = multiprocessing.cpu_count()
@@ -29,12 +31,12 @@ print(f"🖥️ Total Cores: {TOTAL_CORES} | Allocating {AI_CORES} cores strictl
 raw_data = []
 bm25_index = None
 
-# Queueing System using Semaphore (Text & Vision dono ke liye same queue)
 MAX_CONCURRENT_AI = max(1, AI_CORES) 
 ai_semaphore = asyncio.Semaphore(MAX_CONCURRENT_AI)
 
 ai_generator = None
-vision_ai_generator = None # Naya Vision AI
+vision_processor = None
+vision_model = None 
 
 # --- AI CACHING SYSTEM ---
 AI_RESPONSE_CACHE = {}
@@ -104,20 +106,20 @@ def load_and_index_dataset():
         print("❌ Dataset file not found!")
 
 def load_ai_models():
-    global ai_generator, vision_ai_generator
+    global ai_generator, vision_processor, vision_model
     print("🤖 Loading Lightweight AI Models (Text & Vision)...")
     try:
-        # 1. Text AI (Lightweight Qwen)
+        # 1. Text AI
         ai_generator = pipeline("text-generation", model="Qwen/Qwen2.5-0.5B-Instruct", device=-1, dtype=torch.float32)
         
-        # 2. Vision AI (BLIP - Ultra lightweight, perfect for CPU)
-        # Ye model image ko dekh kar batata hai ki usme kya hai (AI Review)
+        # 2. Vision AI (Fix: Direct loading to bypass pipeline task errors)
         print("👁️ Loading Vision AI...")
-        vision_ai_generator = pipeline("image-to-text", model="Salesforce/blip-image-captioning-base", device=-1)
+        vision_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+        vision_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
         
         print("✅ All AI Models Loaded! System ready.")
     except Exception as e:
-        print("❌ Model Loading Failed:", e)
+        print(f"❌ Model Loading Failed: {e}")
 
 
 # --- 3. CORE SEARCH ALGORITHM ---
@@ -162,11 +164,9 @@ async def perform_core_search(raw_query: str):
             doc_tag = str(doc.get("tag", "general")).lower()
             if tag_filter and tag_filter != doc_tag: continue
             
-            # Adult Content Strict Filter
             is_adult = doc_tag in ["adult", "nsfw"] or any(w in str(doc).lower() for w in ADULT_KEYWORDS)
             if not user_wants_adult and is_adult: continue 
             
-            # Ranking Boosts
             final_score = score
             exact_phrase = query.lower()
             doc_title = str(doc.get("title", "")).lower()
@@ -195,20 +195,19 @@ def _generate_text(messages, max_tokens: int):
     return result[0]['generated_text'][-1]['content']
 
 def _generate_image_review(image: Image):
-    result = vision_ai_generator(image, max_new_tokens=50)
-    return result[0]['generated_text']
+    # Direct model generation (Bypasses the buggy pipeline)
+    inputs = vision_processor(image, return_tensors="pt")
+    with torch.no_grad():
+        out = vision_model.generate(**inputs, max_new_tokens=50)
+    return vision_processor.decode(out[0], skip_special_tokens=True).capitalize()
 
 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]
+    if query_key in AI_RESPONSE_CACHE: return AI_RESPONSE_CACHE[query_key]
         
     async with ai_semaphore:
         response = await asyncio.to_thread(_generate_text, messages, max_tokens)
-        
-        if len(AI_RESPONSE_CACHE) > MAX_CACHE_SIZE:
-            AI_RESPONSE_CACHE.clear()
+        if len(AI_RESPONSE_CACHE) > MAX_CACHE_SIZE: AI_RESPONSE_CACHE.clear()
         AI_RESPONSE_CACHE[query_key] = response
         return response
 
@@ -241,56 +240,82 @@ async def rate_limit_middleware(request: Request, call_next):
 
     return await call_next(request)
 
+# NAYA: Frontend Load Karne Ke Liye
+@app.get("/", response_class=HTMLResponse)
+async def serve_ui():
+    if os.path.exists("index.html"):
+        with open("index.html", "r", encoding="utf-8") as f:
+            return f.read()
+    return "<h1>UI Not Found</h1><p>Please save your frontend code as <b>index.html</b> in the same folder as this main.py file.</p>"
+
 @app.get("/search")
 async def pure_search(q: str = Query(...)):
     req_id = str(uuid.uuid4())
     results, prediction, _ = await perform_core_search(q)
     return {"request_id": req_id, "query": q, "results_count": len(results), "results": results}
 
-# --- NEW: ADVANCED SMART IMAGE SEARCH ENDPOINT ---
+@app.get("/aimode")
+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": []}
+        
+    context = " ".join([f"{r.get('title', '')} - {r.get('description', '')}" for r in results[:3]])
+    messages = [
+        {"role": "system", "content": "Provide a brief 2-3 sentence summary based only on the provided context."},
+        {"role": "user", "content": f"Query: {q}\nContext: {context}"}
+    ]
+    
+    ai_reply = await run_ai_task(messages, 150, query_key=f"short_{q}")
+    return {"request_id": req_id, "ai_response": ai_reply, "results": results}
+
+@app.get("/aimodechat")
+async def ai_mode_chat(q: str = Query(...)):
+    req_id = str(uuid.uuid4())
+    results, _, _ = await perform_core_search(q)
+    
+    context_parts = [f"Source [{i+1}]: {r.get('title')} - {r.get('description')}" for i, r in enumerate(results[:5])]
+    messages = [
+        {"role": "system", "content": "You are a Smart AI Search Assistant. Use search sources to answer accurately."},
+        {"role": "user", "content": f"User Query: {q}\n\nContext:\n{chr(10).join(context_parts)}"}
+    ]
+    
+    ai_reply = await run_ai_task(messages, 400, query_key=f"chat_{q}") 
+    return {"request_id": req_id, "ai_response": ai_reply, "results": results}
+
 @app.post("/image-search")
 async def reverse_image_search(file: UploadFile = File(...)):
-    if not vision_ai_generator:
+    if not vision_model: # Yaha updated var check lagaya gaya hai
         return {"error": "Vision AI Model is still loading."}
     
     req_id = str(uuid.uuid4())
     
     try:
-        # 1. Read & Process Uploaded Image
         image_bytes = await file.read()
         img = Image.open(io.BytesIO(image_bytes)).convert("RGB")
         
-        # 2. Get AI Review (Uses Semaphore queue to manage CPU limit)
         async with ai_semaphore:
             ai_review = await asyncio.to_thread(_generate_image_review, img)
             
-        # 3. Use AI Review to search dataset for visually relevant content
         raw_results, _, _ = await perform_core_search(ai_review)
-        
-        # 4. Filter & Format Results (Strictly max 40 results)
         visual_results = []
         
         for doc in raw_results:
-            # Smart SafeSearch Filter (Pehle bhi hota hai, double safety ke liye yaha bhi)
             is_adult = str(doc.get("tag", "general")).lower() in ["adult", "nsfw"] or any(w in str(doc).lower() for w in ADULT_KEYWORDS)
             if is_adult: continue 
             
-            # Result formatting (Hum images, titles aur related URLs return kar rahe hain)
             result_item = {
                 "title": doc.get("title", "No Title"),
                 "url": doc.get("url", ""),
-                "image_url": doc.get("image_url", doc.get("specific_data", {}).get("image", "")), # Try extracting image URL
+                "image_url": doc.get("image_url", doc.get("specific_data", {}).get("image", "")), 
                 "description": doc.get("description", "")
             }
             visual_results.append(result_item)
-            
-            # STRICT LIMIT: Maximum 40 results
-            if len(visual_results) >= 40: 
-                break
+            if len(visual_results) >= 40: break
                 
         return {
             "request_id": req_id,
-            "ai_review": f"Image Analysis: {ai_review.capitalize()}.",
+            "ai_review": f"Image Analysis: {ai_review}.",
             "message": "Found visually related content and images.",
             "results_count": len(visual_results),
             "results": visual_results