o }oi@sddlmZmZddlmZmZmZmZmZm Z ddl Z eGdddZ eGdddZ eGdd d Z d eed eed efd dZdee de jde jdeded eeee ff ddZGdddZGdddZ ddedeed ee fddZdS)) dataclassfield)AnyDictListOptionalTupleUnionNc@s eZdZUdZeed<eeee j fed<dZ e e ed<dZe ee j ed<dZe eeee j ee j fed<eedZeeee j fed <dZe eeeeeefed <dZe eeeeeefed <dZe eeed <dZe eeed <dZeee j fed<dZee jed<dZe eee efeefed<dZe e j ed<dZe eee efeefed<dZe eeee j fed<dZ e e j ed<dZ!e e j ed<dZ"e eed<e#deefddZ$e#dee fddZ%d"ddZ&d d!Z'dS)# Hypothesisa= Hypothesis class for beam search algorithms. score: A float score obtained from an AbstractRNNTDecoder module's score_hypothesis method. y_sequence: Either a sequence of integer ids pointing to some vocabulary, or a packed torch.Tensor behaving in the same manner. dtype must be torch.Long in the latter case. dec_state: A list (or list of list) of LSTM-RNN decoder states. Can be None. text: (Optional) A decoded string after processing via CTC / RNN-T decoding (removing the CTC/RNNT `blank` tokens, and optionally merging word-pieces). Should be used as decoded string for Word Error Rate calculation. timestamp: (Optional) A list of integer indices representing at which index in the decoding process did the token appear. Should be of same length as the number of non-blank tokens. alignments: (Optional) Represents the CTC / RNNT token alignments as integer tokens along an axis of time T (for CTC) or Time x Target (TxU). For CTC, represented as a single list of integer indices. For RNNT, represented as a dangling list of list of integer indices. Outer list represents Time dimension (T), inner list represents Target dimension (U). The set of valid indices **includes** the CTC / RNNT blank token in order to represent alignments. frame_confidence: (Optional) Represents the CTC / RNNT per-frame confidence scores as token probabilities along an axis of time T (for CTC) or Time x Target (TxU). For CTC, represented as a single list of float indices. For RNNT, represented as a dangling list of list of float indices. Outer list represents Time dimension (T), inner list represents Target dimension (U). token_confidence: (Optional) Represents the CTC / RNNT per-token confidence scores as token probabilities along an axis of Target U. Represented as a single list of float indices. word_confidence: (Optional) Represents the CTC / RNNT per-word confidence scores as token probabilities along an axis of Target U. Represented as a single list of float indices. length: Represents the length of the sequence (the original length without padding), otherwise defaults to 0. y: (Unused) A list of torch.Tensors representing the list of hypotheses. lm_state: (Unused) A dictionary state cache used by an external Language Model. lm_scores: (Unused) Score of the external Language Model. ngram_lm_state: (Optional) State of the external n-gram Language Model. tokens: (Optional) A list of decoded tokens (can be characters or word-pieces. last_token (Optional): A token or batch of tokens which was predicted in the last step. last_frame (Optional): Index of the last decoding step hypothesis was updated including blank token prediction. score y_sequenceNtextdec_out dec_state)default_factory timestamp alignmentsframe_confidencetoken_confidenceword_confidencerlengthylm_state lm_scoresngram_lm_statetokens last_tokentoken_duration last_framereturncsg}tjtr jdnj}t|dkrPjdurPtddjDrGd}d}|D]}||kr6|}d}n|d7}|j||q+|Sfdd |D}|S) zGet per-frame confidence for non-blank tokens according to self.timestamp Returns: List with confidence scores. The length of the list is the same as `timestamp`. timesteprNcss|]}t|tVqdSN) isinstancelist.0ir'_/home/ubuntu/.local/lib/python3.10/site-packages/nemo/collections/asr/parts/utils/rnnt_utils.py {sz8Hypothesis.non_blank_frame_confidence..csg|]}j|qSr')r)r%tselfr'r( sz9Hypothesis.non_blank_frame_confidence..)r"rdictlenranyappend)r.non_blank_frame_confidencert_prevoffsetr,r'r-r(r4psz%Hypothesis.non_blank_frame_confidencecCs|jdurgS|jS)zVGet words from self.text Returns: List with words (str). N)r splitr-r'r'r(wordsszHypothesis.wordsothercCs|j|j7_|jdur|j|_nt|jtjr&tj|j|jfdd|_n|j|j|j|_|jdur;|j|_nt|jtjrOtj|j|jfdd|_n|j|j|j |j 7_ |j |_ |j durl|j |_ n|j |j |j dur}|j |_ n|j |j d|_ |S)z4Merge (inplace) current hypothesis with another one.Nrdim)r r r"torchTensorcatextendrrrrrrr )r.r9r'r'r(merge_s,        zHypothesis.merge_cCs d|_dS)z(Clean the decoding state to save memory.N)rr-r'r'r(clean_decoding_state_s z Hypothesis.clean_decoding_state_)r9r rr )(__name__ __module__ __qualname____doc__float__annotations__r rintr<r=r rstrrrrr#rrrrrrrtensorrrrrrrrrrpropertyr4r8r@rAr'r'r'r(r #s6 7( $$$$  r c@s"eZdZUdZeeeed<dS)NBestHypotheseszList of N best hypothesesn_best_hypothesesN)rBrCrDrErrr rGr'r'r'r(rLs rLc@s6eZdZUdZdZeejed<dZ eejed<dS)HATJointOutputzHATJoint outputs for beam search decoding hat_logprobs: standard HATJoint outputs as for RNNTJoint ilm_logprobs: internal language model probabilities (for ILM subtraction) N hat_logprobs ilm_logprobs) rBrCrDrErOrr<r=rGrPr'r'r'r(rNs rNxprefrcCs@t|t|kr dStt|D] }||||krdSqdS)z Obtained from https://github.com/espnet/espnet. Check if pref is a prefix of x. Args: x: Label ID sequence. pref: Prefix label ID sequence. Returns: : Whether pref is a prefix of x. FT)r1range)rQrRr&r'r'r( is_prefixs rThyps topk_idxs topk_logpsgammabetac sg}t|D]H\}fddt||||D}t|ddd}|d} |dttfdd|d dd} t| dkrF|| q|| fgq|S) a Obtained from https://github.com/espnet/espnet Return K hypotheses candidates for expansion from a list of hypothesis. K candidates are selected according to the extended hypotheses probabilities and a prune-by-value method. Where K is equal to beam_size + beta. Args: hyps: Hypotheses. topk_idxs: Indices of candidates hypothesis. Shape = [B, num_candidates] topk_logps: Log-probabilities for hypotheses expansions. Shape = [B, V + 1] gamma: Allowed logp difference for prune-by-value method. beta: Number of additional candidates to store. Return: k_expansions: Best K expansion hypotheses candidates. cs&g|]\}}t|jt|fqSr')rHr rF)r%kv)hypr'r(r/s&z'select_k_expansions..cS|dSNr+r'rQr'r'r(z%select_k_expansions..)keyrr+cs|dkSr^r'r_)rX k_best_expr'r(r`scSr]r^r'r_r'r'r(r`ra) enumeratezipmaxsortedfilterr1r3) rUrVrWrXrY k_expansionsr&hyp_ik_best_exp_valk_best_exp_idx expansionsr')rXr\rcr(select_k_expansionss   rnc @s0eZdZdZ  d%dededeejdeejfddZ d d Z d d Z d&d ej dej dej dej deej f ddZ d&d ej dej dej dej deej f ddZ d&dej dej dej dej deej f ddZ d&dej dej dej dej deej f ddZd'defddZd(d d!Zd)d#d$ZdS)* BatchedHypsz\Class to store batched hypotheses (labels, time_indices, scores) for efficient RNNT decodingN batch_size init_lengthdevice float_dtypecCs|dkr td||dkrtd|||_||_||_||_tj||tjd|_tj||jf|tjd|_ tj||jf|tjd|_ tj||jf|tjd|_ tj|||d|_ tj |fd|tjd|_tj||tjd|_tj||d|_t|j|_dS)a> Args: batch_size: batch size for hypotheses init_length: initial estimate for the length of hypotheses (if the real length is higher, tensors will be reallocated) device: device for storing hypotheses float_dtype: float type for scores rinit_length must be > 0, got batch_size must be > 0, got rrdtyper*rrN) ValueError _max_lengthrprrrsr<zeroslongcurrent_lengths transcript timestampstoken_durationsscoresfulllast_timestamplast_timestamp_lastsarange_batch_indices ones_like _ones_batch)r.rprqrrrsr'r'r(__init__s"zBatchedHyps.__init__cCsX|jd|jd|jd|jd|jd|jd|jddS)2 Clears batched hypotheses state. rr*N)r}fill_r~rrrrrr-r'r'r(clear_=s      zBatchedHyps.clear_cCsltj|jt|jfdd|_tj|jt|jfdd|_tj|jt|jfdd|_|jd9_dS) Allocate 2x space for tensors, similar to common C++ std::vector implementations to maintain O(1) insertion time complexity r*r:N)r<r>r~ zeros_likerrrzr-r'r'r(_allocate_moreIszBatchedHyps._allocate_moreactive_indiceslabels time_indicesrrcCsF|jddkr dS|j|jkr||j|||||ddS)a Add results (inplace) from a decoding step to the batched hypotheses. We assume that all tensors have the same first dimension, and labels are non-blanks. Args: active_indices: tensor with indices of active hypotheses (indices should be within the original batch_size) labels: non-blank labels to add time_indices: tensor of time index for each label scores: label scores rN)rrrrr)shaper}rfitemrzradd_results_no_checks_)r.rrrrrr'r'r( add_results_Ss zBatchedHyps.add_results_cCs|j||7<|j|}||j||f<||j||f<|dur'||j||f<t|j||k|j|dd|j|<||j|<|j|d7<dS)a Add results (inplace) from a decoding step to the batched hypotheses without checks. We assume that all tensors have the same first dimension, and labels are non-blanks. Useful if all the memory is pre-allocated, especially with cuda graphs (otherwise prefer a more safe `add_results_`) Args: active_indices: tensor with indices of active hypotheses (indices should be within the original batch_size) labels: non-blank labels to add time_indices: tensor of time index for each label scores: label scores token_durations: predicted durations for each token by TDT head Nr+) rr}r~rrr<whererr)r.rrrrractive_lengthsr'r'r(rrs  z"BatchedHyps.add_results_no_checks_ active_maskcC4|j||jkr||j|||||ddS)a Add results (inplace) from a decoding step to the batched hypotheses. We assume that all tensors have the same first dimension, and labels are non-blanks. Args: active_mask: tensor with mask for active hypotheses (of batch_size) labels: non-blank labels to add time_indices: tensor of time index for each label scores: label scores token_durations: token durations for TDT )rrrrrNr}rfrzradd_results_masked_no_checks_r.rrrrrr'r'r(add_results_masked_s zBatchedHyps.add_results_masked_cCstj||j||j|jd||j|j|jf<||j|j|jf<|dur-||j|j|jf<tjt||j |k|j d|j |j dtjt||j |k|j |j |j dtj|||j |j d|j|7_dS)af Add results (inplace) from a decoding step to the batched hypotheses without checks. We assume that all tensors have the same first dimension, and labels are non-blanks. Useful if all the memory is pre-allocated, especially with cuda graphs (otherwise prefer a more safe `add_results_`) Args: active_mask: tensor with mask for active hypotheses (of batch_size) labels: non-blank labels to add time_indices: tensor of time index for each label scores: label scores token_durations: token durations for TDT )outNr+) r<rrr~rr}rr logical_andrrrrr'r'r(rs&z)BatchedHyps.add_results_masked_no_checks_r*pad_idcCs&t|jdk|j|j|jdf|S)z7Get last labels. For elements without labels use pad_idrr+)r<rr}r~r)r.rr'r'r(get_last_labelsszBatchedHyps.get_last_labelsrcCs~t|j|j|j|jd}|j|j|j|j|j|j|j |j |j |j |j |j |j |j |S)Return a copy of self)rprqrrrs) rorprzrrrsr}copy_r~rrrrr)r. batched_hypsr'r'r(cloneszBatchedHyps.cloner9cCs"tj|jt|jfdd|_tj|jt|jfdd|_tj|jt|jfdd|_|j|j7_tj|jjd|j j d}|j dddf|dddf}|jj d||jd|jj d||jd|jj d||jd|j |j 7_ |j |j 7_ |j |j |j|j|S)z Merge two batched hypotheses structures. NB: this will reallocate memory Args: other: BatchedHyps r*r:r+rxN)r;indexsrc)r<r>r~rrrrzrrr}rrscatter_rrrr)r.r9indicesshifted_indicesr'r'r(r@s"zBatchedHyps.merge_NNr!)r*)rro)r9rorro)rBrCrDrErHrr<rrrwrrrr=rrrrrrr@r'r'r'r(ros ,  % + " 0 roc@s eZdZdZ     d"dedededeejd eejd e d e d e fd dZ ddZ ddZ   d#dej dej deej deej deej f ddZ   d#dej dej deej deej deej f ddZ   d#dej dej deej deej deej f ddZd$d d!ZdS)%BatchedAlignmentsz Class to store batched alignments (logits, labels, frame_confidence). Size is different from hypotheses, since blank outputs are preserved NTFrp logits_dimrqrrrsstore_alignmentsstore_frame_confidencewith_duration_confidencec Cs@|dkr td||dkrtd|||_||_||_||_||_||_||_||_t j ||jf|t j d|_ t j ||t j d|_ t j d||d|_t j d|t j d|_|jrut j ||j|f||d|_t j ||jf|t j d|_t j d||d|_|jrt j |jr||jdgn||jg||d|_t j||d|_dS)a Args: batch_size: batch size for hypotheses logits_dim: dimension for logits init_length: initial estimate for the lengths of flatten alignments device: device for storing data float_dtype: expected logits/confidence data type store_alignments: if alignments should be stored store_frame_confidence: if frame confidence should be stored rrtrurvrrxN)ryrprrrrswith_frame_confidencerwith_alignmentsrzr<r{r|rr}logitsrrrr) r.rprrqrrrsrrrr'r'r(rs6zBatchedAlignments.__init__cCs@|jd|jd|jd|jd|jddS)rrrN)r}rrrrrr-r'r'r(rVs    zBatchedAlignments.clear_cCstj|jt|jfdd|_|jr0tj|jt|jfdd|_tj|jt|jfdd|_|jrBtj|jt|jfdd|_|j d9_ dS)rr*r:r+rN) r<r>rrrrrrrrzr-r'r'r(r`sz BatchedAlignments._allocate_morerrrr confidencecCs|jddkr dS|j|jkr||j|}||j||f<|jr<|dur<|dur<||j||f<||j ||f<|j rJ|durJ||j ||f<|j|d7<dS)a Add results (inplace) from a decoding step to the batched hypotheses. All tensors must use the same fixed batch dimension. Args: active_mask: tensor with mask for active hypotheses (of batch_size) logits: tensor with raw network outputs labels: tensor with decoded labels (can contain blank) time_indices: tensor of time index for each label confidence: optional tensor with confidence for each item in batch rNr+) rr}rfrrzrrrrrrr)r.rrrrrrr'r'r(rms zBatchedAlignments.add_results_rcCr)a Add results (inplace) from a decoding step to the batched hypotheses. All tensors must use the same fixed batch dimension. Args: active_mask: tensor with indices of active hypotheses (indices should be within the original batch_size) time_indices: tensor of time index for each label logits: tensor with raw network outputs labels: tensor with decoded labels (can contain blank) confidence: optional tensor with confidence for each item in batch )rrrrrNrr.rrrrrr'r'r(rs   z%BatchedAlignments.add_results_masked_cCs||j|j|jf<|jr/|dur/|dur/||j|j|jf<||j|j|jf<||j|j|jf<|jr?|dur?||j|j|jf<|j|7_dS)a Add results (inplace) from a decoding step to the batched hypotheses. All tensors must use the same fixed batch dimension. Useful if all the memory is pre-allocated, especially with cuda graphs (otherwise prefer a more safe `add_results_masked_`) Args: active_mask: tensor with indices of active hypotheses (indices should be within the original batch_size) time_indices: tensor of time index for each label logits: tensor with raw network outputs labels: tensor with decoded labels (can contain blank) confidence: optional tensor with confidence for each item in batch N)rrr}rrrrrrr'r'r(rsz/BatchedAlignments.add_results_masked_no_checks_rc Csrt|j|j|j|j|j|j|j|jd}|j |j |j |j |j |j |j |j |j |j|S)r)rprrqrrrsrrr)rrprrzrrrsrrrr}rrrrr)r.batched_alignmentsr'r'r(rs  zBatchedAlignments.clone)NNTFF)NNN)rr)rBrCrDrErHrr<rrrwboolrrrr=rrrrr'r'r'r(rs  ;  +  !rrrc s|dus|jjdksJ|durjjdn|}jjj j  fddt|D}|dur|j}|jrV|j |j |j r^|j tt |D][g|_|j rsg|_ tj|jd|fdd\}}d|D]3}|jr|jfddt|D|j r|j fddt|D|7qqd|S) aB Convert batched hypotheses to a list of Hypothesis objects. Keep this function separate to allow for jit compilation for BatchedHyps class (see tests) Args: batched_hyps: BatchedHyps object alignments: BatchedAlignments object, optional; must correspond to BatchedHyps if present batch_size: Batch Size to retrieve hypotheses. When working with CUDA graphs the batch size for all tensors is constant, thus we need here the real batch size to return only necessary hypotheses Returns: list of Hypothesis objects Nrc szg|]9}t||d|f|dj|ftj|dj|fdks2ntddddqS)Nr)r r rrrr)r rr}r<allremptyr$)rr} durationsrrr~r'r(r/s z.batched_hyps_to_hypotheses..T) return_countscs,g|]}|f|ffqSr'r'r%j)alignment_labelsalignment_logitsr&startr'r(r/scsg|] }|fqSr'r'r)rr&rr'r(r/%s)rrcpur}r~rrStolistrrrrrr1rr<unique_consecutiver3) rrrpnum_hyps hypothesesalignment_lengths_grouped_counts timestamp_cntr') rrrr}rrr&rrrr~r(batched_hyps_to_hypothesessL               rr) dataclassesrrtypingrrrrrr r<r rLrNrHrrTr=rFrnrorrr'r'r'r(sJ   . N