a hl@sddlZddlZddlmZmZddlZddlZddlm Z ddl m m Z ddlm Z ddlmZddlmZddlmZddlmZdd lmZmZmZmZmZmZdd lmZdd lm Z m!Z!m"Z"d d l#m$Z$e"%e&Z'Gddde j(Z)Gddde j(Z*Gddde j(Z+Gddde j(Z,Gddde j(Z-GdddeZ.GdddeZ/Gddde j(Z0Gddde j(Z1Gd d!d!e j(Z2e Gd"d#d#eZ3Gd$d%d%eZ4Gd&d'd'eZ5Gd(d)d)eZ6Gd*d+d+e j(Z7Gd,d-d-e j(Z8Gd.d/d/e j(Z9dGe:e;e;fed0d1d2Z?eZ@e Gd3d4d4e3ZAd5ZBe d6d7Gd8d9d9e3ZCe d:d7Gd;d<dd>e3ZEGd?d@d@e j(ZFGdAdBdBe j(ZGe dCd7GdDdEdEe3ZHgdFZIdS)HN)OptionalUnion)CrossEntropyLoss)ACT2FN)is_deepspeed_zero3_enabled)is_fsdp_managed_module)GradientCheckpointingLayer)BaseModelOutputCausalLMOutputSequenceClassifierOutputTokenClassifierOutputWav2Vec2BaseModelOutput XVectorOutput)PreTrainedModel)auto_docstringis_peft_availablelogging) WavLMConfigcs$eZdZfddZddZZS)WavLMSamePadLayercs$t|ddkrdnd|_dSNrr)super__init__num_pad_remove)selfnum_conv_pos_embeddings __class__d/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/wavlm/modeling_wavlm.pyr&s zWavLMSamePadLayer.__init__cCs,|jdkr(|ddddd|j f}|S)Nr)rr hidden_statesr r r!forward*s zWavLMSamePadLayer.forward__name__ __module__ __qualname__rr$ __classcell__r r rr!r%s rcs$eZdZfddZddZZS)WavLMPositionalConvEmbeddingcs$ttj|j|j|j|jd|jd|_tjj }t tjj drNtjj j }t rddl }|jj|jjdd"||jddd|_Wdn1s0Yt |jdr|jj jj}|jj jj}n|jj}|jj}|j|||j||n||jddd|_t|j|_t|j|_dS) Nr) kernel_sizepaddinggroups weight_normr)Z modifier_rankweight)namedimparametrizations)rrnnConv1d hidden_sizerZnum_conv_pos_embedding_groupsconvutilsr.hasattrr2r deepspeedzeroZGatheredParametersr/Z original0Z original1weight_gweight_vZregister_external_parameterrr,rfeat_extract_activation activation)rconfigr.r9r;r<rr r!r1s2  0   z%WavLMPositionalConvEmbedding.__init__cCs:|dd}||}||}||}|dd}|SNrr) transposer6r,r>r"r r r!r$Rs      z$WavLMPositionalConvEmbedding.forwardr%r r rr!r*0s !r*cs$eZdZfddZddZZS)WavLMFeatureProjectioncsJttj|jd|jd|_t|jd|j|_ t |j |_ dS)Neps) rrr3 LayerNormconv_dimlayer_norm_eps layer_normLinearr5 projectionDropoutZfeat_proj_dropoutdropoutrr?rr r!r^s zWavLMFeatureProjection.__init__cCs&||}||}||}||fSN)rIrKrM)rr#Znorm_hidden_statesr r r!r$ds   zWavLMFeatureProjection.forwardr%r r rr!rB]s rBc seZdZdZdeeeeeedfdd Zdej e ej e ej ee ej e ej e e ej fd d dZ ej eejejfej eej ej fdddZeeej dddZej ej dddZZS)WavLMAttentionz=Multi-headed attention from 'Attention Is All You Need' paper@ T embed_dim num_headsrM num_buckets max_distancehas_relative_position_biascst||_||_||_|||_|j||jkrNtd|jd|d|jd|_t |||_ t |||_ t |||_ t |||_ ||_||_ttd|jdd|_t |jd|_|rt|j|j|_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).gr)rrrUrVrMZhead_dim ValueErrorZscalingr3rJk_projv_projq_projout_projrWrX Parametertorchonesgru_rel_pos_constgru_rel_pos_linearZ Embeddingrel_attn_embed)rrUrVrMrWrXrYrr r!ros,    zWavLMAttention.__init__NFr)r#attention_mask position_biasoutput_attentionsreturncCs|\}}}|durH|||}|d|ddd||j||}||jdd|jdf} | dddd} || } | | jddd d} t | j ddd\} } | | |j d d } | ||jdd|}|d||f}|||||\}}|||fS) z'Attention layer with relative attentionNrrrCrr)rr1?g@)size compute_bias unsqueezerepeatviewrVshapepermuterdsumraZsigmoidchunkrctorch_multi_head_self_attention)rr#rfrgrhindexZbszZtgt_len_Zgated_hidden_statesZrelative_position_projZgate_aZgate_bZ gate_outputgated_position_bias attn_output attn_weightsr r r!r$s"  $ zWavLMAttention.forward)r#rfryrhricCs|dd}}}|dur&|dnd}d} } d} tj||||j|jtdgt|j j |j j |j j f| | | |j |jj|jj |j|||d|j j|j j|j jd\} } | dd} | dur| dddf| jdd|jf| jdd} | | fS)zCsimple wrapper around torch's multi_head_attention_forward functionrrNFT)Zuse_separate_proj_weightZ q_proj_weightZ k_proj_weightZ v_proj_weight)rAneFZmulti_head_attention_forwardrUrVraemptycatr^biasr\r]rMr_r/training broadcast_torr)rr#rfryrhquerykeyvalueZkey_padding_maskZbias_kZbias_vZ add_zero_attnrzr{r r r!rvsB    "z.WavLMAttention.torch_multi_head_self_attention) query_length key_lengthricCsvtj|tjddddf}tj|tjddddf}||}||}||jjj}||}|gd}|S)Ndtype)rrr) raarangelong_relative_positions_buckettorer/devicers)rrrZcontext_positionZmemory_positionZrelative_positionZrelative_position_bucketvaluesr r r!rns  zWavLMAttention.compute_bias)relative_positionsricCs|jd}|dktj|}t|}|d}||k}t||}|t|j|}|||}||tj}t |t ||d}|t |||7}|Sr) rWrrarabslogfloatmathrXminZ full_likewhere)rrrWZrelative_bucketsZ max_exactZis_smallZrelative_positions_if_largeZrelative_position_if_larger r r!rs   z)WavLMAttention._relative_positions_bucket)rQrRrST)NNFr)r&r'r(__doc__intrboolrraTensorrtupler$ FloatTensorr LongTensorZ BoolTensorrvrnrr)r r rr!rPls@' + 7 rPcs$eZdZfddZddZZS)WavLMFeedForwardcsptt|j|_t|j|j|_ t |j t rDt |j |_n|j |_t|j|j|_t|j|_dSrO)rrr3rLZactivation_dropoutintermediate_dropoutrJr5Zintermediate_sizeintermediate_dense isinstanceZ hidden_actstrrintermediate_act_fn output_densehidden_dropoutoutput_dropoutrNrr r!rs  zWavLMFeedForward.__init__cCs6||}||}||}||}||}|SrO)rrrrrr"r r r!r$ s      zWavLMFeedForward.forwardr%r r rr!rs rcs0eZdZd eedfdd Zd dd ZZS) WavLMEncoderLayerTr?rYcsntt|j|j|j|j|j|d|_t |j |_ t j |j|jd|_t||_t j |j|jd|_dSNrTrDrrrPr5Znum_attention_headsZattention_dropoutrWZmax_bucket_distance attentionr3rLrrMrFrHrIr feed_forwardfinal_layer_normrr?rYrr r!r+s  zWavLMEncoderLayer.__init__NFrc Csl|}|j|||||d\}}}||}||}||}|||}||}||f}|rh||f7}|S)Nrfrgrhrw)rrMrIrr) rr#rfrgrhrw attn_residualr{outputsr r r!r$:s"     zWavLMEncoderLayer.forward)T)NNFrr&r'r(rrrr$r)r r rr!r*srcs0eZdZd eedfdd Zd ddZZS) WavLMEncoderLayerStableLayerNormTrcsntt|j|j|j|j|j|d|_t |j |_ t j |j|jd|_t||_t j |j|jd|_dSrrrrr r!rTs  z)WavLMEncoderLayerStableLayerNorm.__init__NFcCsf|}||}|j||||d\}}}||}||}||||}||f}|rb||f7}|S)N)rfrgrh)rIrrMrr)rr#rfrgrhrr{rr r r!r$cs    z(WavLMEncoderLayerStableLayerNorm.forward)T)NNFrr r rr!rSsrcs&eZdZfddZdddZZS) WavLMEncodercsft|_t|_tjjjd|_ t j |_ t fddtjD|_d|_dS)NrDcsg|]}t|dkdqSr)rY)r.0ir?r r! z)WavLMEncoder.__init__..Frrr?r*pos_conv_embedr3rFr5rHrIrLrrM ModuleListrangenum_hidden_layerslayersgradient_checkpointingrNrrr!rys  zWavLMEncoder.__init__NFTcCsB|rdnd}|rdnd}|durD|ddd|jd}d||<||} || }||}||}tpvt|} d} t|j D]~\} } |r||f}t g}|j o| dko||j jk}|r| r| ||| || d}|dd\}} |rd}|r||df}q|r||f}|s4tdd |||fDSt|||d S) Nr rCrrrrNNNcss|]}|dur|VqdSrOr rvr r r! rz'WavLMEncoder.forward..last_hidden_stater# attentions)rorprrrrIrMrr enumeraterrarandrr? layerdroprr rr#rfrhoutput_hidden_states return_dictZall_hidden_statesZall_self_attentionsZexpand_attention_maskZposition_embeddingsZ synced_gpusrgrlayerZdropout_probabilityZskip_the_layerZ layer_outputsr r r!r$sL         zWavLMEncoder.forward)NFFTr%r r rr!rxs  rcs&eZdZfddZdddZZS) WavLMEncoderStableLayerNormcsft|_t|_tjjjd|_ t j |_ t fddtjD|_d|_dS)NrDcsg|]}t|dkdqSr)rrrr r!rsz8WavLMEncoderStableLayerNorm.__init__..FrrNrrr!rs   z$WavLMEncoderStableLayerNorm.__init__NFTcCs@|rdnd}|rdnd}|durD|ddd|jd}d||<||} || }||}tplt|} d} t|jD]|\} } |r||f}t g}|j o| dko||j j k}|r| r| |||| d}|dd\}} |rd}|r|||df}q|||}|r||f}|s2tdd |||fDSt|||d S) Nr rCrrr)rfrhrgrcss|]}|dur|VqdSrOr rr r r!rrz6WavLMEncoderStableLayerNorm.forward..r)rorprrrrMrrrrrarrr?rrIrr rr r r!r$sF         z#WavLMEncoderStableLayerNorm.forward)NFFTr%r r rr!rs  rcs4eZdZdZfddZeddZddZZS)WavLMGumbelVectorQuantizerz Vector quantization using gumbel softmax. See [CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX](https://huggingface.co/papers/1611.01144) for more information. cst|j|_|j|_|j|jdkrDtd|jd|jdt t d|j|j|j|j|_ t |jd|j|j|_d|_dS)Nrz`config.codevector_dim z5 must be divisible by `config.num_codevector_groups` z for concatenation.rrCr)rrZnum_codevector_groups num_groupsZnum_codevectors_per_groupnum_varsZcodevector_dimr[r3r`rar codevectorsrJrG weight_proj temperaturerNrr r!rs  z#WavLMGumbelVectorQuantizer.__init__cCs8|jdd}ttj|t|ddd }|S)NrrkgHz>rC)meanraexprtr)ZprobsZmarginal_probs perplexityr r r!_compute_perplexity(s (z.WavLMGumbelVectorQuantizer._compute_perplexityc Cs|j\}}}||}||||jd}|jrtjj||j dd}| |}t j ||||jddd}| |}nH|jdd}|j|jd|ddd}||||jd}| |}|||d}|d|j} | |||j|jd} | d||d} | |fS)NrCT)tauhardrkrrl)rrrrqrrr3 functionalZgumbel_softmaxrrZtype_asrasoftmaxrargmaxZ new_zerosZscatter_rorrrt) rr# batch_sizesequence_lengthr5Zcodevector_probsZcodevector_soft_distrZcodevector_idxZcodevectors_per_grouprr r r!r$.s*       z"WavLMGumbelVectorQuantizer.forward) r&r'r(rr staticmethodrr$r)r r rr!r s   rc@sjeZdZUeed<dZdZdZdZdZ dZ ddZ de e jefeed d d Zdee jd d dZdS)WavLMPreTrainedModelr?wavlm input_valuesTFc Cst|tr>|jjjjddd|jjjtj |j njt|t rtj j|j jddtd|j jd|j jdtj |j jdnt|trtd|jj}tj j |jj| |dtj j |jj| |dnt|tjr|jjjd|jjd|jdur|jjnt|tjtjfrN|jj|jjdnZt|tjrtj |j|jdurt|j|j|jd}tj j |j| |ddS) zInitialize the weightsrQr)rstdrr)abNrl)rrrr/dataZnormal_rZzero_r3inituniform_rr*r6rsqrtr+Z in_channelsZ constant_rBrKZ in_featuresrJr?Zinitializer_rangerF GroupNormZfill_r4Zkaiming_normal_r-)rmodulekr r r! _init_weights]s6       z"WavLMPreTrainedModel._init_weightsN) input_lengths add_adaptercCsn|dur|jjn|}dd}t|jj|jjD]\}}||||}q.|rjt|jjD]}||d|jj}qT|S)zH Computes the output length of the convolutional layers NcSstj|||dddS)Nfloor)Z rounding_moder)radiv input_lengthr+strider r r!_conv_out_lengthszOWavLMPreTrainedModel._get_feat_extract_output_lengths.._conv_out_lengthr)r?rzip conv_kernel conv_stridernum_adapter_layersadapter_stride)rrrrr+rrxr r r! _get_feat_extract_output_lengths~sz5WavLMPreTrainedModel._get_feat_extract_output_lengths)feature_vector_lengthrfcCs|jdddddf}|j||d}|tj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dgd dg }|S)NrCrkrr)rrr)r) Zcumsumrrrarrrzerosrrrflipr)rrrfrZnon_padded_lengthsZoutput_lengthsrr r r!"_get_feature_vector_attention_masks  "z7WavLMPreTrainedModel._get_feature_vector_attention_mask)N)N)r&r'r(r__annotations__Zbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_supports_flash_attnZ_supports_sdpaZ_supports_flex_attnrrrarrrrrrr r r r!rSs " rcs&eZdZdfdd ZddZZS)WavLMNoLayerNormConvLayerrcsjt|dkr |j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ dS)Nrrr+rr)rrrG in_conv_dim out_conv_dimr3r4rr conv_biasr6rr=r>rr?layer_idrr r!rs  z"WavLMNoLayerNormConvLayer.__init__cCs||}||}|SrO)r6r>r"r r r!r$s  z!WavLMNoLayerNormConvLayer.forward)rr%r r rr!rsrcs&eZdZdfdd ZddZZS)WavLMLayerNormConvLayerrcs|t|dkr |j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)NrrrT)Zelementwise_affine)rrrGr r r3r4rrr r6rFrIrr=r>r rr r!rs  z WavLMLayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)NrrC)r6rArIr>r"r r r!r$s      zWavLMLayerNormConvLayer.forward)rr%r r rr!rsrcs&eZdZdfdd ZddZZS)WavLMGroupNormConvLayerrcst|dkr |j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)NrrrT)rZ num_channelsZaffine)rrrGr r r3r4rrr r6rr=r>rrIr rr r!rs   z WavLMGroupNormConvLayer.__init__cCs"||}||}||}|SrO)r6rIr>r"r r r!r$s   zWavLMGroupNormConvLayer.forward)rr%r r rr!rsrcs0eZdZdZfddZddZddZZS)WavLMFeatureEncoderz.Construct the features from raw audio waveformcstjdkr@tddgfddtjdD}n6jdkrdfddtjD}ntd jd t||_ d |_ d |_ dS) Ngrouprr csg|]}t|ddqS)rr)rrrr r!rsz0WavLMFeatureEncoder.__init__..rrcsg|]}t|dqS)r)rrrr r!rrz`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) rrZfeat_extract_normrrZnum_feat_extract_layersr[r3r conv_layersr_requires_grad)rr?rrrr!rs      zWavLMFeatureEncoder.__init__cCs|D] }d|_qd|_dS)NF) parameters requires_gradrrparamr r r!_freeze_parameterss z&WavLMFeatureEncoder._freeze_parameterscCs:|dddf}|jr"|jr"d|_|jD] }||}q(|S)NT)rrrr)rrr#Z conv_layerr r r!r$ s    zWavLMFeatureEncoder.forward)r&r'r(rrrr$r)r r rr!rs rcs$eZdZfddZddZZS)WavLMAdapterLayercs0ttj|jd|j|j|jdd|_dS)Nrr)rr,)rrr3r4output_hidden_sizeZadapter_kernel_sizerr6rNrr r!rs zWavLMAdapterLayer.__init__cCs||}tjj|dd}|S)Nrrk)r6r3rZglur"r r r!r$#s zWavLMAdapterLayer.forwardr%r r rr!rs rcs$eZdZfddZddZZS) WavLMAdaptercsptjjkr8tjj|_tj|_n d|_|_t fddt j D|_ j |_ dS)Nc3s|]}tVqdSrO)rrrxrr r!r5rz(WavLMAdapter.__init__..)rrrr5r3rJprojrFproj_layer_normrrrrrrNrrr!r+s    zWavLMAdapter.__init__cCsr|jdur(|jdur(||}||}|dd}|jD]&}tj}|jrX||jkr:||}q:|dd}|Sr@)rrrArnprandomrr)rr#rZlayerdrop_probr r r!r$8s       zWavLMAdapter.forwardr%r r rr!r*s r)rr mask_prob mask_lengthrf min_masksrics|\}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_spanepsilonr#r"r$rr r!compute_num_masked_spanos  z6_compute_mask_indices..compute_num_masked_spanNrCcsg|]}qSr r r)rr r!rrz)_compute_mask_indices..rrF)replace)r[r r!ritemdetachrttolistrrrchoicerlenZ concatenaterbZint32appendarrayrZreshaper&Zput_along_axis)rrr"r#rfr$rr*rZ spec_aug_maskZspec_aug_mask_idxsZmax_num_masked_spanrr'Zspec_aug_mask_idxZdummy_mask_idxoffsetsr r(r!_compute_mask_indicesIs\    r4c seZdZedfdd ZddZddZdeje eje ej d d d Z e de ej e ej e eje ee ee eeeefd d dZZS) WavLMModelrcst|||_t||_t||_|jdks:|jdkrRt t |j |_|jrdt||_n t||_|jr|t|nd|_|dS)NrQ)rrr?rfeature_extractorrBfeature_projectionmask_time_probmask_feature_probr3r`rarr5rmasked_spec_embedZdo_stable_layer_normrencoderrrradapter post_initrNrr r!rs     zWavLMModel.__init__cCstdt|dSz Calling this function will disable the gradient computation for the feature encoder so that its parameters 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_encoderrr r r!freeze_feature_extractors z#WavLMModel.freeze_feature_extractorcCs|jdS Calling this function will disable the gradient computation for the feature encoder so that its parameter will not be updated during training. N)r6rrEr r r!rDsz!WavLMModel.freeze_feature_encoderN)r#mask_time_indicesrfcCst|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)r"r#rfr$)rr)r"r#r$rC)getattrr?rmr:rrr8rr4Zmask_time_lengthZmask_time_min_masksraZtensorrrr9Zmask_feature_lengthZmask_feature_min_masksexpand)rr#rIrfrrr5Zmask_feature_indicesr r r!_mask_hidden_statess4 zWavLMModel._mask_hidden_states)rrfrIrhrrric Cs|dur |n|jj}|dur |n|jj}|dur4|n|jj}||}|dd}|durp|j|jd|dd}||\}}|j |||d}|j |||||d} | d}|j dur| |}|s||f| ddSt ||| j | jd S) 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. NrrFr)rIrfrfrhrrr)rextract_featuresr#r)r?rhruse_return_dictr6rArrrr7rLr;r<WavLMBaseModelOutputr#r) rrrfrIrhrrrNr#Zencoder_outputsr r r!r$s@     zWavLMModel.forward)NN)NNNNN)r&r'r(rrrFrDrarrrrLrrrrrrPr$r)r r rr!r5s2   . r5rzm WavLM 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) WavLMForCTCN) 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 [`WavLMForCTC`] 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: `WavLMForCTC.from_pretrained(..., vocab_size=vocab_size)`. or define `vocab_size` of your model's configuration.r)rrr5rr3rLZ final_dropoutrMrR vocab_sizer[rr8rrr5rJlm_headr=)rr?rRrrr r!r^s    zWavLMForCTC.__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)rRrJr?r[loggerinfoZ load_adapter)rrRr r r! tie_weights{s  zWavLMForCTC.tie_weightscCstdt|dSrHr?Nr@rEr r r!rFs z$WavLMForCTC.freeze_feature_extractorcCs|jjdSrGrr6rrEr r r!rDsz"WavLMForCTC.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. FNrrrrr r r!freeze_base_modelszWavLMForCTC.freeze_base_modelrrfrhrrlabelsric 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: rMrrrC)r1rrF)Zenabled)blankZ reductionZ zero_infinitylosslogitsr#r)r?rOr&rSr[rrMrTraZ ones_likerrrtrZ masked_selectr3rZ log_softmaxZfloat32rAbackendsZcudnnflagsZctc_lossZ pad_token_idZctc_loss_reductionZctc_zero_infinity_HIDDEN_STATES_START_POSITIONr r#r)rrrfrhrrr^rr#rbrarZ labels_maskZtarget_lengthsZflattened_targetsZ log_probsoutputr r r!r$sL     &  zWavLMForCTC.forward)N)NNNNN)r&r'r(rrrrWrFrDr\rrarrrrr r$r)r r rr!rQXs(  rQz WavLM 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)WavLMForSequenceClassificationcst|t|dr$|jr$tdt||_|jd}|jrTt t |||_ t |j|j|_t |j|j|_|dS)Nrz\Sequence classification does not support the use of WavLM adapters (config.add_adapter=True)r)rrr8rr[r5rruse_weighted_layer_sumr3r`rarb layer_weightsrJr5Zclassifier_proj_size projector num_labels classifierr=rr? num_layersrr r!rs   z'WavLMForSequenceClassification.__init__cCstdt|dSr>r@rEr r r!rF s z7WavLMForSequenceClassification.freeze_feature_extractorcCs|jjdSrGrYrEr r r!rDsz5WavLMForSequenceClassification.freeze_feature_encodercCs|jD] }d|_q dSrZr[rr r r!r\sz0WavLMForSequenceClassification.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). NTrMrrkrCrrrQr`)r?rOrhrrerastackr3rrrirqrtrjrrrrrorprlrrkr r#r)rrrfrhrrr^rr# norm_weightsZ pooled_outputZ padding_maskZexpand_padding_maskrbraloss_fctrfr r r!r$&sH     z&WavLMForSequenceClassification.forward)NNNNN)r&r'r(rrFrDr\rrrarrrrr r$r)r r rr!rgs&   rgc szeZdZfddZddZddZddZed ee j ee j ee j ee ee ee e e efd d d ZZS) WavLMForAudioFrameClassificationcszt|t|dr$|jr$tdt||_|jd}|jrTt t |||_ t |j|j|_|j|_|dS)Nrz_Audio frame classification does not support the use of WavLM adapters (config.add_adapter=True)r)rrr8rr[r5rrrhr3r`rarbrirJr5rkrl init_weightsrmrr r!rns   z)WavLMForAudioFrameClassification.__init__cCstdt|dSrXr@rEr r r!rF~s z9WavLMForAudioFrameClassification.freeze_feature_extractorcCs|jjdSrGrYrEr r r!rDsz7WavLMForAudioFrameClassification.freeze_feature_encodercCs|jD] }d|_q dSrZr[rr r r!r\sz2WavLMForAudioFrameClassification.freeze_base_modelN)rrfr^rhrrricCs|dur |n|jj}|jjr dn|}|j|||||d}|jjr|t}tj|dd}tjj |j dd} || dddj dd}n|d}| |} d} |durt} | | d|jtj| d|jdd} |s| f|td} | St| | |j|jd S) roNTrMrrkrCr)Zaxisr`)r?rOrhrrerarpr3rrrirqrtrlrrkrr r#r)rrrfr^rhrrrr#rqrbrarrrfr r r!r$s: (z(WavLMForAudioFrameClassification.forward)NNNNN)r&r'r(rrFrDr\rrrarrrrr r$r)r r rr!rsls&   rscs&eZdZdfdd ZddZZS) AMSoftmaxLoss>@皙?csBt||_||_||_tjt||dd|_ t |_ dS)NT)r) rrscalemarginrkr3r`raZrandnr/rra)rZ input_dimrkrxryrr r!rs  zAMSoftmaxLoss.__init__c Csx|}tjj|jdd}tjj|dd}t||}||j}tj||j }|j t | ||}| ||}|S)Nrrkr)flattenr3r normalizer/rammryZone_hotrkrxrrra) rr#r^r/Z cos_thetapsiZonehotrbrar r r!r$s   zAMSoftmaxLoss.forward)rvrwr%r r rr!rusrucs2eZdZdfdd ZejejdddZZS) TDNNLayerrcsvt|dkr |j|dn|j||_|j||_|j||_|j||_t |j|j|j|_ t |_ dS)Nrr)rrtdnn_dimr r  tdnn_kernelr+Z tdnn_dilationdilationr3rJkernelZReLUr>r rr r!rs "   zTDNNLayer.__init__)r#ricCstrddlm}tr.t|j|r.td|dd}|jj |j |j |j dd}t jj|||jj|jd}|dd}||}|S)Nr) LoraLayerzDetected LoRA on TDNNLayer. LoRA weights won't be applied due to optimization. You should exclude TDNNLayer from LoRA's target modules.rr)r)rZpeft.tuners.lorarrrrArBrAr/rqr r+r r3rZconv1drrr>)rr#rr/r r r!r$s      zTDNNLayer.forward)r)r&r'r(rrarr$r)r r rr!r~s r~zi WavLM Model with an XVector feature extraction head on top for tasks like Speaker Verification. c seZdZfddZddZddZddZeej e fd d d Z e de eje eje ee ee ee ejeeefd ddZZS)WavLMForXVectorcstt|_jd}jr.rCr)rrr5rrrhr3r`rarbrirJr5rrjrr0rtdnnZxvector_output_dimr6rlrurk objectivert)rr?rnZ tdnn_layersrrr!rs    zWavLMForXVector.__init__cCstdt|dSrXr@rEr r r!rF&s z(WavLMForXVector.freeze_feature_extractorcCs|jjdSrGrYrEr r r!rD2sz&WavLMForXVector.freeze_feature_encodercCs|jD] }d|_q dSrZr[rr r r!r\9sz!WavLMForXVector.freeze_base_model)rcCs&dd}|jjD]}|||d}q|S)z? Computes the output length of the TDNN layers cSs|||dS)Nrr rr r r!rFszBWavLMForXVector._get_tdnn_output_lengths.._conv_out_lengthr)r?r)rrrr+r r r!_get_tdnn_output_lengthsAs z(WavLMForXVector._get_tdnn_output_lengthsNr]cCs|dur |n|jj}|jjr dn|}|j|||||d}|jjr|t}tj|dd}tjj |j dd} || dddj dd}n|d}| |}|jD] } | |}q|dur|jdd} |jdd} n||j dd} || }g} g} t|D]D\}}| ||d|fjdd| ||d|fjddqt| } t| } tj| | gdd}||}||}d}|dur|||}|s||f|td}|dur|f|S|St||||j|jdS) roNTrMrrkrCr)rarbZ embeddingsr#r)r?rOrhrrerarpr3rrrirqrtrjrrrrrrr1rr6rlrrr#r)rrrfrhrrr^rr#rqZ tdnn_layerZ mean_featuresZ std_featuresZfeat_extract_output_lengthsZtdnn_output_lengthsrlengthZstatistic_poolingZoutput_embeddingsrbrarfr r r!r$Ps\            zWavLMForXVector.forward)NNNNN)r&r'r(rrFrDr\rrarrrrrrrrrr$r)r r rr!r s(   r)rsrQrgrr5r)Nr)JrrAtypingrrnumpyr raZtorch.nnr3Ztorch.nn.functionalrr}rZ activationsrZintegrations.deepspeedrZintegrations.fsdprZmodeling_layersr Zmodeling_outputsr r r r rrZmodeling_utilsrr7rrrZconfiguration_wavlmrZ get_loggerr&rUModulerr*rBrPrrrrrrrrrrrrrrrrrZndarrayr4rPr5rerQrgrsrur~r__all__r r r r!s           -')%JKFV&#  wsi