Update app.py

app.py

@@ -1,12 +1,17 @@
 from fastapi import FastAPI, Query
 from fastapi.middleware.cors import CORSMiddleware
 from datasets import load_dataset
-from huggingface_hub import HfApi
+from huggingface_hub import HfApi, login
 import os
 
 os.environ["HF_HOME"] = "/tmp/hf"
 
-app = FastAPI(title="Telegram UID Multi-Search")
+# HF_TOKEN se login (fast + authenticated)
+HF_TOKEN = os.environ.get("HF_TOKEN")
+if HF_TOKEN:
+    login(token=HF_TOKEN)
+
+app = FastAPI()
 
 app.add_middleware(
     CORSMiddleware,
@@ -17,121 +22,61 @@ app.add_middleware(
 
 DATASET = "CodeXDevloper/MergedTgDataset"
 
-def get_files(folder=None):
-    """Kisi folder ki saari parquet files nikalo"""
+def get_files():
     api = HfApi()
     files = api.list_repo_files(DATASET, repo_type="dataset")
-    
-    if folder:
-        return [f for f in files if f.startswith(folder + "/") and f.endswith(".parquet")]
-    else:
-        return [f for f in files if "/" not in f and f.endswith(".parquet")]
+    return [f for f in files if f.endswith(".parquet")]
 
 @app.get("/")
 def root():
-    return {
-        "status": "running",
-        "usage": "/search?uid=385167100",
-        "search_order": "Telegram_V2 → TG_DATA_PARTS → Root"
-    }
+    return {"status": "running", "usage": "/search?uid=XXX&offset=0"}
 
 @app.get("/search")
-def search_uid(uid: str = Query(..., description="Telegram UID")):
-    """
-    AUTO MULTI-FOLDER SEARCH:
-    1. Telegram_V2 mein dhoondho
-    2. Nahi mila toh TG_DATA_PARTS mein
-    3. Nahi mila toh Root files mein
-    """
+def search_uid(
+    uid: str = Query(...),
+    offset: int = Query(0, ge=0)
+):
     try:
-        search_log = []
+        files = get_files()
         
-        # Priority order
-        folders = [
-            ("Telegram_V2", "Telegram_V2"),
-            ("TG_DATA_PARTS", "TG_DATA_PARTS"),
-            (None, "Root")
-        ]
+        ds = load_dataset(
+            DATASET,
+            data_files={"train": files[:50]},  # 50 files per request
+            split="train",
+            streaming=True
+        )
         
-        for folder_key, folder_name in folders:
-            # Folder ki files nikalo
-            files = get_files(folder_key)
-            
-            if not files:
-                search_log.append({
-                    "folder": folder_name,
-                    "status": "no_files",
-                    "scanned": 0
-                })
-                continue
-            
-            # Sirf 20 files per folder (RAM safe)
-            files = files[:20]
-            
+        # Offset skip
+        iterator = iter(ds)
+        for _ in range(offset):
             try:
-                ds = load_dataset(
-                    DATASET,
-                    data_files={"train": files},
-                    split="train",
-                    streaming=True
-                )
-            except Exception as e:
-                search_log.append({
-                    "folder": folder_name,
-                    "status": "load_error",
-                    "error": str(e)[:100]
-                })
-                continue
-            
-            # Is folder mein scan karo
-            scanned = 0
-            max_scan = 100000  # 1 lakh rows per folder
-            
-            for row in ds:
-                scanned += 1
-                
-                if uid == str(row.get("user_id", "")).strip():
-                    # ✅ MIL GAYA
-                    return {
-                        "found": True,
-                        "uid": uid,
-                        "found_in": folder_name,
-                        "scanned_in_folder": scanned,
-                        "search_log": search_log + [{
-                            "folder": folder_name,
-                            "status": "FOUND",
-                            "scanned": scanned
-                        }],
-                        "data": {
-                            "user_id": row.get("user_id"),
-                            "username": row.get("username"),
-                            "first_name": row.get("first_name"),
-                            "last_name": row.get("last_name"),
-                            "phone": row.get("phone"),
-                            "email": row.get("email"),
-                            "status": row.get("status"),
-                            "linked_id": row.get("linked_id"),
-                            "linked_name": row.get("linked_name"),
-                            "linked_handle": row.get("linked_handle"),
-                        }
-                    }
-                
-                if scanned >= max_scan:
-                    break
-            
-            search_log.append({
-                "folder": folder_name,
-                "status": "not_found",
-                "scanned": scanned
-            })
+                next(iterator)
+            except StopIteration:
+                return {"found": False, "uid": uid, "message": "End", "next_offset": None}
+        
+        # 2 lakh rows scan
+        scanned = 0
+        max_scan = 200000
+        
+        for row in iterator:
+            scanned += 1
+            if uid == str(row.get("user_id", "")).strip():
+                return {
+                    "found": True,
+                    "uid": uid,
+                    "offset": offset,
+                    "scanned": scanned,
+                    "data": row
+                }
+            if scanned >= max_scan:
+                break
         
-        # Kahin nahi mila
         return {
             "found": False,
             "uid": uid,
-            "search_log": search_log,
-            "message": "Telegram_V2, TG_DATA_PARTS, aur Root - teeno mein nahi mila"
+            "offset": offset,
+            "scanned": scanned,
+            "next_offset": offset + scanned
         }
-    
     except Exception as e:
-        return {"error": str(e), "uid": uid}
+        return {"error": str(e)}