Upload app.py with huggingface_hub

app.py

@@ -1,6 +1,5 @@
 import io
 import os
-import tempfile
 import time
 import datetime
 import threading
@@ -13,22 +12,25 @@ import soundfile as sf
 import torchaudio
 import torch
 from sklearn.cluster import AgglomerativeClustering
-from optimum.onnxruntime import ORTModelForAudioXVector
-from transformers import Wav2Vec2FeatureExtractor
+from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector
 
 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 = 12   # seconds per datasets-server attempt (one try per endpoint)
-STREAMING_TIMEOUT = 120  # seconds before giving up on streaming fallback
+N_WORKERS = max(1, N_CPUS // 2)  # 8 workers on 16-CPU Space
+EMBED_THREADS = 2                  # per-worker thread count — NUMA sweet spot
+API_TIMEOUT = 12
+STREAMING_TIMEOUT = 120
+
+# Set thread count before model load so ORT/MKL picks it up
+torch.set_num_threads(EMBED_THREADS)
 
 
 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"
+MODEL_ID = "microsoft/wavlm-base-plus-sv"
 TARGET_SR = 16000
 
 _feature_extractor = None
@@ -40,13 +42,12 @@ def _load_model():
     global _feature_extractor, _model
     with _init_lock:
         if _model is None:
-            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(ONNX_MODEL_ID)
-            _model = ORTModelForAudioXVector.from_pretrained(ONNX_MODEL_ID)
+            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
+            _model = WavLMForXVector.from_pretrained(MODEL_ID).eval()
     return _feature_extractor, _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)
@@ -55,20 +56,26 @@ def _to_array(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
     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."""
+def _embed_one(array: np.ndarray) -> np.ndarray:
+    """Embed a single clip. Thread-safe: eval()+no_grad(), GIL released in C++."""
     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)
+    inputs = fe(array, sampling_rate=TARGET_SR, return_tensors="pt")
+    with torch.no_grad():
         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)
+    return out.embeddings.squeeze().numpy()
+
+
+def _parallel_embed(arrays: list[np.ndarray], log: list[str]) -> np.ndarray:
+    """Embed all clips using N_WORKERS parallel threads (threads=2 each)."""
+    t0 = time.time()
+    log.append(f"[{_ts()}] --- Phase 2: embed {len(arrays)} clips "
+               f"({N_WORKERS} workers × {EMBED_THREADS} threads) ---")
+    with concurrent.futures.ThreadPoolExecutor(max_workers=N_WORKERS) as ex:
+        futures = [ex.submit(_embed_one, arr) for arr in arrays]
+        results = [f.result() for f in futures]
+    ms = int((time.time() - t0) * 1000)
+    log.append(f"[{_ts()}] embed done: {ms}ms total, {ms//len(arrays)}ms/clip avg")
+    return np.stack(results)
 
 
 # ── Audio fetching ────────────────────────────────────────────────────────────
@@ -261,9 +268,8 @@ def identify_speakers(
     if not repo_arrays:
         return pd.DataFrame(), "No audio downloaded.", "\n".join(log), "", None
 
-    # ── 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…")
+    # ── Phase 2: parallel embed (N_WORKERS workers × EMBED_THREADS each) ───────
+    progress(0.60, desc="Embedding clips…")
     all_arrays: list[np.ndarray] = []
     repo_slices: dict[str, tuple[int, int]] = {}
 
@@ -273,14 +279,7 @@ 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)
-    embed_ms = int((time.time() - t_embed) * 1000)
-    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"
-    )
+    all_embeddings = _parallel_embed(all_arrays, log)
 
     # ── Phase 3: average per repo, cluster ───────────────────────────────────
     log.append(f"[{_ts()}] --- Phase 3: cluster ({len(repo_slices)} repos) ---")
@@ -352,9 +351,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), downloads in parallel, then embeds all clips in one batched forward pass.
+(no full parquet download), downloads in parallel, then embeds clips across 8 parallel workers.
 
-**Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) (ONNX) —
+**Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) —
 language-agnostic speaker embeddings, works for any language.
 
 ---