app.py

import io
import os
import tempfile
import time
import datetime
import threading
import concurrent.futures

import gradio as gr
import numpy as np
import pandas as pd
import requests
import soundfile as sf
import torchaudio
import torch
from sklearn.cluster import AgglomerativeClustering
from transformers import Wav2Vec2FeatureExtractor, WavLMForXVector

os.environ.setdefault("HF_XET_HIGH_PERFORMANCE", "1")

N_CPUS = os.cpu_count() or 2
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

MODEL_ID = "microsoft/wavlm-base-plus-sv"
TARGET_SR = 16000
MIN_AUDIO_SEC = 1                  # WavLM must not receive empty/tiny clips
MIN_AUDIO_SAMPLES = TARGET_SR * MIN_AUDIO_SEC

# 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"

_feature_extractor = None
_model = None
_init_lock = threading.Lock()


def _load_model():
    global _feature_extractor, _model
    with _init_lock:
        if _model is None:
            _feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(MODEL_ID)
            _model = WavLMForXVector.from_pretrained(MODEL_ID).eval()
    return _feature_extractor, _model


def _mono_1d_tensor(audio_array: np.ndarray) -> torch.Tensor:
    """Convert audio from soundfile/torchaudio style arrays to mono 1-D float tensor.

    soundfile normally returns:
      - mono:   shape (frames,)
      - stereo: shape (frames, channels)

    torchaudio often uses:
      - shape (channels, frames)

    The old bug was using mean(0) for soundfile stereo. For shape
    (frames, channels), mean(0) collapses frames and leaves only 1-2 samples,
    which later crashes WavLM conv1d with: kernel size > input size.
    """
    arr = np.asarray(audio_array)

    if arr.size == 0:
        return torch.zeros(0, dtype=torch.float32)

    # Convert integers to float32 in a safe range when necessary.
    if np.issubdtype(arr.dtype, np.integer):
        info = np.iinfo(arr.dtype)
        scale = float(max(abs(info.min), info.max))
        arr = arr.astype(np.float32) / scale
    else:
        arr = arr.astype(np.float32, copy=False)

    arr = np.nan_to_num(arr, nan=0.0, posinf=0.0, neginf=0.0)
    waveform = torch.from_numpy(arr)

    if waveform.ndim == 1:
        return waveform.contiguous()

    if waveform.ndim == 2:
        # Heuristic:
        # - soundfile: (frames, channels), usually channels <= 8 and frames >> channels
        # - torchaudio: (channels, frames), usually first dim <= 8
        if waveform.shape[0] <= 8 and waveform.shape[1] > waveform.shape[0]:
            # channels-first -> average channels
            return waveform.mean(dim=0).contiguous()
        # frames-first -> average channels; this is the important fix
        return waveform.mean(dim=1).contiguous()

    # Fallback for unusual shapes: preserve time axis as the longest axis and
    # average everything else.
    time_axis = int(np.argmax(waveform.shape))
    waveform = waveform.movedim(time_axis, 0)
    waveform = waveform.reshape(waveform.shape[0], -1).mean(dim=1)
    return waveform.contiguous()


def _to_array(audio_array: np.ndarray, sr: int, max_sec: int) -> np.ndarray:
    waveform = _mono_1d_tensor(audio_array)

    if waveform.numel() == 0:
        waveform = torch.zeros(MIN_AUDIO_SAMPLES, dtype=torch.float32)

    if int(sr) <= 0:
        sr = TARGET_SR

    if sr != TARGET_SR:
        waveform = torchaudio.functional.resample(waveform, int(sr), TARGET_SR)

    max_samples = max(MIN_AUDIO_SAMPLES, int(max_sec) * TARGET_SR)
    waveform = waveform[:max_samples]

    # WavLM cannot process empty/tiny clips. Pad to at least 1 second.
    if waveform.numel() < MIN_AUDIO_SAMPLES:
        pad = MIN_AUDIO_SAMPLES - waveform.numel()
        waveform = torch.nn.functional.pad(waveform, (0, pad))

    return waveform.numpy().astype(np.float32, copy=False)


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()

    array = np.asarray(array, dtype=np.float32).reshape(-1)
    array = np.nan_to_num(array, nan=0.0, posinf=0.0, neginf=0.0)

    # Last-resort guard in case a bad array bypassed _to_array().
    if array.size < MIN_AUDIO_SAMPLES:
        array = np.pad(array, (0, MIN_AUDIO_SAMPLES - array.size), mode="constant")

    inputs = fe(array, sampling_rate=TARGET_SR, return_tensors="pt")
    with torch.no_grad():
        out = mdl(**inputs)
    return out.embeddings.squeeze().detach().cpu().numpy().astype(np.float32, copy=False)


def _parallel_embed(arrays: list[np.ndarray], log: list[str]) -> np.ndarray:
    """Embed all clips using N_WORKERS parallel threads (threads=2 each)."""
    if not arrays:
        raise ValueError("No audio arrays to embed.")

    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)
    avg = ms // max(1, len(arrays))
    log.append(f"[{_ts()}] embed done: {ms}ms total, {avg}ms/clip avg")
    return np.stack(results)


