o }oie@sddlZddlZddlmZddlmZddlZddlmZddl m Z ddl m Z m Z ddlmZddlmZmZmZmZdd lmZmZd gZd d ZGd d d e ZdS)N)deepcopy)Dict) DictConfig)ClusteringDiarizer)get_scale_interpolated_embssplit_input_data)OnlineSpeakerClustering)OnlineSegmentoraudio_rttm_mapgenerate_cluster_labelsget_embs_and_timestamps)logging model_utilsOnlineClusteringDiarizercsfdd}|S)z Monitor elapsed time of the corresponding function displaying the method name. Args: method: function that is being measured Return: `timed` function for measuring the elapsed time cstt}|i|}t}d|vr*|dj}t||d|d|<|Std||djf|S)Nlog_timelog_nameiz %2.2fms %r)timeget__name__upperintr info)argskwargstsresulttenamemethod_/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/models/online_diarizer.pytimed0sztimeit..timedr )rr"r rr!timeit%s  r#cs>eZdZdZdeffdd ZddZddZd d Zd d Z d dZ ddZ ddZ ddZ defddZedejfddZdedejdedejfddZeedejdejfd d!Zedejd"ejdejfd#d$Ze %d/d&eeejfdejfd'd(Zdejfd)d*Zed+ejd,ejdejfd-d.ZZS)0ra A class that enables online (streaming) clustering based diarization. - The instance created from `OnlineClusteringDiarizer` sets aside a certain amount of memory to provide the upcoming inference with history information - There are two major modules involved: `OnlineSegmentor` and `OnlineSpeakerClustering`. OnlineSegmentor: Take the VAD-timestamps and generate segments for each scale OnlineSpeakerClustering: Update the entire speaker labels of the given online session while updating the speaker labels of the streaming inputs. - The overall diarization process is done by calling `diarize_step` function. `diarize_step` function goes through the following steps: (1) Segmentation (`OnlineSegmentor` class) (2) Embedding extraction (`_extract_online_embeddings` function call) (3) Online speaker counting and speaker clustering (`OnlineClusteringDiarizer` class) (4) Label generation (`generate_cluster_labels` function call) cfgcst|t||_|jj|_t|jd |_ |j dd|_ |j dd|_ t|jjj|_|jj|_td|jj|_tj|jsOt|jtjr^d|_td|_n d|_td |_||jdS) N scale_dictuniq_iddecimalsrTcudaFcpu) super__init__r#convert_model_config_to_dict_configr$diarizer _cfg_diarizermaxmultiscale_args_dictkeysbase_scale_indexrr&r'r manifest_filepathAUDIO_RTTM_MAP sample_ratetorch manual_seedout_dir_out_dirospathexistsmkdirr) is_availabledevicereset_speaker_modeleval)selfr$ __class__r r!r,Rs&         z!OnlineClusteringDiarizer.__init__cCs@t|j|j|j|j|j|jd|_|j|_|j|_ |jj|_dS)az Initialize online speaker clustering module Attributes: online_clus (OnlineSpeakerClustering): Online clustering diarizer class instance history_n (int): History buffer size for saving history of speaker label inference Total number of embedding vectors saved in the buffer that is kept till the end of the session current_n (int): Current buffer (FIFO queue) size for calculating the speaker label inference Total number of embedding vectors saved in the FIFO queue for clustering inference )max_num_speakersmax_rp_thresholdsparse_search_volumehistory_buffer_sizecurrent_buffer_sizer)N) rrGrHrIrJrKr) online_clus history_n current_nrDclustering_paramsr r r!_init_online_clustering_modulensz7OnlineClusteringDiarizer._init_online_clustering_modulecCst||_dS)z Initialize an online segmentor module Attributes: online_segmentor (OnlineSegmentor): online segmentation module that generates short speech segments from the VAD input N)r online_segmentor)rDr6r r r!_init_online_segmentor_modulesz6OnlineClusteringDiarizer._init_online_segmentor_modulecCsJd|_dd|jdD|_dd|jdD|_tg|_dS)a Variables are kept in memory for future updates Attributes: memory_margin (int): The number of embeddings saved in the memory buffer. This memory margin is dependent on the base scale length: margin = (buffer_length)/(base scale shift) memory margin is automatically calculated to have minimal memory usage memory_segment_ranges (dict): The segment range information kept in the memory buffer memory_segment_indexes (dict): The segment indexes kept in the memory buffer memory_cluster_labels (Tensor): The cluster labels inferred in the previous diarization steps rcSi|]}|gqSr r .0keyr r r! z@OnlineClusteringDiarizer._init_memory_buffer..r%cSrTr r rUr r r!rXrYN) memory_marginr1r2memory_segment_rangesmemory_segment_indexesr7tensormemory_cluster_labelsrDr r r!_init_memory_buffersz,OnlineClusteringDiarizer._init_memory_buffercCs&|dd|_|dd|_i|_dS)a Variables needed for taking majority votes for speaker labels Attributes: use_temporal_label_major_vote (bool): Boolean for whether to use temporal majority voting temporal_label_major_vote_buffer_size (int): buffer size for majority voting base_scale_label_dict (dict): Dictionary containing multiple speaker labels for major voting Speaker labels from multiple steps are saved for each segment index. use_temporal_label_major_voteF%temporal_label_major_vote_buffer_sizeN)rrarbbase_scale_label_dictrOr r r!"_init_temporal_major_voting_modules  z;OnlineClusteringDiarizer._init_temporal_major_voting_modulecCs|i|_i|_i|_i|_i|_|jdD]%}ddg|j|<t g|j|<g|j|<g|j|<g|j|<g|j|<qdS)z Initialize segment variables for each scale. Note that we have `uniq_id` variable in case where multiple sessions are handled. r%N) emb_vectors time_stampssegment_range_tssegment_raw_audiosegment_indexesr1r2$multiscale_embeddings_and_timestampsr7r])rD scale_idxr r r!_init_segment_variabless    z0OnlineClusteringDiarizer._init_segment_variablescCsd|_d|_d|_d|_dS)a Timing variables transferred from OnlineDiarWithASR class. Buffer is window region where input signal is kept for ASR. Frame is window region where the actual inference ASR decoded results are updated Example: buffer_len = 5.0 frame_len = 1.0 |___Buffer___[___________]____________| |____________[ Frame ]____________| | <- buffer_start |____________| <- frame_start |_____________________________________| <- buffer_end buffer_start = 12.0 buffer_end = 17.0 frame_start = 14.0 These timestamps and index variables are updated by OnlineDiarWithASR. Attributes: frame_index (int): Integer index of frame window frame_start (float): The start of the frame window buffer_start (float): The start of the buffer window buffer_end (float): The end of the buffer rN) frame_index frame_start buffer_start buffer_endr_r r r!_init_buffer_frame_timestampss! z6OnlineClusteringDiarizer._init_buffer_frame_timestampscCs"|j|j_|j|j_|j|j_dS)zI Pass the timing information from streaming ASR buffers. N)rprRrqrrr_r r r!!_transfer_timestamps_to_segmentors  z:OnlineClusteringDiarizer._transfer_timestamps_to_segmentorcCsjt|j|jd|jd|_|||jjj | |j j| | |jjj |dS)z Reset all the necessary variables and initialize classes. Attributes: n_embed_seg_len (int): Number of segments needed for 1 second of input time-series signal r%rN)rr6r1r3n_embed_seg_lenrmrQr/ clustering parametersrSr$r`rersr_r r r!rAs zOnlineClusteringDiarizer.resetrlcCs|jd|jd}t|j|j||_ttt|j| tt|j|j |j |j }||j}|j || d|j |<|j || d|j |<|j|| d|j|<|j|| d|j|<dS)a Calculate how many segments should be removed from memory (`memory_margin`) and save the necessary information. `keep_range` determines how many segments and their corresponding embedding, raw audio, timestamps in the memory of the online diarizer instance. Args: scale_idx (int): Scale index in integer type r%rcN)r1r3rrrrqrZlensetscale_mapping_dicttolistrMrNrfrirhrj)rDrlbase_scale_shiftscale_buffer_size keep_ranger r r! _clear_memory s   z&OnlineClusteringDiarizer._clear_memoryreturnc Csg}|j|jD]J}||jvr|j|g|j|<n't|j||jkr5|j|dt|j||jks#|j||j||t t |j|d q|S)aN Take a majority voting for every segment on temporal steps. This feature significantly reduces the error coming from unstable speaker counting in the beginning of sessions. Returns: maj_vote_labels (list): List containing the major-voted speaker labels on temporal domain r) r\r3rdr^rxrbpopappendr7moder]item)rDmaj_vote_labelsseg_idxr r r!_temporal_label_major_vote%s  &z3OnlineClusteringDiarizer._temporal_label_major_votetotal_cluster_labels is_onlinec Cs|}|s%t|j||j|<t|j||j|<||jkr$t||_n|j|d|j|dkrtt|j||j d}t |j|}t ||kdd}t|j||d|j||d<t|j||d|j||d<||jkrt||d|j|d<t |jt |j|krtdt |jdt |j|d||t |j|t |j|krt |j|krt |j|ksntdt |j|dt |j|dt |j|d t |j|d |jr|}|S|j}|S) a Save the temporary input to the class memory buffer. - Clustering is done for (hist_N + curr_N) number of embeddings. - Thus, we need to remove the clustering results on the embedding memory. - If self.diar.history_buffer_seg_end is not None, that indicates streaming diarization system is starting to save embeddings to its memory. Thus, the new incoming clustering label should be separated. - If `is_online = True`, old embeddings outside the window are removed to save GPU memory. Args: scale_idx (int): Scale index in integer total_cluster_labels (Tensor): The speaker labels from the beginning of the session to the current position is_online (bool) Boolean variable that indicates whether the system is currently in online mode or not Returns: cluster_label_hyp (Tensor): Majority voted speaker labels over multiple inferences rNzlself.memory_cluster_labels and self.memory_segment_ranges should always have the same length, but they have z and .zself.emb_vectors, self.segment_raw_audio, self.segment_indexes, and self.segment_range_ts should always have the same length, but they have z, z, and z, respectively.)r{rrhr[rjr\r3r^r0rZr7r]whererx ValueErrorrrfrirar)rDrlrrglobal_stt_idxsegment_indexes_matbuffer_stt_idxcluster_label_hypr r r!save_history_data;sf             z*OnlineClusteringDiarizer.save_history_data audio_signalcsXt|j}tfddt|jdDj}jj ||d\}}|S)a Call `forward` function of the speaker embedding model. Args: audio_signal (Tensor): Torch tensor containing time-series signal Returns: Speaker embedding vectors for the given time-series input `audio_signal`. csg|]}jqSr )ru)rVkr_r r! rYzEOnlineClusteringDiarizer._run_embedding_extractor..r) input_signalinput_signal_length) r7stackfloattor@r]rangeshaperBforward)rDraudio_signal_lens_ torch_embsr r_r!_run_embedding_extractors *z1OnlineClusteringDiarizer._run_embedding_extractorsegment_rangescCs|durdn|jd}t|}||kr:||||}|dus'|jddkr*|}nt|d|ddf|f}n ||krF|dt|}t||jdkrStd|S)aa Incrementally extract speaker embeddings based on `audio_signal` and `segment_ranges` variables. Unlike offline speaker diarization, speaker embedding and subsegment ranges are not saved to disk. Measures the mismatch between `segment_ranges` and `embeddings` then extract the necessary amount of speaker embeddings. Args: audio_signal (Tensor): Torch tensor containing time-series audio signal embeddings (Tensor): Previously existing Torch tensor containing speaker embedding vector segment_ranges(Tensor): Torch tensor containing the start and end of each segment Returns: embeddings (Tensor): Concatenated speaker embedding vectors that match segment range information in `segment_ranges`. Nrz2Segment ranges and embeddings shapes do not match.)rrxrr7vstackr)rDrr embeddingsstt_idxend_idxrr r r!_extract_online_embeddingss z3OnlineClusteringDiarizer._extract_online_embeddingsFuniq_embs_and_timestampsc Cs|rtdntd}t|d|d|dd\}}t|d|||d\}|_t|j|j|j}|j j |||j |d }|j d  D]\} \} } || ||j j} qG| S) a Launch online clustering for `uniq_embs_and_timestamps` input variable. Args: uniq_embs_and_timestamps (dict): Dictionary containing embeddings, timestamps and multiscale weights. If uniq_embs_and_timestamps contains only one scale, single scale diarization is performed. cuda (bool): Boolean indicator for cuda usages r)r*r timestampsmultiscale_segment_counts)embeddings_in_scalestimestamps_in_scalesrmultiscale_weights)rrrr@)curr_embbase_segment_indexesror)r%)r7r@rrrzr]rjr3rrL forward_inferror1itemsrr) rDrr)r@rrrrmerged_clus_labelsrlwindowshiftrr r r!_perform_online_clusterings&   z3OnlineClusteringDiarizer._perform_online_clusteringcCsPt|jdks |jdkrtd|jggdg\}}|St|j|j|j\}}|S)z In case buffer is not filled or there is no speech activity in the input, generate temporary output. Returns: diar_hyp (Tensor): Speaker labels based on the previously saved segments and speaker labels rrn)rxr^rqr total_buffer_in_secsr[r3)rDdiar_hyprr r r!_get_interim_outputsz,OnlineClusteringDiarizer._get_interim_output audio_buffervad_timestampsc Cs||jdkst|dkr|S|jdD]Q\}\}}|jj|||j||j ||j |||d\}}}||j|<||j |<||j |<|j |j||j ||j |d} | |j |<|j | i|j |ig|j|<qt|j|j} |j| |j |jd} t|j|j| \} } | S)a A function for a unit diarization step. Each diarization step goes through the following steps: 1. Segmentation: Using `OnlineSegmentor` class, call `run_online_segmentation` method to get the segments. 2. Embedding Extraction: Extract multiscale embeddings from the extracted speech segments. 3. Online Clustering & Counting Perform online speaker clustering by using `OnlineSpeakerClustering` class. 4. Generate speaker labels: Generate start and end timestamps of speaker labels based on the diarization results. c.f.) Also see method `diarize` in `ClusteringDiarizer` class. Args: audio_buffer (Tensor): Tensor variable containing the time series signal at the current frame Dimensions: (Number of audio time-series samples) x 1 vad_timestamps (Tensor): List containing VAD timestamps. Dimensions: (Number of segments) x 2 Example: >>> vad_timestamps = torch.Tensor([[0.05, 2.52], [3.12, 6.85]]) Returns: diar_hyp (Tensor): Speaker label hypothesis from the start of the session to the current position rr%)rrrirhrjrr)rrr)r))rtrqrxrr1rrRrun_online_segmentationrirhrjrrfr&rkr rr)r r[r3)rDrrrlrr audio_sigsr range_indsrembs_and_timestampsrrrr r r! diarize_steps>    z%OnlineClusteringDiarizer.diarize_step)F)r __module__ __qualname____doc__rr,rQrSr`rermrsrtrArrr#r7Tensorrboolrno_gradrrrstrrrr __classcell__r r rEr!r>sH &J &  (&)r;rcopyrtypingrr7 omegaconfrnemo.collections.asr.modelsr3nemo.collections.asr.parts.utils.offline_clusteringrr2nemo.collections.asr.parts.utils.online_clusteringr.nemo.collections.asr.parts.utils.speaker_utilsr r r r nemo.utilsr r__all__r#rr r r r!s