Upload app.py with huggingface_hub

app.py

@@ -17,6 +17,7 @@ from transformers import Wav2Vec2FeatureExtractor
 os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
 
 N_CPUS = os.cpu_count() or 2
+BATCH_SIZE = 64  # max clips per ONNX forward pass
 
 DATASETS_SERVER = "https://datasets-server.huggingface.co"
 ONNX_MODEL_ID = "fosters/wavlm-base-plus-sv-onnx"
@@ -32,26 +33,39 @@ def _load_model():
     with _init_lock:
         if _model is None:
             _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(ONNX_MODEL_ID)
-            # ONNX Runtime: thread-safe, no GIL concerns
             _model = ORTModelForAudioXVector.from_pretrained(ONNX_MODEL_ID)
     return _feature_extractor, _model
 
 
-def _embed(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
-    fe, mdl = _load_model()
+def _to_array(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
+    """Resample to TARGET_SR, convert to mono, trim to max_sec."""
     waveform = torch.tensor(audio_array, dtype=torch.float32)
     if waveform.ndim == 2:
         waveform = waveform.mean(0)
     if sr != TARGET_SR:
         waveform = torchaudio.functional.resample(waveform, sr, TARGET_SR)
-    waveform = waveform[: max_sec * TARGET_SR]
-    inputs = fe(waveform.numpy(), sampling_rate=TARGET_SR, return_tensors="pt")
-    out = mdl(**inputs)
-    return out.embeddings.squeeze().numpy()
+    return waveform[: max_sec * TARGET_SR].numpy().astype(np.float32)
+
+
+def _batch_embed(arrays: list[np.ndarray]) -> np.ndarray:
+    """Run all clips through the model in batches. Returns (N, D) embeddings."""
+    fe, mdl = _load_model()
+    all_embs = []
+    for i in range(0, len(arrays), BATCH_SIZE):
+        batch = arrays[i : i + BATCH_SIZE]
+        inputs = fe(batch, sampling_rate=TARGET_SR, return_tensors="pt", padding=True)
+        out = mdl(**inputs)
+        # embeddings shape: (batch, D)
+        embs = out.embeddings.detach().numpy()
+        if embs.ndim == 1:
+            embs = embs[np.newaxis]
+        all_embs.append(embs)
+    return np.concatenate(all_embs, axis=0)
+
 
+# ── Audio fetching ────────────────────────────────────────────────────────────
 
 def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> tuple[list[str], str]:
-    """Get direct audio URLs via datasets-server. Returns (urls, debug_info)."""
     headers = {"Authorization": f"Bearer {token}"} if token else {}
     resp = requests.get(
         f"{DATASETS_SERVER}/rows",
@@ -66,104 +80,70 @@ def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> tuple[list[str
     if not rows:
         return [], "datasets-server: empty rows"
 
-    # Inspect actual audio field structure for debugging
     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", {})
-        # audio can be a dict {"src": ...} or a list [{"src": ...}, ...]
         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"datasets-server: no src in audio field (keys={audio_keys})"
-    return urls, f"datasets-server ok (audio keys={audio_keys})"
+        return [], f"no src in audio (keys={audio_keys})"
+    return urls, "ok"
 
 
-def _download_one(url: str, token: str | None, max_sec: int) -> tuple[np.ndarray, dict]:
+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)."""
     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 = (time.time() - t0) * 1000
-
-    t1 = time.time()
     audio_array, sr = sf.read(io.BytesIO(resp.content))
-    decode_ms = (time.time() - t1) * 1000
-
-    t2 = time.time()
-    emb = _embed(audio_array, sr, max_sec)
-    embed_ms = (time.time() - t2) * 1000
-
-    return emb, {
-        "size_kb": len(resp.content) // 1024,
-        "dl_ms": int(dl_ms),
-        "decode_ms": int(decode_ms),
-        "embed_ms": int(embed_ms),
-    }
-
-
-def _streaming_fallback(repo: str, n_samples: int, audio_sec: int, token: str | None) -> list[np.ndarray]:
-    """Fallback: use datasets streaming when datasets-server is unavailable."""
-    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))
-    embs = []
-    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()
-        audio_array, sr = sf.read(io.BytesIO(audio_bytes))
-        embs.append(_embed(audio_array, sr, audio_sec))
-    return embs
-
-
-def _process_repo(
+    return _to_array(audio_array, sr, max_sec), len(resp.content) // 1024
+
+
+def _fetch_repo_audio(
     repo: str, n_samples: int, audio_sec: int, token: str | None
-) -> tuple[str, np.ndarray | None, str]:
+) -> tuple[str, list[np.ndarray], str]:
+    """Fetch + download audio for one repo. Returns (repo, arrays, log)."""
     t0 = time.time()
     try:
-        urls, api_debug = _fetch_audio_urls(repo, n_samples, token)
+        urls, api_status = _fetch_audio_urls(repo, n_samples, token)
         fetch_ms = int((time.time() - t0) * 1000)
 
         if urls:
-            # Fast path: parallel downloads via datasets-server URLs
             with concurrent.futures.ThreadPoolExecutor(max_workers=len(urls)) as ex:
-                futures = [ex.submit(_download_one, u, token, audio_sec) for u in urls]
-                results = [f.result() for f in futures]
-            embs = [r[0] for r in results]
-            s = results[0][1]
-            total_ms = int((time.time() - t0) * 1000)
-            log = (
-                f"{repo.split('/')[-1]}: "
-                f"api={fetch_ms}ms dl={s['dl_ms']}ms "
-                f"decode={s['decode_ms']}ms embed={s['embed_ms']}ms "
-                f"total={total_ms}ms ({s['size_kb']}KB/file) [{api_debug}]"
-            )
-        else:
-            # Slow fallback: streaming (downloads parquet shard)
-            t1 = time.time()
-            embs = _streaming_fallback(repo, n_samples, audio_sec, token)
-            total_ms = int((time.time() - t0) * 1000)
-            log = (
-                f"{repo.split('/')[-1]}: streaming fallback "
-                f"total={total_ms}ms [reason: {api_debug}]"
-            )
-
-        if not embs:
-            return repo, None, f"{repo}: no audio samples extracted"
-
-        return repo, np.mean(embs, axis=0), log
+                results = list(ex.map(lambda u: _download_audio(u, token, audio_sec), urls))
+            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
+
+        # Streaming fallback
+        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))
+        total_ms = int((time.time() - t0) * 1000)
+        return repo, arrays, f"streaming fallback total={total_ms}ms [reason: {api_status}]"
 
     except Exception as exc:
-        return repo, None, f"{repo}: ERROR {exc}"
+        return repo, [], f"ERROR: {exc}"
 
 
+# ── Main ──────────────────────────────────────────────────────────────────────
+
 def identify_speakers(
     repo_ids_text: str,
     samples_per_book: int,
@@ -180,31 +160,56 @@ def identify_speakers(
 
     progress(0, desc="Loading model…")
     _load_model()
-    progress(0.02, desc=f"Processing {len(repos)} datasets in parallel…")
 
-    embeddings: dict[str, np.ndarray] = {}
-    logs: list[str] = []
+    # ── Phase 1: download all audio in parallel (I/O bound) ──────────────────
+    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] = []
-    done = 0
 
-    # Process repos in parallel — capped at N_CPUS since embed is the bottleneck
-    with concurrent.futures.ThreadPoolExecutor(max_workers=N_CPUS) as ex:
-        future_to_repo = {
-            ex.submit(_process_repo, repo, int(samples_per_book), int(audio_sec), token): repo
+    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
             for repo in repos
         }
-        for future in concurrent.futures.as_completed(future_to_repo):
-            repo, emb, log = future.result()
+        done = 0
+        for future in concurrent.futures.as_completed(futures):
+            repo, arrays, log = future.result()
             done += 1
-            progress(done / len(repos), desc=f"[{done}/{len(repos)}] {repo.split('/')[-1]}")
-            if emb is not None:
-                embeddings[repo] = emb
-                logs.append(log)
+            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(log)
+                errors.append(f"{repo.split('/')[-1]}: {log}")
+
+    if not repo_arrays:
+        return pd.DataFrame(), "No audio downloaded.", "\n".join(errors)
+
+    # ── Phase 2: batch embed all clips in one shot ───────────────────────────
+    progress(0.60, desc="Batch embedding all clips…")
+    all_arrays: list[np.ndarray] = []
+    repo_slices: dict[str, tuple[int, int]] = {}
+
+    for repo in repos:
+        if repo in repo_arrays:
+            start = len(all_arrays)
+            all_arrays.extend(repo_arrays[repo])
+            repo_slices[repo] = (start, len(all_arrays))
+
+    t_embed = time.time()
+    all_embeddings = _batch_embed(all_arrays)  # (N_clips, D)
+    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)"
+    )
 
-    if not embeddings:
-        return pd.DataFrame(), "No embeddings extracted.", "\n".join(errors)
+    # ── Phase 3: average per repo, cluster ───────────────────────────────────
+    progress(0.95, desc="Clustering…")
+    embeddings: dict[str, np.ndarray] = {}
+    for repo, (start, end) in repo_slices.items():
+        embeddings[repo] = all_embeddings[start:end].mean(axis=0)
 
     repo_names = list(embeddings.keys())
     emb_matrix = np.stack([embeddings[r] for r in repo_names])
@@ -246,7 +251,7 @@ def identify_speakers(
     n_speakers = len(set(labels))
     summary = f"✅ {len(repo_names)} books → {n_speakers} unique speakers"
 
-    debug = "\n".join(logs)
+    debug = "\n".join(fetch_logs)
     if errors:
         debug += "\n\nERRORS:\n" + "\n".join(errors)
     return df, summary, debug
@@ -257,9 +262,9 @@ DESCRIPTION = """
 
 Finds unique speakers across multiple HF audio datasets. Each dataset is assumed to have
 **one speaker** (e.g. an audiobook). The app fetches audio directly via the datasets-server API
-(no full parquet download), processes datasets in parallel, and clusters by voice similarity.
+(no full parquet download), downloads in parallel, then embeds all clips in one batched forward pass.
 
-**Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) —
+**Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) (ONNX) —
 language-agnostic speaker embeddings, works for any language.
 
 ---
@@ -272,7 +277,7 @@ language-agnostic speaker embeddings, works for any language.
    - *Audio length (sec)* — seconds of each chunk to use for embedding. 5 sec is sufficient for a clear voice.
    - *Same-speaker threshold* — cosine similarity cutoff. Raise if too many books merge into one speaker; lower if one person gets split across IDs.
    - *HF Token* — only needed for **private** repos.
-3. Click **Identify Speakers**. Datasets are processed in parallel — expect ~30–60 sec for 35 books.
+3. Click **Identify Speakers**. Downloads run in parallel, then all clips are embedded in one batch — expect ~20–40 sec for 35 books.
 
 ## Output columns
 
@@ -286,7 +291,7 @@ language-agnostic speaker embeddings, works for any language.
 
 **Tip:** Sort by `speaker_id` to see all books by the same narrator grouped together.
 
-The **Errors / Timing** box shows per-dataset timing breakdown (API fetch / download / decode / embed) — useful for diagnosing slow datasets.
+The **Errors / Timing** box shows per-dataset timing and batch embed stats — useful for diagnosing slow datasets.
 """
 
 with gr.Blocks(title="Speaker Identifier") as demo: