o Zh@sdZddlZddlmZmZmZddlZddlZddlmZddl m Z ddl m Z m Z mZddlmZdd lmZdd lmZdd lmZmZdd lmZmZmZmZdd lmZddlm Z e!e"Z#dZ$dZ%dZ&gdZ'dZ(dZ)Gdddej*Z+Gdddej*Z,Gdddej*Z-Gdddej*Z.Gdddej*Z/Gdd d ej*Z0Gd!d"d"ej*Z1Gd#d$d$ej*Z2Gd%d&d&ej*Z3Gd'd(d(eZ4d)Z5d*Z6Gd+d,d,e4Z7e d-e5Gd.d/d/e4Z8e d0e5Gd1d2d2e4Z9gd3Z:dS)4zPyTorch M-CTC-T model.N)OptionalTupleUnion)nn)ACT2FN)add_code_sample_docstringsadd_start_docstrings%add_start_docstrings_to_model_forward)is_deepspeed_zero3_enabled)is_fsdp_managed_module)_prepare_4d_attention_mask)BaseModelOutputCausalLMOutput)PreTrainedModelapply_chunking_to_forward find_pruneable_heads_and_indicesprune_linear_layer)logging) MCTCTConfigrzspeechbrain/m-ctc-t-large)rizY"Mr. Quilter is the apostle of the middle classes, and we're glad to welcome his gospel."gv@cs(eZdZdZfddZddZZS)MCTCTConv1dSubsamplerz Convolutional subsampler: a stack of 1D convolution (along temporal dimension) followed by non-linear activation via gated linear units (https://arxiv.org/abs/1911.08460) cst|_|j_t|j_|j _ |j |j _ j dkr1|jdur,td|j_nd_|jd_|j_|j_tfddtjD_dS)NrzbNeed to specify `conv_channels` configuration in `MCTCTConfig` to use multiple convolution layers.c3s\|])\}}tj|dkrjnj||jdkrj|nj|j|ddVqdS)rrZvalid) kernel_sizestridepaddingN)rConv1d in_channels mid_channels num_layers out_channelsr).0ikselfb/var/www/auris/lib/python3.10/site-packages/transformers/models/deprecated/mctct/modeling_mctct.py Ys z1MCTCTConv1dSubsampler.__init__..)super__init__configZ conv_glu_dimglu_dimrDropoutZ conv_dropoutdropoutnum_conv_layersr Zinput_feat_per_channelZinput_channelsrZ conv_channels ValueErrorr hidden_sizer! conv_kernelr conv_strider ModuleList enumerate conv_layersr&r, __class__r%r(r+=s&     zMCTCTConv1dSubsampler.__init__cCstdd|jD}tjj|dd||fdd}|dd}|jD]}||}tjj ||j d}| |}q#|dd}|S)NcSsg|]}|dqS)rr')r"sizer'r'r( gz1MCTCTConv1dSubsampler.forward..rZconstantrrdim) sumrtorchr functionalpad transpose contiguousr7Zglur-r/)r&input_featuresr hidden_statesconvr'r'r(forwardds  zMCTCTConv1dSubsampler.forward__name__ __module__ __qualname____doc__r+rI __classcell__r'r'r9r(r7s 'rcs,eZdZdZfddZ dddZZS) MCTCTEmbeddingszGConstruct the embeddings from word, position and token_type embeddings.csttj|j|j|jd|_t|j|j|_ t|j |j|_ t |_ t|j|_|jdt|jddd|jdtj|jtj|jjddddS)N) padding_idx position_ids)rF) persistenttoken_type_idsdtypedevice)r*r+r Embedding vocab_sizer2 pad_token_idword_embeddingsmax_position_embeddingsZposition_embeddingsZtype_vocab_sizetoken_type_embeddingsMCTCTLayerNorm LayerNormr.hidden_dropout_probr/Zregister_bufferrAarangeexpandzerosrRr;longrXr8r9r'r(r+ws  zMCTCTEmbeddings.__init__Nrc Cs|dur|n|dd}|d}|dur%|jdd|||f}|durOt|drD|jddd|f}||d|} | }n tj|tj|jjd}|durX| |}| |} || } | | } | | } | S)NrSrrUrrV) r;rRhasattrrUrcrArdrerXr\r^r`r/) r&rFrUrR inputs_embedsZpast_key_values_lengthZ input_shapeZ seq_lengthZbuffered_token_type_idsZ buffered_token_type_ids_expandedr^Z embeddingsr'r'r(rIs"      zMCTCTEmbeddings.forward)NNNNrrJr'r'r9r(rPts  rPcsDeZdZfddZddZddZddZ d d d ZZS)MCTCTSelfAttentioncst|j|jdkrt|dstd|jd|jd|j|_|j|_|j|j|_t j |j|jdd|_ t j |j|jdd|_ t j |j|jdd|_ t |j|_|j|_t d|jd |j|_|j|_dS) NrZembedding_sizezThe hidden size (z6) is not a multiple of the number of attention heads ()Fbiasrr)r*r+r2num_attention_headsrfr1Zattention_head_dimattention_head_size all_head_sizerLinearquerykeyvaluer.Zattention_probs_dropout_probr/r]rYdistance_embedding is_decoderr8r9r'r(r+s"   zMCTCTSelfAttention.__init__cCs6|dd|j|jf}|j|}|ddddS)NrSrrr)r;rlrmviewpermute)r&xZ new_x_shaper'r'r(transpose_for_scoress z'MCTCTSelfAttention.transpose_for_scorescCsFt|jdkr|jttt|j}|jt|jttt|S)Nr)lenshaperwreversedrangeZreshape)r&rxr{r'r'r(reshape_fortrans z"MCTCTSelfAttention.reshape_fortrancCs|dddd}|j\}}}}tj|tj||||f|jdfdd}||||||d|g}|ddd||d|f}|||||d||g}|d}|dd|||fdd}|ddddS)NrrrurrXr>)rwr{rAcatrdrXr~rD)r&ZscoresbatchZ hidden_stateZseq_lenheadsZ halfpointr'r'r("relative_position_embedding_rotates&  z5MCTCTSelfAttention.relative_position_embedding_rotateNFcCs||}|t|j}|||}|||}||}t|| dd} |j j } t d| | dd} | | } | | } |durL| |} tjj| dd} || } |dura| |} t| |} | ddddjdd } |r{| | f}|S| f}|S) NrSzlh, bche -> bclerrur>rr)Z start_dim)rpmathsqrtrmryrqrrrAmatmulrDrsweightZeinsumrrrBZsoftmaxr/rwflatten)r&rGattention_mask head_maskoutput_attentionsZmixed_query_layerZ key_layerZ value_layerZ query_layerZattention_scoresZpositional_embeddingZrelative_position_scoresZattention_probsZ context_layeroutputsr'r'r(rIs,      zMCTCTSelfAttention.forwardNNF) rKrLrMr+ryr~rrIrOr'r'r9r(rhs rhc$eZdZfddZddZZS)r_cs2tttd|_ttd|_dSNr) r*r+r ParameterrAonessingleton_weightrdsingleton_biasr%r9r'r(r+s zMCTCTLayerNorm.__init__cCs||j|jSN)rrr&rGr'r'r(rI szMCTCTLayerNorm.forwardrKrLrMr+rIrOr'r'r9r(r_s r_cr)MCTCTSelfOutputcsLt||_tj|j|jdd|_tj|j|jd|_t |j |_ dSNFrj)Zeps) r*r+r,rror2denser`layer_norm_epsr.rar/r8r9r'r(r+%s zMCTCTSelfOutput.__init__cC&||}||}|||}|Srrr/r`r&rGZ input_tensorr'r'r(rI,  zMCTCTSelfOutput.forwardrr'r'r9r(r$s rcs4eZdZfddZddZ   d ddZZS) MCTCTAttentioncs*tt||_t||_t|_dSr)r*r+rhr&routputset pruned_headsr8r9r'r(r+4s    zMCTCTAttention.__init__cCst|dkrdSt||jj|jj|j\}}t|jj||j_t|jj||j_t|jj ||j_ t|j j |dd|j _ |jjt||j_|jj|jj|j_ |j ||_dS)Nrrr>)rzrr&rlrmrrrprqrrrrrnunion)r&rindexr'r'r( prune_heads:s zMCTCTAttention.prune_headsNFcCs6|||||}||d|}|f|dd}|S)Nrr)r&r)r&rGrrrZ self_outputsattention_outputrr'r'r(rILszMCTCTAttention.forwardr)rKrLrMr+rrIrOr'r'r9r(r3s rcr)MCTCTIntermediatecsHttj|j|jdd|_t|jt rt |j|_ dS|j|_ dS)NFrj) r*r+rror2intermediate_sizer isinstanceZ hidden_actstrrintermediate_act_fnr8r9r'r(r+`s   zMCTCTIntermediate.__init__cCs||}||}|Sr)rrrr'r'r(rIhs  zMCTCTIntermediate.forwardrr'r'r9r(r_s rcr) MCTCTOutputcsFttj|j|jdd|_tj|j|jd|_t |j |_ dSr) r*r+rrorr2rr`rr.rar/r8r9r'r(r+os zMCTCTOutput.__init__cCrrrrr'r'r(rIurzMCTCTOutput.forwardrr'r'r9r(rns rcs:eZdZdeffdd Z   d ddZdd ZZS) MCTCTLayerr,csBtd|_|j|_t||_t||_|j|_t ||_ dSr) r*r+ seq_len_dimchunk_size_feed_forwardr intermediater attentionrtrrr8r9r'r(r+}s   zMCTCTLayer.__init__NFc CsH|j||||d}|d}|dd}t|j|j|j|}|f|}|S)N)rrr)rrfeed_forward_chunkrr) r&rGrrrZself_attention_outputsrr layer_outputr'r'r(rIs  zMCTCTLayer.forwardcCs||}|||}|Sr)rr)r&rZintermediate_outputrr'r'r(rs  zMCTCTLayer.feed_forward_chunkr)rKrLrMrr+rIrrOr'r'r9r(r|s rc@s@eZdZdZeZdZdZdZddZ de j fdd Z d d Z d S) MCTCTPreTrainedModelz An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained models. mctctrFTcCs |jj}t|tjr|jjjd|d|jdur|jj nDt|tj r=|jjjd|d|j dur<|jj|j  n&t|tj rQ|jj |jj dnt|trc|jj d|jj t|tjtjfr|jjjd|d|jdur|jj dSdSdS)zInitialize the weightsg)meanstdNg?)r,Zinitializer_rangerrrordataZnormal_rkZzero_rYrQr`Zfill_r_rrr)r&modulerr'r'r( _init_weightss.          z"MCTCTPreTrainedModel._init_weights input_lengthscCshd}tt|jj|jj|jjD]!\}}}|d}|d|||dd}tj||ddd}q|S)zH Computes the output length of the convolutional layers rrtrunc)Z rounding_mode)zipr}r,r0r3r4rAdiv)r&rZdilation_Z kernel_szrrr'r'r( _get_feat_extract_output_lengthssz5MCTCTPreTrainedModel._get_feat_extract_output_lengthscCst|jdkr|dddddf}||d}|d}tj||f|j|jd}d|tj ||jd|df<| dg d dg }|S)NrrSrrVrr) rzr{rr@r;rArdrWrXrbflipZcumsumre)r&Zfeature_vector_lengthrZsubsampled_lengthsZbszr'r'r("_get_feature_vector_attention_masks z7MCTCTPreTrainedModel._get_feature_vector_attention_maskN)rKrLrMrNr config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingrrA LongTensorrrr'r'r'r(rs raH This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) sub-class. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`MCTCTConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. a Args: input_features (`torch.LongTensor` of shape `({0})`): Indices of input sequence tokens in the vocabulary. Indices can be obtained using [`Wav2Vec2CTCTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.FloatTensor` of shape `({0})`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple. cs^eZdZdeffdd Z   ddejdejdejd ed ed ed ee e ffd dZ Z S) MCTCTEncoderr,csPtj|_t|_t|_tfddt j D|_ d|_ dS)Ncsg|]}tqSr')r)r"rr,r'r(r<r=z)MCTCTEncoder.__init__..F) r*r+rar_ layer_normrrHrr5r}Znum_hidden_layerslayersgradient_checkpointingr8r9rr(r+s    zMCTCTEncoder.__init__FTrFrrroutput_hidden_states return_dictreturncCs|dur|n|jj}|dur|n|jj}|dur|n|jj}||}||}|dur5||jd|}tj j ||j |j d}|durJt ||j}|rNdnd} |rTdnd} |durw|dt|jkrwtdt|jd|ddtp}t|} t|jD]R\} } |r| |f} tg}|j r||jjkrdnd }|r| r|jr|j r|| j|||dur|| nd|}n| |||d }|d}|rd }|r| |df} q|r| |f} |std d || | fDSt|| | dS)Nr)ptrainingr'rz&The head_mask should be specified for z layers, but it is for .TF)rGrr)NNcss|] }|dur|VqdSrr')r"vr'r'r(r)msz'MCTCTEncoder.forward..Zlast_hidden_staterG attentions)r,rruse_return_dictrrHrr{rrBr/rarr rWr;rzrr1r r r6rAZrandZ layerdroprZ_gradient_checkpointing_func__call__tupler)r&rFrrrrrrgrGZencoder_statesZall_attentionsZ synced_gpusidxZ encoder_layerZdropout_probabilityZskip_the_layerZ layer_outputsr'r'r(rIsj           zMCTCTEncoder.forward)FFT) rKrLrMrr+rATensorboolrrrrIrOr'r'r9r(rs( rzaThe bare M-CTC-T Model transformer outputting raw hidden-states without any specific head on top.cseZdZfddZeedeee e de d     dde j dee j d ee j d eed eed eed eee ffddZZS) MCTCTModelcs(t|||_t||_|dSr)r*r+r,rencoder post_initr8r9r'r(r+xs   zMCTCTModel.__init__zbatch_size, sequence_lengthZaudio) checkpoint output_typerZmodalityexpected_outputNrFrrrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}|dur&td|j||||||d}|d}|s@|f|ddSt||j|jdS)Nz#You have to specify input_features.rrrrrrrr) r,rrrr1rrrGr) r&rFrrrrrZencoder_outputsZsequence_outputr'r'r(rIs,zMCTCTModel.forward)NNNNN)rKrLrMr+r MCTCT_INPUTS_DOCSTRINGformatr_CHECKPOINT_FOR_DOCr_CONFIG_FOR_DOC_EXPECTED_OUTPUT_SHAPErArrrrrrIrOr'r'r9r(rss<    rzcMCTCT Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC).cseZdZfddZeeeeee e e d      dde j dee j dee j deed eed eed ee jd eeeffd dZZS) MCTCTForCTCcsTt|t||_|jdurtd|jd|j}t ||j|_ | dS)NzYou are trying to instantiate z with a configuration that does not define the vocabulary size of the language model head. Please instantiate the model as follows: `MCTCTForCTC.from_pretrained(..., vocab_size=vocab_size)`. or define `vocab_size` of your model's configuration.) r*r+rrrZr1r:r2rroctc_headr)r&r,Zoutput_hidden_sizer9r'r(r+s     zMCTCTForCTC.__init__)rrrrZ expected_lossNrFrrrrrlabelsrc Cs~|dur||jjkrtd|jj|dur|n|jj}|j||||||d}|d} || } d} |dur|dur?|n tj|j ddtj d}| | d tj } |dk} | d}|| }tjj| dtjddd}tjjjd d tjj||| ||jj|jj|jjd } Wdn1swY|s| f|td}| dur| f|S|St| | |j|jd S) a labels (`torch.LongTensor` of shape `(batch_size, target_length)`, *optional*): Labels for connectionist temporal classification. Note that `target_length` has to be smaller or equal to the sequence length of the output logits. Indices are selected in `[-100, 0, ..., config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size - 1]`. Nz$Label values must be <= vocab_size: rrrS)rW)r?rWrF)enabled)blankZ reductionZ zero_infinity)losslogitsrGr)maxr,rZr1rrrrArr{rerr@toZ masked_selectrrBZ log_softmaxZfloat32rDbackendsZcudnnflagsZctc_lossr[Zctc_loss_reductionZctc_zero_infinity_HIDDEN_STATES_START_POSITIONrrGr)r&rFrrrrrrrrGrrrZ labels_maskZtarget_lengthsZflattened_targetsZ log_probsrr'r'r(rIsR     zMCTCTForCTC.forward)NNNNNN)rKrLrMr+r rrrrr_CTC_EXPECTED_OUTPUT_CTC_EXPECTED_LOSSrArrrrrrrIrOr'r'r9r(rsB   r)rrr);rNrtypingrrrrAZtorch.utils.checkpointrZ activationsrZ file_utilsrr r Zintegrations.deepspeedr Zintegrations.fsdpr Zmodeling_attn_mask_utilsr Zmodeling_outputsrrZmodeling_utilsrrrrutilsrZconfiguration_mctctrZ get_loggerrKloggerrrrrrrModulerrPrhr_rrrrrrZMCTCT_START_DOCSTRINGrrrr__all__r'r'r'r(s\        =:l ,'E  `8 d