Spaces:
Sleeping
Sleeping
| from fastapi import FastAPI, Query | |
| from fastapi.middleware.cors import CORSMiddleware | |
| from datasets import load_dataset | |
| from huggingface_hub import HfApi, login | |
| from concurrent.futures import ThreadPoolExecutor, as_completed | |
| import os | |
| import time | |
| os.environ["HF_HOME"] = "/tmp/hf" | |
| # HF login (fast + authenticated) | |
| HF_TOKEN = os.environ.get("HF_TOKEN") | |
| if HF_TOKEN: | |
| login(token=HF_TOKEN) | |
| app = FastAPI(title="Telegram UID Search API") | |
| app.add_middleware( | |
| CORSMiddleware, | |
| allow_origins=["*"], | |
| allow_methods=["*"], | |
| allow_headers=["*"], | |
| ) | |
| DATASET = "CodeXDevloper/MergedTgDataset" | |
| WORKERS = 20 # 20 parallel workers | |
| BATCH_SIZE = 10 # 10 batches | |
| def get_all_files(): | |
| api = HfApi() | |
| files = api.list_repo_files(DATASET, repo_type="dataset") | |
| return sorted([f for f in files if f.endswith(".parquet")]) | |
| def search_in_file(file_path: str, uid: str): | |
| """Ek file mein UID search karo (worker function)""" | |
| try: | |
| ds = load_dataset( | |
| DATASET, | |
| data_files={"train": [file_path]}, | |
| split="train", | |
| streaming=True | |
| ) | |
| scanned = 0 | |
| for row in ds: | |
| scanned += 1 | |
| if uid == str(row.get("user_id", "")).strip(): | |
| return { | |
| "found": True, | |
| "file": file_path, | |
| "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"), | |
| } | |
| } | |
| return {"found": False, "file": file_path, "scanned": scanned} | |
| except Exception as e: | |
| return {"found": False, "file": file_path, "error": str(e), "scanned": 0} | |
| def root(): | |
| return { | |
| "status": "running", | |
| "usage": "/search?uid=385167100", | |
| "note": f"Parallel search with {WORKERS} workers, batch size {BATCH_SIZE}, no timeout" | |
| } | |
| def search_uid(uid: str = Query(..., description="Telegram UID")): | |
| """ | |
| Poore dataset mein UID dhoondho - 20 parallel workers, no timeout. | |
| """ | |
| try: | |
| files = get_all_files() | |
| if not files: | |
| return {"error": "Koi parquet file nahi mili", "uid": uid} | |
| start_time = time.time() | |
| total_scanned = 0 | |
| # 20 workers, batch me files process karo | |
| with ThreadPoolExecutor(max_workers=WORKERS) as executor: | |
| # batch_size ke hisaab se chunks banao | |
| for i in range(0, len(files), BATCH_SIZE): | |
| batch = files[i:i + BATCH_SIZE] | |
| futures = {executor.submit(search_in_file, f, uid): f for f in batch} | |
| for future in as_completed(futures): | |
| result = future.result() | |
| total_scanned += result.get("scanned", 0) | |
| if result.get("found"): | |
| # Baaki futures cancel karne ki koshish | |
| for f in futures: | |
| f.cancel() | |
| return { | |
| "found": True, | |
| "uid": uid, | |
| "scanned": total_scanned, | |
| "time_taken": round(time.time() - start_time, 2), | |
| "file": result.get("file"), | |
| "data": result.get("data"), | |
| } | |
| return { | |
| "found": False, | |
| "uid": uid, | |
| "scanned": total_scanned, | |
| "time_taken": round(time.time() - start_time, 2), | |
| "message": "Poore dataset mein nahi mila" | |
| } | |
| except Exception as e: | |
| return {"error": str(e), "uid": uid} |