diarization.py

"""Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.

Spectral clustering implementation adapted from FunASR/3D-Speaker:
https://github.com/alibaba-damo-academy/FunASR
MIT License (https://opensource.org/licenses/MIT)
"""

import warnings

import numpy as np
import scipy
import sklearn.metrics.pairwise
import torch
from sklearn.cluster._kmeans import k_means
from sklearn.preprocessing import normalize


def _get_device() -> torch.device:
    """Get best available device for inference."""
    if torch.cuda.is_available():
        return torch.device("cuda")
    if torch.backends.mps.is_available():
        return torch.device("mps")
    return torch.device("cpu")


class SpectralCluster:
    """Spectral clustering using unnormalized Laplacian of affinity matrix.

    Adapted from FunASR/3D-Speaker and SpeechBrain implementations.
    Uses eigenvalue gap to automatically determine number of speakers.
    """

    def __init__(self, min_num_spks: int = 1, max_num_spks: int = 15, pval: float = 0.06):
        self.min_num_spks = min_num_spks
        self.max_num_spks = max_num_spks
        self.pval = pval

    def __call__(self, embeddings: np.ndarray, oracle_num: int | None = None) -> np.ndarray:
        """Run spectral clustering on embeddings.

        Args:
            embeddings: Speaker embeddings of shape [N, D]
            oracle_num: Optional known number of speakers

        Returns:
            Cluster labels of shape [N]
        """
        # Similarity matrix computation
        sim_mat = self.get_sim_mat(embeddings)

        # Refining similarity matrix with pval
        prunned_sim_mat = self.p_pruning(sim_mat)

        # Symmetrization
        sym_prund_sim_mat = 0.5 * (prunned_sim_mat + prunned_sim_mat.T)

        # Laplacian calculation
        laplacian = self.get_laplacian(sym_prund_sim_mat)

        # Get Spectral Embeddings
        emb, num_of_spk = self.get_spec_embs(laplacian, oracle_num)

        # Perform clustering
        return self.cluster_embs(emb, num_of_spk)

    def get_sim_mat(self, embeddings: np.ndarray) -> np.ndarray:
        """Compute cosine similarity matrix."""
        return sklearn.metrics.pairwise.cosine_similarity(embeddings, embeddings)

    def p_pruning(self, affinity: np.ndarray) -> np.ndarray:
        """Prune low similarity values in affinity matrix (keep top pval fraction)."""
        n = affinity.shape[0]
        pval = max(self.pval, 6.0 / n)
        k_keep = max(1, int(pval * n))

        # Vectorized: find top-k indices per row and zero out the rest
        top_k_idx = np.argpartition(affinity, -k_keep, axis=1)[:, -k_keep:]
        mask = np.zeros_like(affinity, dtype=bool)
        np.put_along_axis(mask, top_k_idx, True, axis=1)
        affinity[~mask] = 0
        return affinity

    def get_laplacian(self, sim_mat: np.ndarray) -> np.ndarray:
        """Compute unnormalized Laplacian matrix."""
        from scipy.sparse.csgraph import laplacian

        np.fill_diagonal(sim_mat, 0)
        return laplacian(sim_mat, normed=False)

    def get_spec_embs(
        self, laplacian: np.ndarray, k_oracle: int | None = None
    ) -> tuple[np.ndarray, int]:
        """Extract spectral embeddings from Laplacian.

        Uses the eigengap heuristic to estimate the number of clusters:
        The number of clusters k is chosen where the gap between consecutive
        eigenvalues is largest, indicating a transition from "cluster" eigenvalues
        (near 0) to "noise" eigenvalues.
        """
        lambdas, eig_vecs = scipy.linalg.eigh(laplacian)

        num_of_spk = k_oracle if k_oracle is not None else self._estimate_num_speakers(lambdas)

        emb = eig_vecs[:, :num_of_spk]
        return emb, num_of_spk

    def _estimate_num_speakers(self, lambdas: np.ndarray) -> int:
        """Estimate number of speakers using refined eigengap heuristic.

        For spectral clustering, we look for the largest gap in eigenvalues.
        The eigenvalues corresponding to clusters are close to 0, and there
        should be a significant jump to the remaining eigenvalues.
        """
        # Consider eigenvalues from index 1 to max_num_spks (skip first, it's always ~0)
        # We need gaps between positions, so look at indices 1 to max_num_spks+1
        max_idx = min(self.max_num_spks + 1, len(lambdas))
        relevant_lambdas = lambdas[1:max_idx]  # Skip first eigenvalue

        if len(relevant_lambdas) < 2:
            return self.min_num_spks

        # Compute absolute gaps (not ratios - ratios are unstable near 0)
        gaps = np.diff(relevant_lambdas)

        # Find the largest gap - the index gives us (k-1) since we skipped first
        # Add 1 to convert from gap index to number of speakers
        # Add 1 again because we skipped the first eigenvalue
        max_gap_idx = int(np.argmax(gaps))
        num_of_spk = max_gap_idx + 2  # +1 for gap->count, +1 for skipped eigenvalue

        # Clamp between min and max
        return max(self.min_num_spks, min(num_of_spk, self.max_num_spks))

    def cluster_embs(self, emb: np.ndarray, k: int) -> np.ndarray:
        """Cluster spectral embeddings using k-means."""
        _, labels, _ = k_means(emb, k, n_init=10)
        return labels

    def get_eigen_gaps(self, eig_vals: np.ndarray) -> np.ndarray:
        """Compute gaps between consecutive eigenvalues."""
        return np.diff(eig_vals)


