Upload app.py with huggingface_hub

app.py

@@ -1,6 +1,7 @@
 import io
 import os
 import time
+import datetime
 import threading
 import concurrent.futures
 import gradio as gr
@@ -18,6 +19,12 @@ 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_RETRIES = 2
+
+
+def _ts() -> str:
+    return datetime.datetime.now().strftime("%H:%M:%S.%f")[:-3]
 
 DATASETS_SERVER = "https://datasets-server.huggingface.co"
 ONNX_MODEL_ID = "fosters/wavlm-base-plus-sv-onnx"
@@ -66,64 +73,85 @@ def _batch_embed(arrays: list[np.ndarray]) -> np.ndarray:
 # ── Audio fetching ────────────────────────────────────────────────────────────
 
 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)."""
     headers = {"Authorization": f"Bearer {token}"} if token else {}
-    resp = requests.get(
-        f"{DATASETS_SERVER}/rows",
-        params={"dataset": repo_id, "config": "default", "split": "train", "offset": 0, "length": n},
-        headers=headers,
-        timeout=30,
-    )
-    if not resp.ok:
-        return [], f"datasets-server {resp.status_code}"
-
-    rows = resp.json().get("rows", [])
-    if not rows:
-        return [], "datasets-server: empty rows"
-
-    sample_audio = rows[0]["row"].get("audio", {})
-    audio_keys = list(sample_audio.keys()) if isinstance(sample_audio, dict) else type(sample_audio).__name__
-
-    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 in audio (keys={audio_keys})"
-    return urls, "ok"
-
-
-def _download_audio(url: str, token: str | None, max_sec: int) -> tuple[np.ndarray, int]:
-    """Download one audio file. Returns (array_at_TARGET_SR, size_kb)."""
+    params = {"dataset": repo_id, "config": "default", "split": "train", "offset": 0, "length": n}
+    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})"
+    return [], last_err
+
+
+def _download_audio(url: str, token: str | None, max_sec: int) -> tuple[np.ndarray, int, int]:
+    """Download one audio file. Returns (array, size_kb, dl_ms)."""
     headers = {"Authorization": f"Bearer {token}"} if token else {}
+    t0 = time.time()
     resp = requests.get(url, headers=headers, timeout=30)
     resp.raise_for_status()
+    dl_ms = int((time.time() - t0) * 1000)
     audio_array, sr = sf.read(io.BytesIO(resp.content))
-    return _to_array(audio_array, sr, max_sec), len(resp.content) // 1024
+    return _to_array(audio_array, sr, max_sec), len(resp.content) // 1024, dl_ms
 
 
 def _fetch_repo_audio(
-    repo: str, n_samples: int, audio_sec: int, token: str | None
-) -> tuple[str, list[np.ndarray], str]:
-    """Fetch + download audio for one repo. Returns (repo, arrays, log)."""
+    repo: str, n_samples: int, audio_sec: int, token: str | None,
+    log: list[str],
+) -> tuple[str, list[np.ndarray]]:
+    """Fetch + download audio for one repo. Appends timestamped entries to log."""
+    short = repo.split("/")[-1]
     t0 = time.time()
+
     try:
+        log.append(f"[{_ts()}] {short}: fetching URLs from datasets-server…")
         urls, api_status = _fetch_audio_urls(repo, n_samples, token)
         fetch_ms = int((time.time() - t0) * 1000)
 
         if urls:
+            log.append(f"[{_ts()}] {short}: {api_status} → {len(urls)} URLs")
             with concurrent.futures.ThreadPoolExecutor(max_workers=len(urls)) as ex:
-                results = list(ex.map(lambda u: _download_audio(u, token, audio_sec), urls))
+                futures = [ex.submit(_download_audio, u, token, audio_sec) for u in urls]
+                results = [f.result() for f in futures]
             arrays = [r[0] for r in results]
-            avg_kb = int(np.mean([r[1] for r in results]))
-            dl_ms = int((time.time() - t0) * 1000) - fetch_ms
-            log = f"api={fetch_ms}ms dl={dl_ms}ms ({avg_kb}KB/file) [{api_status}]"
-            return repo, arrays, log
+            sizes  = [r[1] for r in results]
+            dls    = [r[2] for r in results]
+            total_ms = int((time.time() - t0) * 1000)
+            log.append(
+                f"[{_ts()}] {short}: ✓ {len(arrays)} clips "
+                f"dl=[{', '.join(str(d)+'ms' for d in dls)}] "
+                f"size=[{', '.join(str(s)+'KB' for s in sizes)}] "
+                f"total={total_ms}ms"
+            )
+            return repo, arrays
 
         # Streaming fallback
+        log.append(f"[{_ts()}] {short}: ⚠ datasets-server failed ({api_status}) → streaming fallback")
+        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))
@@ -135,11 +163,14 @@ def _fetch_repo_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)
-        return repo, arrays, f"streaming fallback total={total_ms}ms [reason: {api_status}]"
+        log.append(f"[{_ts()}] {short}: ✓ streaming done, {len(arrays)} clips, total={total_ms}ms")
+        return repo, arrays
 
     except Exception as exc:
-        return repo, [], f"ERROR: {exc}"
+        log.append(f"[{_ts()}] {short}: ✗ {exc} (total={int((time.time()-t0)*1000)}ms)")
+        return repo, []
 
 
 # ── Main ──────────────────────────────────────────────────────────────────────
@@ -158,35 +189,44 @@ def identify_speakers(
 
     token = hf_token.strip() or os.environ.get("HF_TOKEN") or None
 
+    log: list[str] = []
+    t_total = time.time()
+
+    log.append(f"[{_ts()}] === START: {len(repos)} repos, {samples_per_book} sample/book, {audio_sec}s/clip ===")
+    log.append(f"[{_ts()}] CPU count: {N_CPUS}, batch size: {BATCH_SIZE}")
+
     progress(0, desc="Loading model…")
+    t_model = time.time()
     _load_model()
+    log.append(f"[{_ts()}] model loaded in {int((time.time()-t_model)*1000)}ms")
 
     # ── Phase 1: download all audio in parallel (I/O bound) ──────────────────
+    log.append(f"[{_ts()}] --- Phase 1: download ({N_CPUS*2} workers) ---")
     progress(0.05, desc=f"Downloading audio from {len(repos)} datasets…")
     repo_arrays: dict[str, list[np.ndarray]] = {}
-    fetch_logs: list[str] = []
-    errors: list[str] = []
 
     with concurrent.futures.ThreadPoolExecutor(max_workers=N_CPUS * 2) as ex:
         futures = {
-            ex.submit(_fetch_repo_audio, repo, int(samples_per_book), int(audio_sec), token): repo
+            ex.submit(_fetch_repo_audio, repo, int(samples_per_book), int(audio_sec), token, log): repo
             for repo in repos
         }
         done = 0
         for future in concurrent.futures.as_completed(futures):
-            repo, arrays, log = future.result()
+            repo, arrays = future.result()
             done += 1
             progress(0.05 + 0.55 * done / len(repos), desc=f"[{done}/{len(repos)}] {repo.split('/')[-1]}")
             if arrays:
                 repo_arrays[repo] = arrays
-                fetch_logs.append(f"{repo.split('/')[-1]}: {log}")
-            else:
-                errors.append(f"{repo.split('/')[-1]}: {log}")
+
+    ok = len(repo_arrays)
+    failed = len(repos) - ok
+    log.append(f"[{_ts()}] Phase 1 done: {ok} ok, {failed} failed, phase_total={int((time.time()-t_total)*1000)}ms")
 
     if not repo_arrays:
-        return pd.DataFrame(), "No audio downloaded.", "\n".join(errors)
+        return pd.DataFrame(), "No audio downloaded.", "\n".join(log)
 
     # ── Phase 2: batch embed all clips in one shot ───────────────────────────
+    log.append(f"[{_ts()}] --- Phase 2: batch embed ---")
     progress(0.60, desc="Batch embedding all clips…")
     all_arrays: list[np.ndarray] = []
     repo_slices: dict[str, tuple[int, int]] = {}
@@ -197,15 +237,17 @@ def identify_speakers(
             all_arrays.extend(repo_arrays[repo])
             repo_slices[repo] = (start, len(all_arrays))
 
+    log.append(f"[{_ts()}] embedding {len(all_arrays)} clips in batches of {BATCH_SIZE}…")
     t_embed = time.time()
-    all_embeddings = _batch_embed(all_arrays)  # (N_clips, D)
+    all_embeddings = _batch_embed(all_arrays)
     embed_ms = int((time.time() - t_embed) * 1000)
-    fetch_logs.append(
-        f"batch embed: {len(all_arrays)} clips in {embed_ms}ms "
-        f"({embed_ms // len(all_arrays)}ms/clip avg)"
+    log.append(
+        f"[{_ts()}] batch embed done: {len(all_arrays)} clips, "
+        f"{embed_ms}ms total, {embed_ms // max(len(all_arrays),1)}ms/clip avg"
     )
 
     # ── Phase 3: average per repo, cluster ───────────────────────────────────
+    log.append(f"[{_ts()}] --- Phase 3: cluster ({len(repo_slices)} repos) ---")
     progress(0.95, desc="Clustering…")
     embeddings: dict[str, np.ndarray] = {}
     for repo, (start, end) in repo_slices.items():
@@ -249,12 +291,11 @@ def identify_speakers(
 
     df = pd.DataFrame(rows).sort_values(["speaker_id", "dataset"]).reset_index(drop=True)
     n_speakers = len(set(labels))
-    summary = f"✅ {len(repo_names)} books → {n_speakers} unique speakers"
+    total_ms = int((time.time() - t_total) * 1000)
+    summary = f"✅ {len(repo_names)} books → {n_speakers} unique speakers  ({total_ms/1000:.1f}s total)"
+    log.append(f"[{_ts()}] === DONE: {total_ms}ms total ===")
 
-    debug = "\n".join(fetch_logs)
-    if errors:
-        debug += "\n\nERRORS:\n" + "\n".join(errors)
-    return df, summary, debug
+    return df, summary, "\n".join(log)
 
 
 DESCRIPTION = """