Update app.py

app.py

@@ -1,6 +1,7 @@
 from fastapi import FastAPI, Query
 from fastapi.middleware.cors import CORSMiddleware
 from datasets import load_dataset
+from huggingface_hub import HfApi
 import os
 
 os.environ["HF_HOME"] = "/tmp/hf"
@@ -16,6 +17,13 @@ app.add_middleware(
 
 DATASET = "CodeXDevloper/MergedTgDataset"
 
+def get_all_parquet_files():
+    """Saari parquet files nikalo - saare folders se"""
+    api = HfApi()
+    files = api.list_repo_files(DATASET, repo_type="dataset")
+    parquet_files = [f for f in files if f.endswith(".parquet")]
+    return parquet_files
+
 @app.get("/")
 def root():
     return {
@@ -27,17 +35,25 @@ def root():
 @app.get("/search")
 def search_uid(
     uid: str = Query(..., description="Telegram UID"),
-    offset: int = Query(0, ge=0, description="Kahan se start karna hai (default 0)")
+    offset: int = Query(0, ge=0, description="Kahan se start karna hai")
 ):
     """
-    Poore dataset mein UID dhoondho.
-    - offset=0 se start karo
-    - Agar nahi mila toh next_offset milega
-    - Woh next_offset daal ke dobara call karo
-    - Aise karte karte poora dataset cover hoga
+    POORE dataset mein UID dhoondho - saare folders se
     """
     try:
-        ds = load_dataset(DATASET, split="train", streaming=True)
+        # Saari parquet files nikalo (saare folders)
+        parquet_files = get_all_parquet_files()
+        
+        if not parquet_files:
+            return {"error": "Koi parquet file nahi mili"}
+        
+        # Saari files ko load karo (streaming mode)
+        ds = load_dataset(
+            DATASET,
+            data_files=parquet_files,
+            split="train",
+            streaming=True
+        )
         
         # Offset tak skip karo
         iterator = iter(ds)
@@ -56,7 +72,7 @@ def search_uid(
                     "next_offset": None
                 }
         
-        # Ab 5 lakh rows scan karo (HF Space 30 second timeout ke liye safe)
+        # Ab 5 lakh rows scan karo
         scanned = 0
         max_scan = 500000
         
@@ -92,6 +108,7 @@ def search_uid(
             "offset": offset,
             "scanned": scanned,
             "next_offset": offset + scanned,
+            "total_files": len(parquet_files),
             "message": f"Is chunk mein nahi mila. Ab offset={offset + scanned} try karo."
         }