# ── Audio fetching ────────────────────────────────────────────────────────────

def _parse_audio_urls(rows: list) -> list[str]:
    urls = []
    for row in rows:
        audio = row.get("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"])
    return urls


def _try_endpoint(endpoint: str, params: dict, headers: dict) -> tuple[list[str], str]:
    """Single request attempt. Returns (urls, status_str)."""
    try:
        t0 = time.time()
        resp = requests.get(
            f"{DATASETS_SERVER}{endpoint}",
            params=params,
            headers=headers,
            timeout=API_TIMEOUT,
        )
        elapsed = int((time.time() - t0) * 1000)

        if not resp.ok:
            return [], f"{endpoint} HTTP {resp.status_code} ({elapsed}ms)"

        rows = resp.json().get("rows", [])
        if not rows:
            return [], f"{endpoint} empty ({elapsed}ms)"

        urls = _parse_audio_urls(rows)
        if urls:
            return urls, f"{endpoint} {elapsed}ms"

        return [], f"{endpoint} no src ({elapsed}ms)"
    except requests.Timeout:
        return [], f"{endpoint} timeout>{API_TIMEOUT}s"
    except Exception as exc:
        return [], f"{endpoint} {exc}"


def _fetch_audio_urls(repo_id: str, n: int, token: str | None) -> tuple[list[str], str]:
    """Fire /rows and /first-rows in parallel, return first successful result.

    Uses shutdown(wait=False) so the losing request doesn't block the caller.
    """
    headers = {"Authorization": f"Bearer {token}"} if token else {}
    calls = [
        (
            "/rows",
            {
                "dataset": repo_id,
                "config": "default",
                "split": "train",
                "offset": 0,
                "length": n,
            },
        ),
        (
            "/first-rows",
            {
                "dataset": repo_id,
                "config": "default",
                "split": "train",
            },
        ),
    ]

    ex = concurrent.futures.ThreadPoolExecutor(max_workers=2)
    futs = {ex.submit(_try_endpoint, ep, params, headers): ep for ep, params in calls}
    errs = []

    try:
        for fut in concurrent.futures.as_completed(futs, timeout=API_TIMEOUT + 2):
            urls, status = fut.result()
            if urls:
                ex.shutdown(wait=False, cancel_futures=True)
                return urls[:n], status
            errs.append(status)
    except concurrent.futures.TimeoutError:
        errs.append("both endpoints timed out")
    finally:
        ex.shutdown(wait=False, cancel_futures=True)

    return [], " | ".join(errs)


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), always_2d=False)
    array = _to_array(audio_array, sr, max_sec)

    return array, len(resp.content) // 1024, dl_ms


def _fetch_repo_audio(
    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, int(n_samples), token)

        if urls:
            log.append(f"[{_ts()}] {short}: {api_status} → {len(urls)} URLs")
            with concurrent.futures.ThreadPoolExecutor(max_workers=max(1, len(urls))) as ex:
                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 if r[0].size >= MIN_AUDIO_SAMPLES]
            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 with timeout
        log.append(
            f"[{_ts()}] {short}: ⚠ datasets-server failed ({api_status}) "
            f"→ streaming fallback (timeout={STREAMING_TIMEOUT}s)"
        )
        t1 = time.time()

        def _stream() -> list[np.ndarray]:
            from datasets import Audio as HFAudio
            from datasets import load_dataset

            ds = load_dataset(repo, split="train", streaming=True, token=token)
            ds = ds.cast_column("audio", HFAudio(decode=False))

            result = []
            for j, row in enumerate(ds):
                if j >= int(n_samples):
                    break

                raw = row["audio"]
                if raw.get("bytes") is not None:
                    audio_bytes = raw["bytes"]
                else:
                    with open(raw["path"], "rb") as fh:
                        audio_bytes = fh.read()

                a, sr = sf.read(io.BytesIO(audio_bytes), always_2d=False)
                result.append(_to_array(a, sr, audio_sec))

                elapsed = int((time.time() - t1) * 1000)
                log.append(
                    f"[{_ts()}] {short}: streaming clip {j + 1}/{n_samples} "
                    f"({len(audio_bytes) // 1024}KB, {elapsed}ms so far)"
                )

            return result

        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as ex:
            fut = ex.submit(_stream)
            try:
                arrays = fut.result(timeout=STREAMING_TIMEOUT)
                total_ms = int((time.time() - t0) * 1000)
                log.append(
                    f"[{_ts()}] {short}: ✓ streaming done, "
                    f"{len(arrays)} clips, total={total_ms}ms"
                )
                return repo, arrays
            except concurrent.futures.TimeoutError:
                total_ms = int((time.time() - t0) * 1000)
                log.append(
                    f"[{_ts()}] {short}: ✗ streaming timeout after "
                    f"{STREAMING_TIMEOUT}s — skipping (total={total_ms}ms)"
                )
                return repo, []

    except Exception as exc:
        log.append(f"[{_ts()}] {short}: ✗ {exc} (total={int((time.time() - t0) * 1000)}ms)")
        return repo, []


