Upload app.py with huggingface_hub

app.py

@@ -19,8 +19,9 @@ os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
 
 N_CPUS = os.cpu_count() or 2
 BATCH_SIZE = 64  # max clips per ONNX forward pass
-API_TIMEOUT = 20   # seconds per datasets-server attempt
+API_TIMEOUT = 15   # seconds per datasets-server attempt
 API_RETRIES = 2
+STREAMING_TIMEOUT = 45  # seconds before giving up on streaming fallback
 
 
 def _ts() -> str:
@@ -72,39 +73,51 @@ def _batch_embed(arrays: list[np.ndarray]) -> np.ndarray:
 
 # ── Audio fetching ────────────────────────────────────────────────────────────
 
+def _parse_audio_urls(rows: list) -> list[str]:
+    urls = []
+    for row in rows:
+        audio = row["row"].get("audio", {})
+        if isinstance(audio, list):
+            audio = audio[0] if audio else {}
+        if isinstance(audio, dict) and "src" in audio:
+            urls.append(audio["src"])
+    return urls
+
+
 def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> tuple[list[str], str]:
-    """Fetch audio URLs from datasets-server with retry. Returns (urls, status)."""
+    """Try /rows then /first-rows endpoints with retry. Returns (urls, status)."""
     headers = {"Authorization": f"Bearer {token}"} if token else {}
-    params = {"dataset": repo_id, "config": "default", "split": "train", "offset": 0, "length": n}
+    endpoints = [
+        ("/rows",       {"dataset": repo_id, "config": "default", "split": "train", "offset": 0, "length": n}),
+        ("/first-rows", {"dataset": repo_id, "config": "default", "split": "train"}),
+    ]
     last_err = ""
-    for attempt in range(1, API_RETRIES + 1):
-        try:
-            t0 = time.time()
-            resp = requests.get(
-                f"{DATASETS_SERVER}/rows", params=params,
-                headers=headers, timeout=API_TIMEOUT,
-            )
-            elapsed = int((time.time() - t0) * 1000)
-            if not resp.ok:
-                last_err = f"HTTP {resp.status_code} (attempt {attempt}, {elapsed}ms)"
-                continue
-            rows = resp.json().get("rows", [])
-            if not rows:
-                return [], f"empty rows ({elapsed}ms)"
-            urls = []
-            for row in rows:
-                audio = row["row"].get("audio", {})
-                if isinstance(audio, list):
-                    audio = audio[0] if audio else {}
-                if isinstance(audio, dict) and "src" in audio:
-                    urls.append(audio["src"])
-            if not urls:
-                return [], f"no src field ({elapsed}ms)"
-            return urls, f"api={elapsed}ms"
-        except requests.Timeout:
-            last_err = f"timeout>{API_TIMEOUT}s (attempt {attempt})"
-        except Exception as exc:
-            last_err = f"{exc} (attempt {attempt})"
+    for endpoint, params in endpoints:
+        for attempt in range(1, API_RETRIES + 1):
+            try:
+                t0 = time.time()
+                resp = requests.get(
+                    f"{DATASETS_SERVER}{endpoint}", params=params,
+                    headers=headers, timeout=API_TIMEOUT,
+                )
+                elapsed = int((time.time() - t0) * 1000)
+                if not resp.ok:
+                    last_err = f"{endpoint} HTTP {resp.status_code} (attempt {attempt}, {elapsed}ms)"
+                    break  # 4xx/5xx — no point retrying same endpoint
+                key = "rows" if endpoint == "/rows" else "rows"
+                rows = resp.json().get(key, [])
+                if not rows:
+                    last_err = f"{endpoint} empty ({elapsed}ms)"
+                    break
+                urls = _parse_audio_urls(rows[:n])
+                if urls:
+                    return urls, f"{endpoint} {elapsed}ms"
+                last_err = f"{endpoint} no src ({elapsed}ms)"
+                break
+            except requests.Timeout:
+                last_err = f"{endpoint} timeout>{API_TIMEOUT}s (attempt {attempt})"
+            except Exception as exc:
+                last_err = f"{endpoint} {exc} (attempt {attempt})"
     return [], last_err
 
 
@@ -149,24 +162,37 @@ def _fetch_repo_audio(
             )
             return repo, arrays
 
-        # Streaming fallback
-        log.append(f"[{_ts()}] {short}: ⚠ datasets-server failed ({api_status}) → streaming fallback")
+        # Streaming fallback with timeout
+        log.append(f"[{_ts()}] {short}: ⚠ datasets-server failed ({api_status}) → streaming fallback (timeout={STREAMING_TIMEOUT}s)")
         t1 = time.time()
-        from datasets import load_dataset, Audio as HFAudio
-        ds = load_dataset(repo, split="train", streaming=True, token=token)
-        ds = ds.cast_column("audio", HFAudio(decode=False))
-        arrays = []
-        for j, row in enumerate(ds):
-            if j >= n_samples:
-                break
-            raw = row["audio"]
-            audio_bytes = raw.get("bytes") or open(raw["path"], "rb").read()
-            a, sr = sf.read(io.BytesIO(audio_bytes))
-            arrays.append(_to_array(a, sr, audio_sec))
-            log.append(f"[{_ts()}] {short}: streaming clip {j+1}/{n_samples} ({len(audio_bytes)//1024}KB, {int((time.time()-t1)*1000)}ms so far)")
-        total_ms = int((time.time() - t0) * 1000)
-        log.append(f"[{_ts()}] {short}: ✓ streaming done, {len(arrays)} clips, total={total_ms}ms")
-        return repo, arrays
+
+        def _stream() -> list[np.ndarray]:
+            from datasets import load_dataset, Audio as HFAudio
+            ds = load_dataset(repo, split="train", streaming=True, token=token)
+            ds = ds.cast_column("audio", HFAudio(decode=False))
+            result = []
+            for j, row in enumerate(ds):
+                if j >= n_samples:
+                    break
+                raw = row["audio"]
+                audio_bytes = raw.get("bytes") or open(raw["path"], "rb").read()
+                a, sr = sf.read(io.BytesIO(audio_bytes))
+                result.append(_to_array(a, sr, audio_sec))
+                elapsed = int((time.time() - t1) * 1000)
+                log.append(f"[{_ts()}] {short}: streaming clip {j+1}/{n_samples} ({len(audio_bytes)//1024}KB, {elapsed}ms so far)")
+            return result
+
+        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
+            fut = ex.submit(_stream)
+            try:
+                arrays = fut.result(timeout=STREAMING_TIMEOUT)
+                total_ms = int((time.time() - t0) * 1000)
+                log.append(f"[{_ts()}] {short}: ✓ streaming done, {len(arrays)} clips, total={total_ms}ms")
+                return repo, arrays
+            except concurrent.futures.TimeoutError:
+                total_ms = int((time.time() - t0) * 1000)
+                log.append(f"[{_ts()}] {short}: ✗ streaming timeout after {STREAMING_TIMEOUT}s — skipping")
+                return repo, []
 
     except Exception as exc:
         log.append(f"[{_ts()}] {short}: ✗ {exc} (total={int((time.time()-t0)*1000)}ms)")