Upload app.py with huggingface_hub

app.py

@@ -0,0 +1,167 @@
+import os
+import gradio as gr
+import numpy as np
+import pandas as pd
+import torch
+import torchaudio
+from datasets import load_dataset
+from sklearn.cluster import AgglomerativeClustering
+from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector
+
+MODEL_ID = "microsoft/wavlm-base-plus-sv"
+TARGET_SR = 16000
+MAX_AUDIO_SEC = 20  # trim long clips to keep inference fast
+
+_feature_extractor = None
+_model = None
+
+
+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()
+    return _feature_extractor, _model
+
+
+def _embed(audio_array: np.ndarray, sr: int) -> np.ndarray:
+    fe, mdl = _load_model()
+    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_AUDIO_SEC * TARGET_SR]
+    inputs = fe(waveform.numpy(), sampling_rate=TARGET_SR, return_tensors="pt")
+    with torch.no_grad():
+        out = mdl(**inputs)
+    return out.embeddings.squeeze().numpy()
+
+
+def identify_speakers(
+    repo_ids_text: str,
+    samples_per_book: int,
+    threshold: float,
+    hf_token: str,
+    progress=gr.Progress(),
+):
+    repos = [r.strip() for r in repo_ids_text.strip().splitlines() if r.strip()]
+    if not repos:
+        return pd.DataFrame(), "No repos provided.", ""
+
+    token = hf_token.strip() or os.environ.get("HF_TOKEN") or None
+
+    progress(0, desc="Loading model…")
+    _load_model()
+
+    embeddings: dict[str, np.ndarray] = {}
+    errors: list[str] = []
+
+    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)
+            embs = []
+            for j, row in enumerate(ds):
+                if j >= int(samples_per_book):
+                    break
+                audio = row["audio"]
+                embs.append(_embed(np.array(audio["array"]), audio["sampling_rate"]))
+            if embs:
+                embeddings[repo] = np.mean(embs, axis=0)
+            else:
+                errors.append(f"{repo}: no rows found")
+        except Exception as exc:
+            errors.append(f"{repo}: {exc}")
+
+    if not embeddings:
+        return pd.DataFrame(), "No embeddings extracted.", "\n".join(errors)
+
+    repo_names = list(embeddings.keys())
+    emb_matrix = np.stack([embeddings[r] for r in repo_names])
+    emb_matrix = emb_matrix / np.linalg.norm(emb_matrix, axis=1, keepdims=True)
+    sim_matrix = np.clip(emb_matrix @ emb_matrix.T, -1.0, 1.0)
+    dist_matrix = 1.0 - sim_matrix
+    np.fill_diagonal(dist_matrix, 0.0)
+
+    n = len(repo_names)
+    if n == 1:
+        labels = [0]
+    else:
+        labels = AgglomerativeClustering(
+            n_clusters=None,
+            distance_threshold=1.0 - float(threshold),
+            metric="precomputed",
+            linkage="average",
+        ).fit_predict(dist_matrix).tolist()
+
+    rows = []
+    for i, repo in enumerate(repo_names):
+        cluster = labels[i]
+        same_idx = [j for j, l in enumerate(labels) if l == cluster and j != i]
+        intra_sim = float(np.mean([sim_matrix[i][j] for j in same_idx])) if same_idx else 1.0
+        other_sorted = sorted([j for j in range(n) if j != i], key=lambda j: -sim_matrix[i][j])
+        closest = (
+            f"{repo_names[other_sorted[0]].split('/')[-1]} ({sim_matrix[i][other_sorted[0]]:.2f})"
+            if other_sorted else "-"
+        )
+        rows.append({
+            "dataset": repo.split("/")[-1],
+            "speaker_id": f"speaker_{cluster + 1:02d}",
+            "books_with_speaker": sum(1 for l in labels if l == cluster),
+            "intra_sim": round(intra_sim, 3),
+            "closest_match": closest,
+        })
+
+    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"
+
+
+with gr.Blocks(title="Speaker Identifier") as demo:
+    gr.Markdown(
+        """
+# 🎙️ Speaker Identifier
+Extract speaker embeddings from HF audio datasets and cluster into unique speakers.
+Uses **WavLM-Base+** (`microsoft/wavlm-base-plus-sv`) — language-agnostic, works for any language.
+"""
+    )
+    with gr.Row():
+        with gr.Column(scale=2):
+            repo_input = gr.Textbox(
+                label="Dataset repo IDs (one per line, owner/name)",
+                placeholder="fosters/some_audiobook_output\nfosters/another_audiobook_output",
+                lines=14,
+            )
+        with gr.Column(scale=1):
+            samples = gr.Slider(1, 10, value=3, step=1, label="Samples per book")
+            threshold = gr.Slider(
+                0.60, 0.98, value=0.82, step=0.01,
+                label="Same-speaker threshold (cosine similarity)",
+                info="Higher = stricter matching, more clusters",
+            )
+            hf_token = gr.Textbox(
+                label="HF Token (only for private repos)",
+                type="password",
+                placeholder="hf_…",
+            )
+            run_btn = gr.Button("Identify Speakers", variant="primary")
+
+    summary_out = gr.Textbox(label="Summary", interactive=False)
+    table_out = gr.Dataframe(
+        label="Results — sorted by speaker_id",
+        headers=["dataset", "speaker_id", "books_with_speaker", "intra_sim", "closest_match"],
+        wrap=True,
+    )
+    errors_out = gr.Textbox(label="Errors / Skipped", interactive=False)
+
+    run_btn.click(
+        identify_speakers,
+        inputs=[repo_input, samples, threshold, hf_token],
+        outputs=[table_out, summary_out, errors_out],
+    )
+
+demo.launch()