Upload app.py with huggingface_hub

app.py

@@ -1,28 +1,36 @@
 import io
 import os
+import time
+import threading
+import concurrent.futures
 import gradio as gr
 import numpy as np
 import pandas as pd
+import requests
 import soundfile as sf
 import torch
 import torchaudio
-from datasets import load_dataset, Audio
 from sklearn.cluster import AgglomerativeClustering
 from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector
 
+os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")
+
+DATASETS_SERVER = "https://datasets-server.huggingface.co"
 MODEL_ID = "microsoft/wavlm-base-plus-sv"
 TARGET_SR = 16000
 
 _feature_extractor = None
 _model = None
+_embed_lock = threading.Lock()  # WavLM inference is not thread-safe
 
 
 def _load_model():
     global _feature_extractor, _model
-    if _model is None:
-        _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
-        _model = WavLMForXVector.from_pretrained(MODEL_ID)
-        _model.eval()
+    with _embed_lock:
+        if _model is None:
+            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
+            _model = WavLMForXVector.from_pretrained(MODEL_ID)
+            _model.eval()
     return _feature_extractor, _model
 
 
@@ -35,11 +43,84 @@ def _embed(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
         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")
-    with torch.no_grad():
-        out = mdl(**inputs)
+    with _embed_lock:
+        with torch.no_grad():
+            out = mdl(**inputs)
     return out.embeddings.squeeze().numpy()
 
 
+def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> list[str]:
+    """Get direct audio URLs via datasets-server — no parquet download needed."""
+    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,
+    )
+    resp.raise_for_status()
+    urls = []
+    for row in resp.json().get("rows", []):
+        audio = row["row"].get("audio", {})
+        if isinstance(audio, dict) and "src" in audio:
+            urls.append(audio["src"])
+    return urls
+
+
+def _download_one(url: str, token: str | None, max_sec: int) -> tuple[np.ndarray, dict]:
+    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 _process_repo(
+    repo: str, n_samples: int, audio_sec: int, token: str | None
+) -> tuple[str, np.ndarray | None, str]:
+    t0 = time.time()
+    try:
+        urls = _fetch_audio_urls(repo, n_samples, token)
+        if not urls:
+            return repo, None, f"{repo}: no audio URLs from datasets-server"
+        fetch_ms = int((time.time() - t0) * 1000)
+
+        # Download audio files in parallel (I/O bound)
+        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]
+        # Log stats from first sample as representative
+        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)"
+        )
+        return repo, np.mean(embs, axis=0), log
+
+    except Exception as exc:
+        return repo, None, f"{repo}: ERROR {exc}"
+
+
 def identify_speakers(
     repo_ids_text: str,
     samples_per_book: int,
@@ -56,30 +137,28 @@ 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] = []
     errors: list[str] = []
+    done = 0
 
-    for i, repo in enumerate(repos):
-        short = repo.split("/")[-1]
-        progress((i + 0.5) / len(repos), desc=f"[{i+1}/{len(repos)}] {short}")
-        try:
-            ds = load_dataset(repo, split="train", streaming=True, token=token)
-            ds = ds.cast_column("audio", Audio(decode=False))
-            embs = []
-            for j, row in enumerate(ds):
-                if j >= int(samples_per_book):
-                    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, int(audio_sec)))
-            if embs:
-                embeddings[repo] = np.mean(embs, axis=0)
+    # Process all repos in parallel (downloads are I/O bound)
+    with concurrent.futures.ThreadPoolExecutor(max_workers=8) as ex:
+        future_to_repo = {
+            ex.submit(_process_repo, 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 += 1
+            progress(done / len(repos), desc=f"[{done}/{len(repos)}] {repo.split('/')[-1]}")
+            if emb is not None:
+                embeddings[repo] = emb
+                logs.append(log)
             else:
-                errors.append(f"{repo}: no rows found")
-        except Exception as exc:
-            errors.append(f"{repo}: {exc}")
+                errors.append(log)
 
     if not embeddings:
         return pd.DataFrame(), "No embeddings extracted.", "\n".join(errors)
@@ -123,15 +202,19 @@ 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"
-    return df, summary, "\n".join(errors) if errors else "None"
+
+    debug = "\n".join(logs)
+    if errors:
+        debug += "\n\nERRORS:\n" + "\n".join(errors)
+    return df, summary, debug
 
 
 DESCRIPTION = """
 # 🎙️ Speaker Identifier
 
 Finds unique speakers across multiple HF audio datasets. Each dataset is assumed to have
-**one speaker** (e.g. an audiobook). The app extracts voice embeddings and clusters datasets
-by voice similarity.
+**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.
 
 **Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) —
 language-agnostic speaker embeddings, works for any language.
@@ -142,23 +225,25 @@ language-agnostic speaker embeddings, works for any language.
 
 1. **Paste dataset repo IDs** (one per line, `owner/name` format) into the left box.
 2. **Adjust parameters** if needed (defaults work well for audiobooks):
-   - *Samples per book* — how many audio chunks to average per dataset. More = more robust, slower.
-   - *Audio length (sec)* — how many seconds of each chunk to use. 5 sec is enough for a clear voice; increase if results look wrong.
-   - *Same-speaker threshold* — cosine similarity cutoff. Raise it if too many books are merged into one speaker; lower it if one person gets split into multiple IDs.
+   - *Samples per book* — audio chunks to average per dataset. More = more robust, slower. 3 is usually enough.
+   - *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**.
+3. Click **Identify Speakers**. Datasets are processed in parallel — expect ~30–60 sec for 35 books.
 
 ## Output columns
 
 | Column | Meaning |
 |---|---|
 | `dataset` | Repo name (short) |
-| `speaker_id` | Cluster label — same ID means same voice |
+| `speaker_id` | Cluster label — same ID = same voice |
 | `books_with_speaker` | How many books share this speaker |
-| `intra_sim` | Avg cosine similarity between books with the same speaker (1.0 = only one book; lower = cluster is less tight) |
-| `closest_match` | Most similar other book and its similarity score |
+| `intra_sim` | Avg cosine similarity within cluster (1.0 = only one book; lower = cluster is less tight) |
+| `closest_match` | Most similar other book and similarity score |
 
 **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.
 """
 
 with gr.Blocks(title="Speaker Identifier") as demo:
@@ -196,7 +281,7 @@ with gr.Blocks(title="Speaker Identifier") as demo:
         headers=["dataset", "speaker_id", "books_with_speaker", "intra_sim", "closest_match"],
         wrap=True,
     )
-    errors_out = gr.Textbox(label="Errors / Skipped", interactive=False)
+    errors_out = gr.Textbox(label="Errors / Timing", interactive=False)
 
     run_btn.click(
         identify_speakers,