a hQ @shddlZddlZddlmZmZmZddlZddlZddlm Z ddl m Z ddl m Z ddlmZddlmZdd lmZdd lmZdd lmZmZmZdd lmZmZdd lmZddlm Z m!Z!ddl"m#Z#e!$e%Z&GdddeZ'GdddeZ(GdddeZ)Gddde j*Z+Gddde j*Z,Gddde j*Z-Gddde j*Z.d:e j*ej/ej/ej/eej/ee0e0eej/d d!d"Z1Gd#d$d$e j*Z2Gd%d&d&e j*Z3Gd'd(d(eZ4Gd)d*d*e j*Z5e Gd+d,d,eZ6d;e7e8e8fe0e8eej9e8ej:d-d.d/Z;e Gd0d1d1e6Ze d6d3Gd7d8d8e6Z?gd9Z@dS)<N)CallableOptionalUnion)nn)CrossEntropyLoss)ACT2FN)is_deepspeed_zero3_enabled)is_fsdp_managed_module)FlashAttentionKwargs)GradientCheckpointingLayer)BaseModelOutputCausalLMOutputSequenceClassifierOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringlogging) SEWConfigcs&eZdZdfdd ZddZZS)SEWNoLayerNormConvLayerrcsjt|dkr |j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ dS)Nrr kernel_sizestridebias)super__init__conv_dim in_conv_dim out_conv_dimrConv1d conv_kernel conv_stride conv_biasconvrfeat_extract_activation activationselfconfiglayer_id __class__`/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/sew/modeling_sew.pyr/s  z SEWNoLayerNormConvLayer.__init__cCs||}||}|SN)r%r'r) hidden_statesr.r.r/forward=s  zSEWNoLayerNormConvLayer.forward)r__name__ __module__ __qualname__rr3 __classcell__r.r.r,r/r.srcs&eZdZdfdd ZddZZS)SEWLayerNormConvLayerrcs|t|dkr |j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)NrrrT)Zelementwise_affine)rrrrr rr!r"r#r$r% LayerNorm layer_normrr&r'r(r,r.r/rDs  zSEWLayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)N)r% transposer;r'r1r.r.r/r3Ss      zSEWLayerNormConvLayer.forward)rr4r.r.r,r/r9Csr9cs&eZdZdfdd ZddZZS)SEWGroupNormConvLayerrcst|dkr |j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)NrrrT)Z num_groupsZ num_channelsZaffine)rrrrr rr!r"r#r$r%rr&r' GroupNormr;r(r,r.r/r_s   zSEWGroupNormConvLayer.__init__cCs"||}||}||}|Sr0)r%r;r'r1r.r.r/r3os   zSEWGroupNormConvLayer.forward)rr4r.r.r,r/r?^sr?cs$eZdZfddZddZZS)SEWPositionalConvEmbeddingcs(ttj|j|j|j|jd|j|jd|_tj j }t tj j drRtj j j }t rddl}|jj|jjdd"||jddd|_Wdn1s0Yt |jdr|jj jj}|jj jj}n|jj}|jj}|j|||j||n||jddd|_t|j|_t|j|_dS) N)rpaddinggroupsr weight_normrZ modifier_rankweight)namedimparametrizations)rrrr! hidden_sizenum_conv_pos_embeddingsZnum_conv_pos_embedding_groupssqueeze_factorr%utilsrEhasattrrJr deepspeedzeroGatheredParametersrGZ original0Z original1weight_gweight_vZregister_external_parameterSEWSamePadLayerrCrr&r')r)r*rErPrSrTr,r.r/rws4   0   z#SEWPositionalConvEmbedding.__init__cCs"||}||}||}|Sr0)r%rCr'r1r.r.r/r3s   z"SEWPositionalConvEmbedding.forwardr4r.r.r,r/rAvs "rAcs$eZdZfddZddZZS)rUcs$t|ddkrdnd|_dS)NrBrr)rrnum_pad_remove)r)rLr,r.r/rs zSEWSamePadLayer.__init__cCs,|jdkr(|ddddd|j f}|S)Nr)rVr1r.r.r/r3s zSEWSamePadLayer.forwardr4r.r.r,r/rUs rUcs$eZdZfddZddZZS) SEWUpsamplingcs:tt|j|j|j|_t|j|_ |j|_dSr0) rrrLinearrKrM projectionrr&r'r)r*r,r.r/rs  zSEWUpsampling.__init__cCsd||}||}|jdkr`|\}}}||j}||j}||||j|}||||}|S)Nr)rYr'rMsizereshape)r)r2bszsrc_lenZ src_embed_dimtgt_lenZ tgt_embed_dimr.r.r/r3s     zSEWUpsampling.forwardr4r.r.r,r/rWs rWcs0eZdZdZfddZddZddZZS)SEWFeatureEncoderz.Construct the features from raw audio waveformcstjdkr@tddgfddtjdD}n6jdkrdfddtjD}ntd jd t||_ d |_ d |_ dS) Ngrouprr+csg|]}t|ddqS)rrb)r.0ir*r.r/ sz.SEWFeatureEncoder.__init__..rlayercsg|]}t|dqS)rb)r9rcrfr.r/rgz`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) rrZfeat_extract_normr?rangeZnum_feat_extract_layers ValueErrorr ModuleList conv_layersgradient_checkpointing_requires_grad)r)r*rmr,rfr/rs      zSEWFeatureEncoder.__init__cCs|D] }d|_qd|_dS)NF) parameters requires_gradror)paramr.r.r/_freeze_parameterss z$SEWFeatureEncoder._freeze_parameterscCs:|dddf}|jr"|jr"d|_|jD] }||}q(|S)NT)rotrainingrqrm)r) input_valuesr2Z conv_layerr.r.r/r3s    zSEWFeatureEncoder.forward)r5r6r7__doc__rrtr3r8r.r.r,r/r`s r`)modulequerykeyvalueattention_maskscalingdropout head_maskc Ks|dur|dd}t||dd|} |dur>| |} tjj| dd} |durj| |dddd} tjj| ||j d} t| |} | dd } | | fS)Nr=rBrrIr)pru) r[torchmatmulr>r functionalsoftmaxviewrru contiguous) ryrzr{r|r}r~rrkwargs attn_weights attn_outputr.r.r/eager_attention_forwards  rc seZdZdZd eeeeeeeedfdd Z d e j ee j ee j ee j eee e ee j ee j eee j fd d d ZZS) SEWAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrxFTN) embed_dim num_headsr is_decoderr is_causalr*cst||_||_||_|||_||_|j||jkrTtd|jd|d|jd|_||_ ||_ t j |||d|_ t j |||d|_t j |||d|_t j |||d|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).r)r)rrrrrhead_dimr*rkr~rrrrXk_projv_projq_projout_proj)r)rrrrrrr*r,r.r/r s&    zSEWAttention.__init__)r2key_value_statesr}layer_head_maskoutput_attentionsrreturncKs|du}|jdd\}} |r(|jdn| } || d|jf} || d|jf} ||j| dd} |rh|n|}||j| dd}||j| dd}t}|jj dkrt |jj }||| |||f|j sdn|j |j ||d|\}}||| d}||}||dfS)z#Input shape: Batch x Time x ChannelNr=rrBeagerrx)rr~rr)shaperrrr>rrrr*_attn_implementationrrurr~r\rr)r)r2rr}rrrZis_cross_attentionr]r_r^Z q_input_shapeZkv_input_shapeZ query_statesZcurrent_statesZ key_statesZ value_statesZattention_interfacerrr.r.r/r3(s:      zSEWAttention.forward)rxFTFN)NNNF)r5r6r7rwintfloatboolrrrrTensorrr tupler3r8r.r.r,r/rs8"rcs$eZdZfddZddZZS)SEWFeedForwardcsptt|j|_t|j|j|_ t |j t rDt |j |_n|j |_t|j|j|_t|j|_dSr0)rrrDropoutZactivation_dropoutintermediate_dropoutrXrKZintermediate_sizeintermediate_dense isinstanceZ hidden_actstrrintermediate_act_fn output_densehidden_dropoutoutput_dropoutrZr,r.r/r_s  zSEWFeedForward.__init__cCs6||}||}||}||}||}|Sr0)rrrrrr1r.r.r/r3ls      zSEWFeedForward.forwardr4r.r.r,r/r^s rcs&eZdZfddZdddZZS)SEWEncoderLayercshtt|j|j|jd|d|_t|j |_ tj |j|j d|_ t||_tj |j|j d|_dS)NF)rrrrr*eps)rrrrKZnum_attention_headsZattention_dropout attentionrrrrr:layer_norm_epsr;r feed_forwardfinal_layer_normrZr,r.r/rws  zSEWEncoderLayer.__init__NFcCsf|}|j|||d\}}}||}||}||}|||}||}|f}|rb||f7}|S)Nr}r)rrr;rr)r)r2r}rZ attn_residualr_outputsr.r.r/r3s     zSEWEncoderLayer.forward)NFr4r.r.r,r/rvs rcs&eZdZfddZdddZZS) SEWEncodercst|_t|_tjj|_tj j j d|_ t j|_tfddtjD|_t|_d|_dS)Nrcsg|] }tqSr.)rrdrrfr.r/rgriz'SEWEncoder.__init__..F)rrr*rApos_conv_embedrZ AvgPool1drMpoolr:rKrr;rrrrlrjnum_hidden_layerslayersrWupsamplernrZr,rfr/rs    zSEWEncoder.__init__NFTc Cs|rdnd}|rdnd}|dur8|ddd|jd}|jjdkrld||<|durfd|vrf|nd}nd||<|d} | |jj} |jd|jj} tj d| | j d dd | jdd} | | ddk}d |ddddddfj |jd }|t|jj}| |jdd|jd|jd}|jd} |dd}||}||}t|d|d}|d d|f|d d|f}|dd}||}||}tpt|}|jD]t}|r||f}tg}|jo||jjk}|r |r"||||d }|d}|r,d }|r||df}q|rT||f}||}|jd| krtj |ddd| |jdf}|st!dd|||fDSt"|||dS)Nr.r=rrBZflash_attention_2rxrdevice?dtype.r)NNcss|]}|dur|VqdSr0r.)rdvr.r.r/ riz%SEWEncoder.forward..Zlast_hidden_stater2 attentions)# unsqueezerepeatrr*rlongsumrMrarangerrexpandtorZfinfominr>rrr[r;rr r rrandruZ layerdroprrrpadrr )r)r2r}routput_hidden_states return_dictZall_hidden_statesZall_self_attentionsZexpand_attention_mask input_lengthsoutput_lengthsZmax_encoder_lengthZ attention_idsZn_input_timestepsZposition_embeddingsZpooled_hidden_statesZ min_lengthZ synced_gpusrhZdropout_probabilityZskip_the_layerZ layer_outputsr.r.r/r3sv        &               zSEWEncoder.forward)NFFTr4r.r.r,r/rs  rc@s`eZdZUeed<dZdZdZdZdZ dZ ddZ e e jefdd d Zee jd d d ZdS)SEWPreTrainedModelr*sewrvTFc Cst|trTtjj|jjddtd|jj d|jj dtj |jj dn,t|tj rz|jjjd|jjdnt|tjtjfr|j j|jjdnt|tjrtrpddl}t|dr*t|d r*|jj|j|jgdd  tj|jjWdn1s0YnD|jj|jdd  tj|jjWdn1sd0Yntj|jjt|tj tjfr|j dur|j jdS) zInitialize the weightsrrBr)meanZstdrxrNrTrSrF)rrArinitZnormal_r%rGmathsqrtrZ in_channelsZ constant_rrXdatar*Zinitializer_ranger:r@Zzero_Zfill_r!r rPrOrQrRrTrSZkaiming_normal_)r)ryrPr.r.r/ _init_weights s.    22 z SEWPreTrainedModel._init_weights)rcCs4dd}t|jj|jjD]\}}||||}q|S)zH Computes the output length of the convolutional layers cSstj|||dddS)Nfloor)Z rounding_moder)rdiv) input_lengthrrr.r.r/_conv_out_length/szMSEWPreTrainedModel._get_feat_extract_output_lengths.._conv_out_length)zipr*r"r#)r)rrrrr.r.r/ _get_feat_extract_output_lengths*sz3SEWPreTrainedModel._get_feat_extract_output_lengths)feature_vector_lengthr}cCs~||dtj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dg d dg }|S)Nr=r)rrrr) rrrrrrzerosrrrflipZcumsumr)r)rr}r batch_sizer.r.r/"_get_feature_vector_attention_mask9s "z5SEWPreTrainedModel._get_feature_vector_attention_maskN)r5r6r7r__annotations__Zbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_supports_flash_attnZ_supports_sdpaZ_supports_flex_attnrrr LongTensorrrrr.r.r.r/rs  r)r mask_prob mask_lengthr} min_masksrcs|\}dkrtdkr6tdddtjdfdd}|durt|d nfd d t|D}tj |ft d }g} |} | d kr|S|D]v} || } tjj t | d| dd} t | d krd}n| d }t| tj| | tjd |g} | | qt| } t| dddddf|| f} | || } t ddddf}t||| f|| }| |} | dkr҈d| | dk<t|| dd |S)an Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for ASR](https://huggingface.co/papers/1904.08779). Note that this method is not optimized to run on TPU and should be run on CPU as part of the preprocessing during training. Args: shape: The shape for which to compute masks. This should be of a tuple of size 2 where the first element is the batch size and the second element is the length of the axis to span. mask_prob: The percentage of the whole axis (between 0 and 1) which will be masked. The number of independently generated mask spans of length `mask_length` is computed by `mask_prob*shape[1]/mask_length`. Note that due to overlaps, `mask_prob` is an upper bound and the actual percentage will be smaller. mask_length: size of the mask min_masks: minimum number of masked spans attention_mask: A (right-padded) attention mask which independently shortens the feature axis of each batch dimension. rz&`mask_length` has to be bigger than 0.zO`mask_length` has to be smaller than `sequence_length`, but got `mask_length`: z and `sequence_length`: `csXt|}t|}|kr2}|d|krTt|dd}|S)z;Given input length, compute how many spans should be maskedrr)rmax)rnum_masked_spanepsilonrrrsequence_lengthr.r/compute_num_masked_spanls  z6_compute_mask_indices..compute_num_masked_spanNr=csg|]}qSr.r.r)rr.r/rgriz)_compute_mask_indices..rrF)replace)rknprandomritemdetachrtolistrjrrchoicerlenZ concatenateonesZint32appendarrayZ broadcast_tor\rZput_along_axis)rrrr}rrrrZ spec_aug_maskZspec_aug_mask_idxsZmax_num_masked_spanrrZspec_aug_mask_idxZdummy_mask_idxoffsetsr.rr/_compute_mask_indicesFs\    rc seZdZedfdd Zd ejeejeejdddZ e d eej eej eejee ee ee e eefdd d ZZS) SEWModelrfcst|||_t||_tj|jd|jd|_ |jd|j k|_ |j rbt |jd|j |_ t|j|_|jdks|jdkrtt|j |_t||_|dS)Nr=rrx)rrr*r`feature_extractorrr:rrr;rKproject_featuresrXfeature_projectionrZfeat_proj_dropoutfeature_dropoutmask_time_probmask_feature_prob ParameterrrZuniform_masked_spec_embedrencoder post_initrZr,r.r/rs   zSEWModel.__init__N)r2mask_time_indicesr}cCst|jdds|S|\}}}|dur<|j|j||<nZ|jjdkr|jrt||f|jj|jj ||jj d}t j ||j t jd}|j|j||<|jjdkr|jrt||f|jj|jj|jjd}t j ||j t jd}|dddfd|d}d||<|S) z Masks extracted features along time axis and/or along feature axis according to [SpecAugment](https://huggingface.co/papers/1904.08779). Zapply_spec_augmentTNr)rrr}r)rr)rrrr=)getattrr*r[r rrrrurZmask_time_lengthZmask_time_min_masksrZtensorrrrZmask_feature_lengthZmask_feature_min_masksr)r)r2r r}rrrKZmask_feature_indicesr.r.r/_mask_hidden_statess4 zSEWModel._mask_hidden_states)rvr}r rrrrc Cs|dur |n|jj}|dur |n|jj}|dur4|n|jj}||}|dd}||}|jrl||}| |}|dur| |j d|}|j ||d}|j |||||d} | d}|s|f| ddSt|| j| jdS)a/ mask_time_indices (`torch.BoolTensor` of shape `(batch_size, sequence_length)`, *optional*): Indices to mask extracted features for contrastive loss. When in training mode, model learns to predict masked extracted features in *config.proj_codevector_dim* space. NrrB)r r}rrrrr)r*rruse_return_dictrr>r;rrrrrrr r r2r) r)rvr}r rrrZextract_featuresr2Zencoder_outputsr.r.r/r3s8     zSEWModel.forward)NN)NNNNN)r5r6r7rrrZ FloatTensorrrrrrrrrr r3r8r.r.r,r/rs. . rzk SEW Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). )Z custom_introc seZdZdeedfdd ZddZddZd d Zd d Z e dee j ee j ee ee ee ee j eeefd ddZZS) SEWForCTCN) target_langcs~t|t||_t|j|_||_|j durFt d|j dt |dr\|j r\|jn|j}t||j |_|dS)a- target_lang (`str`, *optional*): Language id of adapter weights. Adapter weights are stored in the format adapter..safetensors or adapter..bin. Only relevant when using an instance of [`SEWForCTC`] with adapters. Uses 'eng' by default. 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: `SEWForCTC.from_pretrained(..., vocab_size=vocab_size)`. or define `vocab_size` of your model's configuration. add_adapter)rrrrrrZ final_dropoutrr vocab_sizerkr-rOroutput_hidden_sizerKrXlm_headr )r)r*rrr,r.r/rAs    zSEWForCTC.__init__cCsr|j}|dur2t|jdddur2td|dn<|durXt|jdddurXtdn|durn|j|dddS)a' This method overwrites [`~PreTrainedModel.tie_weights`] so that adapter weights can be correctly loaded when passing `target_lang=...` to `from_pretrained(...)`. This method is **not** supposed to be called by the user and is prone to be changed in the future. NZadapter_attn_dimzCannot pass `target_lang`: z- if `config.adapter_attn_dim` is not defined.z)By default `target_lang` is set to 'eng'.T)Z force_load)rrr*rkloggerinfoZ load_adapter)r)rr.r.r/ tie_weights^s  zSEWForCTC.tie_weightscCstdt|dS) Calling this function will disable the gradient computation for the feature encoder so that its parameter will not be updated during training. The method `freeze_feature_extractor` is deprecated and will be removed in Transformers v5. Please use the equivalent `freeze_feature_encoder` method instead.Nwarningswarn FutureWarningfreeze_feature_encoderr)r.r.r/freeze_feature_extractorss z"SEWForCTC.freeze_feature_extractorcCs|jjdSrNrrrtr"r.r.r/r!sz SEWForCTC.freeze_feature_encodercCs|jD] }d|_q dSz Calling this function will disable the gradient computation for the base model so that its parameters will not be updated during training. Only the classification head will be updated. FNrrprqrrr.r.r/freeze_base_modelszSEWForCTC.freeze_base_modelrvr}rrrlabelsrc Cs|dur |n|jj}|dur>||jjkr>td|jj|j|||||d}|d}||}||} d} |dur@|dur|ntj |tj d}| | d tj } |dk} | d} || }tjj| dtjddd}tjjjd d 6tjj||| | |jj|jj|jjd } Wdn1s60Y|sp| f|td}| durl| 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: rrrr=)rIrrF)Zenabled)blankZ reductionZ zero_infinitylosslogitsr2r)r*rrrrkrrrrZ ones_likerrrrZ masked_selectrrZ log_softmaxZfloat32r>backendsZcudnnflagsZctc_lossZ pad_token_idZctc_loss_reductionZctc_zero_infinity_HIDDEN_STATES_START_POSITIONrr2r)r)rvr}rrrr*rr2r.r-rZ labels_maskZtarget_lengthsZflattened_targetsZ log_probsoutputr.r.r/r3sL     &  zSEWForCTC.forward)N)NNNNN)r5r6r7rrrrr#r!r(rrrrrrrr3r8r.r.r,r/r;s(  rz SEW Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like SUPERB Keyword Spotting. c szeZdZfddZddZddZddZed ee j ee j ee ee ee ee j e e efd d d ZZS)SEWForSequenceClassificationcst|t|dr$|jr$tdt||_|jd}|jrTt t |||_ t |j|j|_t |j|j|_|dS)NrzZSequence classification does not support the use of SEW adapters (config.add_adapter=True)r)rrrOrrkrrruse_weighted_layer_sumrr rr layer_weightsrXrKZclassifier_proj_size projector num_labels classifierr )r)r*Z num_layersr,r.r/rs   z%SEWForSequenceClassification.__init__cCstdt|dS)z Calling this function will disable the gradient computation for the feature encoder so that its parameters will not be updated during training. rNrr"r.r.r/r#s z5SEWForSequenceClassification.freeze_feature_extractorcCs|jjdSr$r%r"r.r.r/r!sz3SEWForSequenceClassification.freeze_feature_encodercCs|jD] }d|_q dSr&r'rrr.r.r/r(sz.SEWForSequenceClassification.freeze_base_modelNr)cCs|dur |n|jj}|jjr dn|}|j|||||d}|jjr|t}tj|dd}tjj |j dd} || dddj dd}n|d}| |}|dur|jdd} nV||jd|} | ddd|jd} d || <|j dd| j dd dd} || } d}|dur 1` a classification loss is computed (Cross-Entropy). NTrrrr=rrBrxr,)r*rr4rr1rstackrrrr5rrr6rrrrrr8rr7rr2r)r)rvr}rrrr*rr2Z norm_weightsZ pooled_outputZ padding_maskZexpand_padding_maskr.r-Zloss_fctr2r.r.r/r3 sH     z$SEWForSequenceClassification.forward)NNNNN)r5r6r7rr#r!r(rrrrrrrrr3r8r.r.r,r/r3s&   r3)rr3rr)NrxN)Nr)ArrtypingrrrnumpyrrrZtorch.nnrZ activationsrZintegrations.deepspeedr Zintegrations.fsdpr Zmodeling_flash_attention_utilsr Zmodeling_layersr Zmodeling_outputsr rrZmodeling_utilsrrZprocessing_utilsrrNrrZconfiguration_sewrZ get_loggerr5rrr9r?ModulerArUrWr`rrrrrrrrrrrZndarrayrrr1rr3__all__r.r.r.r/s          + , X$fI  wzs