class SpeakerClusterer:
    """Speaker clustering backend using spectral clustering with speaker merging.

    Features:
    - Spectral clustering with eigenvalue gap for auto speaker count detection
    - P-pruning for affinity matrix refinement
    - Post-clustering speaker merging by cosine similarity
    """

    def __init__(
        self,
        min_num_spks: int = 2,
        max_num_spks: int = 10,
        merge_thr: float = 0.90,  # Moderate merging
    ):
        self.min_num_spks = min_num_spks
        self.max_num_spks = max_num_spks
        self.merge_thr = merge_thr
        self._spectral_cluster: SpectralCluster | None = None

    def _get_spectral_cluster(self) -> SpectralCluster:
        """Lazy-load spectral clusterer."""
        if self._spectral_cluster is None:
            self._spectral_cluster = SpectralCluster(
                min_num_spks=self.min_num_spks,
                max_num_spks=self.max_num_spks,
            )
        return self._spectral_cluster

    def __call__(self, embeddings: np.ndarray, num_speakers: int | None = None) -> np.ndarray:
        """Cluster speaker embeddings and return labels.

        Args:
            embeddings: Speaker embeddings of shape [N, D]
            num_speakers: Optional oracle number of speakers

        Returns:
            Cluster labels of shape [N]
        """
        import warnings

        if len(embeddings.shape) != 2:
            raise ValueError(f"Expected 2D array, got shape {embeddings.shape}")

        # Handle edge cases
        if embeddings.shape[0] == 0:
            return np.array([], dtype=int)
        if embeddings.shape[0] == 1:
            return np.array([0], dtype=int)
        if embeddings.shape[0] < 6:
            return np.zeros(embeddings.shape[0], dtype=int)

        # Normalize embeddings and replace NaN/inf
        embeddings = np.nan_to_num(embeddings, nan=0.0, posinf=0.0, neginf=0.0)
        embeddings = normalize(embeddings)

        # Run spectral clustering (suppress numerical warnings)
        spectral = self._get_spectral_cluster()

        # Update min/max for oracle case
        if num_speakers is not None:
            spectral.min_num_spks = num_speakers
            spectral.max_num_spks = num_speakers

        with warnings.catch_warnings():
            warnings.filterwarnings("ignore", category=RuntimeWarning)
            labels = spectral(embeddings, oracle_num=num_speakers)

        # Reset min/max
        if num_speakers is not None:
            spectral.min_num_spks = self.min_num_spks
            spectral.max_num_spks = self.max_num_spks

        # Merge similar speakers if no oracle
        if num_speakers is None:
            labels = self._merge_by_cos(labels, embeddings, self.merge_thr)

        # Re-index labels sequentially
        _, labels = np.unique(labels, return_inverse=True)

        return labels

    def _merge_by_cos(self, labels: np.ndarray, embs: np.ndarray, cos_thr: float) -> np.ndarray:
        """Merge similar speakers by cosine similarity of centroids."""
        from scipy.cluster.hierarchy import fcluster, linkage
        from scipy.spatial.distance import pdist

        unique_labels = np.unique(labels)
        if len(unique_labels) <= 1:
            return labels

        # Compute normalized speaker centroids
        centroids = np.array([embs[labels == lbl].mean(0) for lbl in unique_labels])
        centroids = normalize(centroids)

        # Hierarchical clustering with cosine distance
        distances = pdist(centroids, metric="cosine")
        linkage_matrix = linkage(distances, method="average")
        merged_labels = fcluster(linkage_matrix, t=1.0 - cos_thr, criterion="distance") - 1

        # Map original labels to merged labels
        label_map = dict(zip(unique_labels, merged_labels))
        return np.array([label_map[lbl] for lbl in labels])