# ── Main ──────────────────────────────────────────────────────────────────────

def identify_speakers(
    repo_ids_text: str,
    samples_per_book: int,
    audio_sec: 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.", "", "", None

    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, "
        f"{samples_per_book} sample/book, {audio_sec}s/clip ==="
    )
    log.append(f"[{_ts()}] CPU count: {N_CPUS}, workers: {N_WORKERS}, embed_threads: {EMBED_THREADS}")

    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]] = {}

    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, log): repo
            for repo in repos
        }
        done = 0
        for future in concurrent.futures.as_completed(futures):
            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
            else:
                log.append(f"[{_ts()}] {repo.split('/')[-1]}: no usable audio clips after loading")

    ok = len(repo_arrays)
    failed = len(repos) - ok
    log.append(
        f"[{_ts()}] Phase 1 done: {ok} ok, {failed} failed, "
        f"phase_total={int((time.time() - t_total) * 1000)}ms"
    )

    if not repo_arrays:
        return pd.DataFrame(), "No audio downloaded.", "\n".join(log), "", None

    # ── 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]] = {}

    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))

    try:
        all_embeddings = _parallel_embed(all_arrays, log)
    except Exception as exc:
        log.append(f"[{_ts()}] ✗ embedding failed: {exc}")
        return pd.DataFrame(), f"Embedding failed: {exc}", "\n".join(log), "", None

    # ── 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():
        if end > start:
            embeddings[repo] = all_embeddings[start:end].mean(axis=0)

    if not embeddings:
        return pd.DataFrame(), "No embeddings created.", "\n".join(log), "", None

    repo_names = list(embeddings.keys())
    emb_matrix = np.stack([embeddings[r] for r in repo_names])

    norms = np.linalg.norm(emb_matrix, axis=1, keepdims=True)
    norms = np.where(norms == 0, 1.0, norms)
    emb_matrix = emb_matrix / norms

    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, label in enumerate(labels) if label == 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 label in labels if label == 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))
    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 ===")

    # Plain-text copy-friendly output (space-separated, matches log format)
    text_lines = []
    for _, r in df.iterrows():
        text_lines.append(
            f"{r['dataset']}  {r['speaker_id']}  {r['books_with_speaker']}  "
            f"{r['intra_sim']}  {r['closest_match']}"
        )
    plain_text = "\n".join(text_lines)

    # CSV file for download
    tmp = tempfile.NamedTemporaryFile(mode="w", suffix=".csv", delete=False, encoding="utf-8")
    df.to_csv(tmp.name, index=False)
    tmp.close()

    return df, summary, "\n".join(log), plain_text, tmp.name


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 fetches audio directly via the datasets-server API
(no full parquet download), downloads in parallel, then embeds clips across parallel workers.

**Model:** [microsoft/wavlm-base-plus-sv](https://huggingface.co/microsoft/wavlm-base-plus-sv) —
language-agnostic speaker embeddings, works for any language.

---

## How to use

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* — 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**. Downloads run in parallel, then clips are embedded in parallel.

## Output columns

| Column | Meaning |
|---|---|
| `dataset` | Repo name (short) |
| `speaker_id` | Cluster label — same ID = same voice |
| `books_with_speaker` | How many books share this speaker |
| `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 and batch embed stats — useful for diagnosing slow datasets.
"""


with gr.Blocks(title="Speaker Identifier") as demo:
    gr.Markdown(DESCRIPTION)
    gr.Markdown("---")

    with gr.Row():
        with gr.Column(scale=2):
            repo_input = gr.Textbox(
                label="Dataset repo IDs (one per line)",
                placeholder="fosters/some_audiobook_output\nfosters/another_audiobook_output",
                lines=16,
            )
        with gr.Column(scale=1):
            samples = gr.Slider(1, 10, value=3, step=1, label="Samples per book")
            audio_sec = gr.Slider(
                2,
                30,
                value=5,
                step=1,
                label="Audio length per sample (sec)",
                info="5 sec is usually enough; longer = more accurate but slower",
            )
            threshold = gr.Slider(
                0.60,
                0.98,
                value=0.82,
                step=0.01,
                label="Same-speaker threshold",
                info="Higher = stricter matching → more clusters",
            )
            hf_token = gr.Textbox(
                label="HF Token (private repos only)",
                type="password",
                placeholder="hf_…",
            )
            run_btn = gr.Button("Identify Speakers", variant="primary", size="lg")

    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,
    )

    with gr.Row():
        text_out = gr.Textbox(
            label="Plain text (copy-friendly)",
            interactive=False,
            lines=12,
            info="dataset  speaker_id  n_books  intra_sim  closest_match",
        )
        csv_out = gr.File(label="Download CSV", file_types=[".csv"])

    errors_out = gr.Textbox(label="Errors / Timing", interactive=False)

    run_btn.click(
        identify_speakers,
        inputs=[repo_input, samples, audio_sec, threshold, hf_token],
        outputs=[table_out, summary_out, errors_out, text_out, csv_out],
    )


demo.launch()