tools/pyannote_cluster.py

import json
import argparse
import torch
import numpy as np
from pathlib import Path
from tqdm import tqdm
from pyannote.audio import Model, Inference
from sklearn.cluster import AgglomerativeClustering
# import umap.umap_ as umap # Optional for visualizing, skipping for speed

"""
tools/cluster_with_pyannote.py

WHAT THIS SCRIPT DOES:
----------------------
1. LOADS DATA: Reads your 'train_manifest.jsonl' to find all your audio files.
2. EXTRACTS "FINGERPRINTS": Uses the `pyannote/embedding` model to listen to each file 
   and convert the voice into a mathematical vector (embedding). 
   This is far more accurate than the previous Method.
3. CLUSTERS THEM: Instead of just comparing A vs B, it collects ALL fingerprints 
   and groups them into "Families" (Speakers) using Agglomerative Clustering.
4. SAVES RESULT: Writes a new manifest with corrected 'speaker' IDs (e.g., 'speaker_0', 'speaker_1').

REQUIREMENTS:
-------------
- pip install pyannote.audio sklearn
- HuggingFace Token (exported as env var `HF_TOKEN` or logged in via `huggingface-cli login`)
- Accepted license for `pyannote/embedding` on HuggingFace Hub.
"""

def parse_args():
    parser = argparse.ArgumentParser(description="Cluster speakers using Pyannote embeddings.")
    parser.add_argument("--input", type=Path, required=True, help="Input manifest path.")
    parser.add_argument("--output", type=Path, required=True, help="Output manifest path.")
    parser.add_argument("--token", type=str, default=None, help="HuggingFace Token (optional if logged in).")
    parser.add_argument("--batch-size", type=int, default=32, help="Inference batch size.")
    parser.add_argument("--num-clusters", type=int, default=20000, 
                        help="Number of speakers to group into. Higher = stricter (more speakers). Lower = looser (fewer).")
    parser.add_argument("--detect-multi-speaker", action="store_true", 
                        help="Detect files with >1 speaker using Diarization pipeline (WARNING: Very Slow!).")
    return parser.parse_args()