class LocalSpeakerDiarizer:
    """Local speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.

    Pipeline:
    1. TEN-VAD detects speech segments
    2. Sliding window (1.0s, 75% overlap) for uniform embedding extraction
    3. ECAPA-TDNN extracts speaker embeddings per window
    4. Spectral clustering with eigenvalue gap for auto speaker detection
    5. Frame-level consensus voting for segment reconstruction
    6. Post-processing merges short segments to reduce flicker

    Tunable Parameters (class attributes):
    - WINDOW_SIZE: Embedding extraction window size in seconds
    - STEP_SIZE: Sliding window step size (overlap = WINDOW_SIZE - STEP_SIZE)
    - VAD_THRESHOLD: Speech detection threshold (lower = more sensitive)
    - VAD_MIN_DURATION: Minimum speech segment duration
    - VAD_MAX_GAP: Maximum gap to bridge between segments
    - VAD_PAD_ONSET/OFFSET: Padding added to speech segments
    - VOTING_RATE: Frame resolution for consensus voting
    - MIN_SEGMENT_DURATION: Minimum final segment duration
    - SAME_SPEAKER_GAP: Maximum gap to merge same-speaker segments
    - TAIL_COVERAGE_RATIO: Minimum tail coverage to add extra window
    """

    _ten_vad_model = None
    _ecapa_model = None
    _device = None

    # ==================== TUNABLE PARAMETERS ====================

    # Sliding window for embedding extraction
    WINDOW_SIZE = 0.75  # seconds - shorter window for finer resolution
    STEP_SIZE = 0.15  # seconds (80% overlap for more votes)
    TAIL_COVERAGE_RATIO = 0.1  # Add extra window if tail > this ratio of window

    # VAD hysteresis parameters
    VAD_THRESHOLD = 0.25  # Balanced threshold
    VAD_MIN_DURATION = 0.05  # Minimum speech segment duration (seconds)
    VAD_MAX_GAP = 0.50  # Bridge gaps shorter than this (seconds)
    VAD_PAD_ONSET = 0.05  # Padding at segment start (seconds)
    VAD_PAD_OFFSET = 0.05  # Padding at segment end (seconds)

    # Frame-level voting
    VOTING_RATE = 0.01  # 10ms resolution for consensus voting

    # Post-processing
    MIN_SEGMENT_DURATION = 0.15  # Minimum final segment duration (seconds)
    SHORT_SEGMENT_GAP = 0.1  # Gap threshold for merging short segments
    SAME_SPEAKER_GAP = 0.5  # Gap threshold for merging same-speaker segments

    # ===========================================================

    @classmethod
    def _get_ten_vad_model(cls):
        """Lazy-load TEN-VAD model (singleton)."""
        if cls._ten_vad_model is None:
            from ten_vad import TenVad

            cls._ten_vad_model = TenVad(hop_size=256, threshold=cls.VAD_THRESHOLD)
        return cls._ten_vad_model

    @classmethod
    def _get_device(cls) -> torch.device:
        """Get the best available device."""
        if cls._device is None:
            cls._device = _get_device()
        return cls._device

    @classmethod
    def _get_ecapa_model(cls):
        """Lazy-load ECAPA-TDNN speaker embedding model (singleton)."""
        if cls._ecapa_model is None:
            # Suppress torchaudio deprecation warning from SpeechBrain
            with warnings.catch_warnings():
                warnings.filterwarnings("ignore", message="torchaudio._backend")
                from speechbrain.inference.speaker import EncoderClassifier

                device = cls._get_device()
                cls._ecapa_model = EncoderClassifier.from_hparams(
                    source="speechbrain/spkrec-ecapa-voxceleb",
                    run_opts={"device": str(device)},
                )

        return cls._ecapa_model

    @classmethod
    def diarize(
        cls,
        audio: np.ndarray | str,
        sample_rate: int = 16000,
        num_speakers: int | None = None,
        min_speakers: int = 2,
        max_speakers: int = 10,
        **_kwargs,
    ) -> list[dict]:
        """Run speaker diarization on audio.

        Args:
            audio: Audio waveform as numpy array or path to audio file
            sample_rate: Audio sample rate (default 16000)
            num_speakers: Exact number of speakers (if known)
            min_speakers: Minimum number of speakers
            max_speakers: Maximum number of speakers

        Returns:
            List of dicts with 'speaker', 'start', 'end' keys
        """
        # Handle file path input
        if isinstance(audio, str):
            import librosa

            audio, sample_rate = librosa.load(audio, sr=16000)

        # Ensure correct sample rate
        if sample_rate != 16000:
            import librosa

            audio = librosa.resample(audio, orig_sr=sample_rate, target_sr=16000)
            sample_rate = 16000

        audio = audio.astype(np.float32)
        total_duration = len(audio) / sample_rate

        # Step 1: VAD (returns segments and raw frame-level decisions)
        segments, vad_frames = cls._get_speech_segments(audio, sample_rate)
        if not segments:
            return []

        # Step 2: Extract embeddings
        embeddings, window_segments = cls._extract_embeddings(audio, segments, sample_rate)
        if len(embeddings) == 0:
            return []

        # Step 3: Cluster
        clusterer = SpeakerClusterer(min_num_spks=min_speakers, max_num_spks=max_speakers)
        labels = clusterer(embeddings, num_speakers)

        # Step 4: Post-process with consensus voting (VAD-aware)
        return cls._postprocess_segments(window_segments, labels, total_duration, vad_frames)

    @classmethod
    def _get_speech_segments(
        cls, audio_array: np.ndarray, sample_rate: int = 16000
    ) -> tuple[list[dict], list[bool]]:
        """Get speech segments using TEN-VAD.

        Returns:
            Tuple of (segments list, vad_frames list of per-frame speech decisions)
        """
        vad_model = cls._get_ten_vad_model()

        # Convert to int16 as required by TEN-VAD
        # Clip to prevent integer overflow
        if audio_array.dtype != np.int16:
            audio_int16 = (np.clip(audio_array, -1.0, 1.0) * 32767).astype(np.int16)
        else:
            audio_int16 = audio_array

        # Process frame by frame
        hop_size = 256
        frame_duration = hop_size / sample_rate
        speech_frames: list[bool] = []

        for i in range(0, len(audio_int16) - hop_size, hop_size):
            frame = audio_int16[i : i + hop_size]
            _, is_speech = vad_model.process(frame)
            speech_frames.append(is_speech)

        # Convert frame-level decisions to segments
        segments = []
        in_speech = False
        start_idx = 0

        for i, is_speech in enumerate(speech_frames):
            if is_speech and not in_speech:
                start_idx = i
                in_speech = True
            elif not is_speech and in_speech:
                start_time = start_idx * frame_duration
                end_time = i * frame_duration
                segments.append(
                    {
                        "start": start_time,
                        "end": end_time,
                        "start_sample": int(start_time * sample_rate),
                        "end_sample": int(end_time * sample_rate),
                    }
                )
                in_speech = False

        # Handle trailing speech
        if in_speech:
            start_time = start_idx * frame_duration
            end_time = len(speech_frames) * frame_duration
            segments.append(
                {
                    "start": start_time,
                    "end": end_time,
                    "start_sample": int(start_time * sample_rate),
                    "end_sample": int(end_time * sample_rate),
                }
            )

        return cls._apply_vad_hysteresis(segments, sample_rate), speech_frames

    @classmethod
    def _apply_vad_hysteresis(cls, segments: list[dict], sample_rate: int = 16000) -> list[dict]:
        """Apply hysteresis-like post-processing to VAD segments."""
        if not segments:
            return segments

        segments = sorted(segments, key=lambda x: x["start"])

        # Fill short gaps
        merged = [segments[0].copy()]
        for seg in segments[1:]:
            gap = seg["start"] - merged[-1]["end"]
            if gap <= cls.VAD_MAX_GAP:
                merged[-1]["end"] = seg["end"]
                merged[-1]["end_sample"] = seg["end_sample"]
            else:
                merged.append(seg.copy())

        # Remove short segments
        filtered = [seg for seg in merged if (seg["end"] - seg["start"]) >= cls.VAD_MIN_DURATION]

        # Dilate segments (add padding)
        for seg in filtered:
            seg["start"] = max(0.0, seg["start"] - cls.VAD_PAD_ONSET)
            seg["end"] = seg["end"] + cls.VAD_PAD_OFFSET
            seg["start_sample"] = int(seg["start"] * sample_rate)
            seg["end_sample"] = int(seg["end"] * sample_rate)

        return filtered

    @classmethod
    def _extract_embeddings(
        cls, audio_array: np.ndarray, segments: list[dict], sample_rate: int
    ) -> tuple[np.ndarray, list[dict]]:
        """Extract speaker embeddings using sliding windows."""
        speaker_model = cls._get_ecapa_model()

        window_samples = int(cls.WINDOW_SIZE * sample_rate)
        step_samples = int(cls.STEP_SIZE * sample_rate)

        embeddings = []
        window_segments = []

        with torch.no_grad():
            for seg in segments:
                seg_start = seg["start_sample"]
                seg_end = seg["end_sample"]
                seg_len = seg_end - seg_start

                # Generate window positions
                if seg_len <= window_samples:
                    starts = [seg_start]
                    ends = [seg_end]
                else:
                    starts = list(range(seg_start, seg_end - window_samples + 1, step_samples))
                    ends = [s + window_samples for s in starts]

                    # Cover tail if > TAIL_COVERAGE_RATIO of window remains
                    if ends and ends[-1] < seg_end:
                        remainder = seg_end - ends[-1]
                        if remainder > (window_samples * cls.TAIL_COVERAGE_RATIO):
                            starts.append(seg_end - window_samples)
                            ends.append(seg_end)

                for c_start, c_end in zip(starts, ends):
                    chunk = audio_array[c_start:c_end]

                    # Pad short chunks with reflection
                    if len(chunk) < window_samples:
                        pad_width = window_samples - len(chunk)
                        chunk = np.pad(chunk, (0, pad_width), mode="reflect")

                    # Extract embedding using SpeechBrain's encode_batch
                    chunk_tensor = torch.from_numpy(chunk).float().unsqueeze(0)
                    embedding = (
                        speaker_model.encode_batch(chunk_tensor).squeeze(0).squeeze(0).cpu().numpy()
                    )

                    # Validate embedding
                    if np.isfinite(embedding).all() and np.linalg.norm(embedding) > 1e-8:
                        embeddings.append(embedding)
                        window_segments.append(
                            {
                                "start": c_start / sample_rate,
                                "end": c_end / sample_rate,
                            }
                        )

        # Normalize all embeddings at once
        if embeddings:
            return normalize(np.array(embeddings)), window_segments
        return np.array([]), []

    @classmethod
    def _resample_vad(cls, vad_frames: list[bool], num_frames: int) -> np.ndarray:
        """Resample VAD frame decisions to match voting grid resolution.

        VAD operates at 256 samples / 16000 Hz = 16ms per frame.
        Voting operates at VOTING_RATE (default 10ms) per frame.
        This maps VAD decisions to the finer voting grid.
        """
        if not vad_frames:
            return np.zeros(num_frames, dtype=bool)

        vad_rate = 256 / 16000  # 16ms per VAD frame
        vad_arr = np.array(vad_frames)

        # Vectorized: compute VAD frame indices for each voting frame
        voting_times = np.arange(num_frames) * cls.VOTING_RATE
        vad_indices = np.clip((voting_times / vad_rate).astype(int), 0, len(vad_arr) - 1)
        return vad_arr[vad_indices]

    @classmethod
    def _postprocess_segments(
        cls,
        window_segments: list[dict],
        labels: np.ndarray,
        total_duration: float,
        vad_frames: list[bool],
    ) -> list[dict]:
        """Post-process using frame-level consensus voting with VAD-aware silence."""
        if not window_segments or len(labels) == 0:
            return []

        # Correct labels to be contiguous
        unique_labels = np.unique(labels)
        label_map = {old: new for new, old in enumerate(unique_labels)}
        clean_labels = np.array([label_map[lbl] for lbl in labels])
        num_speakers = len(unique_labels)

        if num_speakers == 0:
            return []

        # Create voting grid
        num_frames = int(np.ceil(total_duration / cls.VOTING_RATE)) + 1
        votes = np.zeros((num_frames, num_speakers), dtype=np.float32)

        # Accumulate votes
        for win, label in zip(window_segments, clean_labels):
            start_frame = int(win["start"] / cls.VOTING_RATE)
            end_frame = int(win["end"] / cls.VOTING_RATE)
            end_frame = min(end_frame, num_frames)
            if start_frame < end_frame:
                votes[start_frame:end_frame, label] += 1.0

        # Determine winner per frame
        frame_speakers = np.argmax(votes, axis=1)
        max_votes = np.max(votes, axis=1)

        # Resample VAD to voting grid resolution for silence-aware voting
        vad_resampled = cls._resample_vad(vad_frames, num_frames)

        # Convert frames to segments
        final_segments = []
        current_speaker = -1
        seg_start = 0.0

        for f in range(num_frames):
            speaker = int(frame_speakers[f])
            score = max_votes[f]

            # Force silence if VAD says no speech OR no votes
            if score == 0 or not vad_resampled[f]:
                speaker = -1

            if speaker != current_speaker:
                if current_speaker != -1:
                    final_segments.append(
                        {
                            "speaker": f"SPEAKER_{current_speaker}",
                            "start": seg_start,
                            "end": f * cls.VOTING_RATE,
                        }
                    )
                current_speaker = speaker
                seg_start = f * cls.VOTING_RATE

        # Close last segment
        if current_speaker != -1:
            final_segments.append(
                {
                    "speaker": f"SPEAKER_{current_speaker}",
                    "start": seg_start,
                    "end": num_frames * cls.VOTING_RATE,
                }
            )

        return cls._merge_short_segments(final_segments)

    @classmethod
    def _merge_short_segments(cls, segments: list[dict]) -> list[dict]:
        """Merge short segments to reduce flicker."""
        if not segments:
            return []

        clean: list[dict] = []
        for seg in segments:
            dur = seg["end"] - seg["start"]
            if dur < cls.MIN_SEGMENT_DURATION:
                if (
                    clean
                    and clean[-1]["speaker"] == seg["speaker"]
                    and seg["start"] - clean[-1]["end"] < cls.SHORT_SEGMENT_GAP
                ):
                    clean[-1]["end"] = seg["end"]
                continue

            if (
                clean
                and clean[-1]["speaker"] == seg["speaker"]
                and seg["start"] - clean[-1]["end"] < cls.SAME_SPEAKER_GAP
            ):
                clean[-1]["end"] = seg["end"]
            else:
                clean.append(seg)

        return clean

    @classmethod
    def assign_speakers_to_words(
        cls,
        words: list[dict],
        speaker_segments: list[dict],
    ) -> list[dict]:
        """Assign speaker labels to words based on timestamp overlap.

        Args:
            words: List of word dicts with 'word', 'start', 'end' keys
            speaker_segments: List of speaker dicts with 'speaker', 'start', 'end' keys

        Returns:
            Words list with 'speaker' key added to each word
        """
        for word in words:
            word_mid = (word["start"] + word["end"]) / 2

            # Find the speaker segment that contains this word's midpoint
            best_speaker = None
            for seg in speaker_segments:
                if seg["start"] <= word_mid <= seg["end"]:
                    best_speaker = seg["speaker"]
                    break

            # If no exact match, find closest segment
            if best_speaker is None and speaker_segments:
                min_dist = float("inf")
                for seg in speaker_segments:
                    seg_mid = (seg["start"] + seg["end"]) / 2
                    dist = abs(word_mid - seg_mid)
                    if dist < min_dist:
                        min_dist = dist
                        best_speaker = seg["speaker"]

            word["speaker"] = best_speaker

        return words


