o oi@sfddlZddlmZddlZddlmZmZddlm Z m Z ddl m Z ddl mZGdd d ZdS) N)List)DatasetAudio)EmbedderConfig WavLMEmbedder)SpeakerMakerUnsupervised)SpeakerEmbeddingGrouperc@s@eZdZdZddZdededefddZdeefd d Z d S) SpeakerEmbeddingProcessora= Orchestrates speaker embedding generation and optional downstream steps (unsupervised clustering, per-speaker averaged embeddings) for a HuggingFace dataset. Receives a dataset (already loaded at 22050 Hz by DatasetProcessor), generates WavLM speaker embeddings using a separate 16 kHz audio channel, then optionally runs: - UMAP + HDBSCAN clustering to assign cluster-based speaker IDs - Per-speaker embedding averaging (group_sp_emb) Returns the dataset with all produced columns added, ready for the tokenization step. cCs4||_t|j|j|j|j|j|jd}t||_ dS)zc Args: settings: SpeakerEmbeddingSettings instance from ConfigManager. ) model_name target_sr max_audio_sec batch_sizeembedding_column_nameuse_multiprocessingN) settingsrr r r rembedding_columnrrembedder)selfrembedder_configrF/home/ubuntu/kanitts-2-dataset-pipeline/utils/speaker_emb/processor.py__init__sz"SpeakerEmbeddingProcessor.__init__dataset audio_columnreturncsTdtjd<|jj}|tdd}|jfddddd }|d t|jjd d }|t|d d }|j d dddd }|j j |d dd}| d }|jj t|}|j fddddd }|t|}|rvtd|dt||jjrt|jj|jjd|jj |jjd}||}|jjr|jjjrt|jj|jj d}||}|S)u Add embedding column and optionally speaker/cluster column and grouped embedding column to the dataset. Audio resampling strategy: 1. Switch audio column to raw-bytes mode (decode=False) — avoids double decode/resample quality loss. 2. Copy raw bytes to a temporary "audio_spk" column. 3. Decode "audio_spk" at target_sr (16 kHz) for embedding. 4. Decode original audio column back at 22050 Hz for tokenization. 5. Filter samples shorter than 0.7 s (too short for reliable embedding). 6. Generate embeddings; remove "audio_spk". 7. Filter samples with NaN embeddings (mandatory, before all downstream steps). 8. Optionally run UMAP + HDBSCAN clustering. 9. Optionally run per-speaker embedding averaging (group_sp_emb). Args: dataset: HuggingFace Dataset loaded at 22050 Hz. audio_column: Name of the audio column (from DatasetConfig). Returns: Dataset with added embedding (and optional clustering / grouping) columns. The audio column is still at 22050 Hz — ready for tokenization. torchaudioDATASETS_AUDIO_BACKENDF)decodecs d|iS)N audio_spkri)rrrNs z3SpeakerEmbeddingProcessor.process.. z#Copying audio for speaker embedding)num_procdescrT) sampling_ratercSs t|dd|dddkS)Nrarrayr&gffffff?)lenr rrrr"as zFiltering short audiospeaker_embedding)rr split_namecstt|dkS)Nr)npisnanr'sumr )emb_colrrr"vszFiltering NaN embeddingsz$[SpeakerEmbeddingProcessor] Removed z) samples with NaN embeddings. Remaining: )UMAPHDBSCAN)cluster_configrspeaker_column)rr)osenvironfeaturesr& cast_columnrmaprr filterr embed_datasetremove_columnsrr(print do_clustersr umap_paramshdbscan_paramsclustering_speaker_column group_embdo_thisr group)rrr original_sr initial_sizeremoved speaker_makergrouperr)rr/rprocess,sr         z!SpeakerEmbeddingProcessor.processcCsF|jjg}|jjr||jj|jjr!|jjjr!||jjj|S)z Return the list of column names added by this processor. These must be passed through the tokenization step (preserve_columns). )rrr=appendr@rArBgrouped_embedding_column)rcolsrrrget_preserved_columnss z/SpeakerEmbeddingProcessor.get_preserved_columnsN) __name__ __module__ __qualname____doc__rrstrrIrrMrrrrr s kr )r4typingrnumpyr,datasetsrrrrrspeaker_clusteringrgroup_embeddingsr r rrrrs