def main():
    args = parse_args()
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    print(f"Using device: {device}")

    # 1. Load Model
    print("Loading Pyannote Embedding Model (Wespeaker)...")
    try:
        # Fix for PyTorch 2.6+ safe globals issue
        try:
            import pyannote.audio.core.task
            torch.serialization.add_safe_globals([
                torch.torch_version.TorchVersion,
                pyannote.audio.core.task.Specifications,
                pyannote.audio.core.task.Problem,
                pyannote.audio.core.task.Resolution,
            ])
        except (AttributeError, ImportError):
            pass # Older torch versions or missing imports

        # Using the modern Wespeaker model which is the successor to pyannote/embedding
        model = Model.from_pretrained("pyannote/wespeaker-voxceleb-resnet34-LM", use_auth_token=args.token)
        inference = Inference(model, window="whole", device=device)
    except Exception as e:
        print(f"\n[ERROR] Could not load Pyannote model: {e}")
        print("Make sure you accepted the license at: https://huggingface.co/pyannote/wespeaker-voxceleb-resnet34-LM")
        print("And ran: huggingface-cli login\n")
        return

    # Load Diarization Pipeline ONLY if requested (very slow!)
    diarization_pipeline = None
    if args.detect_multi_speaker:
        print("\n[WARNING] Multi-speaker detection is ENABLED. This will drastically increase processing time!")
        print("Loading Pyannote Diarization Pipeline (Offline Cache)...")
        from pyannote.audio import Pipeline
        import os
        
        # Point directly to the cached config file to bypass auth errors
        cache_dir = Path(os.path.expanduser("~")) / ".cache/huggingface/hub/models--pyannote--speaker-diarization-3.1/snapshots"
        try:
            # Find the first snapshot folder
            if cache_dir.exists():
                snap_dirs = [d for d in cache_dir.iterdir() if d.is_dir()]
                if snap_dirs:
                    config_path = snap_dirs[0] / "config.yaml"
                    diarization_pipeline = Pipeline.from_pretrained(str(config_path))
                    diarization_pipeline.to(device)
                    print("Diarization Pipeline loaded successfully from cache.")
                else:
                    print("[ERROR] No snapshot found in diarization cache.")
            else:
                print("[ERROR] Diarization cache directory not found.")
        except Exception as e:
            print(f"[ERROR] Could not load Diarization Pipeline: {e}")
            return

    # 2. Read Manifest
    print(f"Reading manifest: {args.input}")
    entries = []
    with args.input.open("r", encoding="utf-8") as f:
        for line in f:
            if line.strip():
                entries.append(json.loads(line))
    
    # 3. Extract Embeddings (with caching)
    cache_path = args.input.parent / f"{args.input.stem}_embeddings.npy"
    cache_ids_path = args.input.parent / f"{args.input.stem}_ids.json"
    
    if cache_path.exists() and cache_ids_path.exists():
        print(f"Loading cached embeddings from {cache_path}...")
        embeddings = np.load(cache_path)
        with open(cache_ids_path, 'r', encoding='utf-8') as f:
            valid_indices = json.load(f)
        
        # Filter entries to match the cache
        # We need to ensure 'entries' aligns with 'embeddings'
        # The cache stores indices of the ORIGINAL entries list that were valid
        valid_entries = [entries[i] for i in valid_indices]
        
    else:
        print(f"Extracting embeddings for {len(entries)} files...")
        embeddings = []
        valid_entries = []
        valid_indices = []
        multi_speaker_files = []
        
        # Resolving relative paths
        dataset_root = args.input.parent.parent 

        for i, entry in enumerate(tqdm(entries)):
            audio_path = Path(entry['audio_path'])
            
            # FIX: The audio is likely in the 'dataset' subfolder
            possible_paths = [
                audio_path,
                Path("dataset") / audio_path,
                args.input.parent / audio_path,
                Path("thaitts_emilia_dataset") / audio_path.name
            ]

            found_path = None
            for p in possible_paths:
                if p.exists():
                    found_path = p
                    break
            
            if not found_path:
                continue

            try:
                embedding = inference(str(found_path))
                embeddings.append(embedding)
                valid_entries.append(entry)
                valid_indices.append(i)
                
                # Check for multiple speakers if pipeline is active
                if diarization_pipeline is not None:
                    # Diarize the file
                    diar = diarization_pipeline(str(found_path))
                    speakers_found = diar.labels()
                    if len(speakers_found) > 1:
                        multi_speaker_files.append(str(found_path))
                        
            except Exception:
                pass

        if not embeddings:
            print("No embeddings extracted.")
            return

        embeddings = np.array(embeddings)
        print(f"Extracted {len(embeddings)} embeddings. Saving cache...")
        np.save(cache_path, embeddings)
        with open(cache_ids_path, 'w', encoding='utf-8') as f:
            json.dump(valid_indices, f)
            
        if diarization_pipeline is not None:
            multi_out = args.input.parent / "multi_speaker_detected.txt"
            with open(multi_out, "w", encoding="utf-8") as f:
                for mf in multi_speaker_files:
                    f.write(mf + "\n")
            print(f"Found {len(multi_speaker_files)} files with multiple speakers. List saved to {multi_out}")

    print(f"Embeddings shape: {embeddings.shape}")

    # 4. Clustering (Scalable)
    # AgglomerativeClustering is O(N^2) memory, which crashes on 300k files (requires 300GB+ RAM).
    # We MUST use MiniBatchKMeans for this size.
    from sklearn.cluster import MiniBatchKMeans
    
    # Heuristic: Estimate number of speakers.
    n_clusters = args.num_clusters
    if n_clusters > len(embeddings):
        print(f"[WARNING] You asked for {n_clusters} clusters, but there are only {len(embeddings)} audio files!")
        n_clusters = len(embeddings) // 2
        if n_clusters < 1: n_clusters = 1
        print(f"Auto-adjusting clusters to: {n_clusters}")

    print(f"Clustering into {n_clusters} estimated groups using MiniBatchKMeans...")
    
    clustering = MiniBatchKMeans(
        n_clusters=n_clusters,
        batch_size=4096,
        random_state=42,
        init='random',
        n_init=1
    ).fit(embeddings)

    labels = clustering.labels_
    num_speakers = len(set(labels))
    print(f"Found {num_speakers} cluster groups.")

    # 5. Save Result
    print(f"Saving to {args.output}...")
    args.output.parent.mkdir(parents=True, exist_ok=True)
    
    with args.output.open("w", encoding="utf-8") as f:
        for entry, label in zip(valid_entries, labels):
            entry['speaker'] = f"speaker_{label}" # Assign the new Cluster ID
            f.write(json.dumps(entry, ensure_ascii=False) + "\n")

    print("Done! You now have a properly labeled manifest.")

if __name__ == "__main__":
    main()