class SpeakerDiarizer:
    """Speaker diarization using TEN-VAD + ECAPA-TDNN + spectral clustering.

    Example:
        >>> segments = SpeakerDiarizer.diarize(audio_array)
        >>> for seg in segments:
        ...     print(f"{seg['speaker']}: {seg['start']:.2f} - {seg['end']:.2f}")
    """

    @classmethod
    def diarize(
        cls,
        audio: np.ndarray | str,
        sample_rate: int = 16000,
        num_speakers: int | None = None,
        min_speakers: int | None = None,
        max_speakers: int | None = None,
        **_kwargs,
    ) -> list[dict]:
        """Run speaker diarization on audio.

        Args:
            audio: Audio waveform as numpy array or path to audio file
            sample_rate: Audio sample rate (default 16000)
            num_speakers: Exact number of speakers (if known)
            min_speakers: Minimum number of speakers
            max_speakers: Maximum number of speakers

        Returns:
            List of dicts with 'speaker', 'start', 'end' keys
        """
        return LocalSpeakerDiarizer.diarize(
            audio,
            sample_rate=sample_rate,
            num_speakers=num_speakers,
            min_speakers=min_speakers or 2,
            max_speakers=max_speakers or 10,
        )

    @classmethod
    def assign_speakers_to_words(
        cls,
        words: list[dict],
        speaker_segments: list[dict],
    ) -> list[dict]:
        """Assign speaker labels to words based on timestamp overlap."""
        return LocalSpeakerDiarizer.assign_speakers_to_words(words, speaker_segments)