o Zh @sddlZddlZddlmZmZmZddlZddlZddl m Z ddl m m Z ddl mZddlmZddlmZddlmZddlmZmZmZmZmZmZdd lmZdd lmZm Z m!Z!d d l"m#Z#e!$e%Z&Gd dde j'Z(Gddde j'Z)Gddde j'Z*Gddde j'Z+Gddde j'Z,Gddde j'Z-Gddde j'Z.Gddde j'Z/Gddde j'Z0Gdd d e j'Z1eGd!d"d"eZ2Gd#d$d$e j'Z3Gd%d&d&e j'Z4Gd'd(d(e j'Z5Gd)d*d*e j'Z6Gd+d,d,e j'Z7Gd-d.d.e j'Z8  dKd/ee9e9fd0e:d1e9d2eej;d3e9d4ejeGd7d8d8e2Z?d9Z@ed:d;Gdd;Gd?d@d@e2ZBeGdAdBdBe2ZCGdCdDdDe j'ZDGdEdFdFe j'ZEedGd;GdHdIdIe2ZFgdJZGdS)LN)OptionalTupleUnion)CrossEntropyLoss)ACT2FN)is_deepspeed_zero3_enabled)is_fsdp_managed_module)BaseModelOutputCausalLMOutputSequenceClassifierOutputTokenClassifierOutputWav2Vec2BaseModelOutput XVectorOutput)PreTrainedModel)auto_docstringis_peft_availablelogging) WavLMConfigc$eZdZfddZddZZS)WavLMSamePadLayercs*t|ddkrd|_dSd|_dSNrr)super__init__num_pad_remove)selfnum_conv_pos_embeddings __class__W/var/www/auris/lib/python3.10/site-packages/transformers/models/wavlm/modeling_wavlm.pyr%s  zWavLMSamePadLayer.__init__cCs,|jdkr|ddddd|j f}|SNr)rr hidden_statesr!r!r"forward)s zWavLMSamePadLayer.forward__name__ __module__ __qualname__rr& __classcell__r!r!rr"r$s rcr)WavLMPositionalConvEmbeddingcs$ttj|j|j|j|jd|jd|_tjj }t tjj dr'tjj j }t r{ddl }|jj|jjdd||jddd|_Wdn1sLwYt |jdrd|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_groupsconvutilsr0hasattrr4r deepspeedzeroZGatheredParametersr1Z original0Z original1weight_gweight_vZregister_external_parameterrr.rfeat_extract_activation activation)rconfigr0r;r=r>rr!r"r0s4     z%WavLMPositionalConvEmbedding.__init__cCs:|dd}||}||}||}|dd}|SNrr) transposer8r.r@r$r!r!r"r&Qs     z$WavLMPositionalConvEmbedding.forwardr'r!r!rr"r,/s !r,cr)WavLMFeatureProjectioncsJttj|jd|jd|_t|jd|j|_ t |j |_ dS)NZeps) rrr5 LayerNormconv_dimlayer_norm_eps layer_normLinearr7 projectionDropoutZfeat_proj_dropoutdropoutrrArr!r"r]s zWavLMFeatureProjection.__init__cCs&||}||}||}||fSN)rJrLrN)rr%Znorm_hidden_statesr!r!r"r&cs   zWavLMFeatureProjection.forwardr'r!r!rr"rD\s rDcseZdZdZ    d"dededed ed ed ef fd d Z    d#dej de ej de ej dede ej e ej e e ej ff ddZ dej deejejfdej dedej ej ff ddZdededej fddZdej dej fd d!ZZS)$WavLMAttentionz=Multi-headed attention from 'Attention Is All You Need' paper@ T embed_dim num_headsrN num_buckets max_distancehas_relative_position_biascst||_||_||_|||_|j||jkr'td|jd|d|jd|_t |||_ t |||_ t |||_ t |||_ ||_||_ttd|jdd|_t |jd|_|rqt|j|j|_dSdS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).gr)rrrUrVrNZhead_dim ValueErrorZscalingr5rKk_projv_projq_projout_projrWrX Parametertorchonesgru_rel_pos_constgru_rel_pos_linearZ Embeddingrel_attn_embed)rrUrVrNrWrXrYrr!r"rns.    zWavLMAttention.__init__NFrr%attention_mask position_biasoutput_attentionsreturncCs|\}}}|dur$|||}|d|ddd||j||}||jdd|jdf} | dddd} || } | | jddd d} t | j ddd\} } | | |j d d } | ||jdd|}|d||f}|||||\}}|||fS) z'Attention layer with relative attentionNrrrErr)rr3?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.forwardrycCs|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} | durz| 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)rCneFZmulti_head_attention_forwardrUrVraemptycatr^biasr\r]rNr_r1training 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_lengthcCsvtj|tjddddf}tj|tjddddf}||}||}||jjj}||}|gd}|S)Ndtype)rrr) raarangelong_relative_positions_buckettorer1devicers)rrrZcontext_positionZmemory_positionZrelative_positionZrelative_position_bucketvaluesr!r!r"rns  zWavLMAttention.compute_biasrelative_positionscCs|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)rRrSrTTNNFr)r(r)r*__doc__intrboolrraTensorrrr& FloatTensorr LongTensorZ BoolTensorrvrnrr+r!r!rr"rQks^' )  7 rQcr)WavLMFeedForwardcsptt|j|_t|j|j|_ t |j t r"t |j |_n|j |_t|j|j|_t|j|_dSrP)rrr5rMZactivation_dropoutintermediate_dropoutrKr7Zintermediate_sizeintermediate_dense isinstanceZ hidden_actstrrintermediate_act_fn output_densehidden_dropoutoutput_dropoutrOrr!r"rs  zWavLMFeedForward.__init__cCs6||}||}||}||}||}|SrP)rrrrrr$r!r!r"r&s     zWavLMFeedForward.forwardr'r!r!rr"r  rcs2eZdZd dedeffdd Zd d d ZZS) WavLMEncoderLayerTrArYcntt|j|j|j|j|j|d|_t |j |_ t j |j|jd|_t||_t j |j|jd|_dSN)rUrVrNrWrXrYrFrrrQr7Znum_attention_headsZattention_dropoutrWZmax_bucket_distance attentionr5rMrrNrGrIrJr feed_forwardfinal_layer_normrrArYrr!r"r*  zWavLMEncoderLayer.__init__NFrc Csl|}|j|||||d\}}}||}||}||}|||}||}||f}|r4||f7}|S)Nrfrgrhrw)rrNrJrr) rr%rfrgrhrw attn_residualr{outputsr!r!r"r&9s"    zWavLMEncoderLayer.forwardTrr(r)r*rrrr&r+r!r!rr"r)rcs2eZdZd dedeffdd Zd dd ZZS) WavLMEncoderLayerStableLayerNormTrArYcrrrrrr!r"rSrz)WavLMEncoderLayerStableLayerNorm.__init__NFcCsf|}||}|j||||d\}}}||}||}||||}||f}|r1||f7}|S)N)rfrgrh)rJrrNrr)rr%rfrgrhrr{rr!r!r"r&bs   z(WavLMEncoderLayerStableLayerNorm.forwardr)NNFrr!r!rr"rRrrc.eZdZfddZ    dddZZS) WavLMEncodercft|_t|_tjjjd|_ t j |_ t fddtjD|_d|_dS)NrFcg|] }t|dkdqSr)rY)r.0irAr!r" sz)WavLMEncoder.__init__..FrrrAr,pos_conv_embedr5rGr7rIrJrMrrN ModuleListrangenum_hidden_layerslayersgradient_checkpointingrOrrr"rxs   zWavLMEncoder.__init__NFTcCs`|rdnd}|r dnd}|dur"|ddd|jd}d||<||} || }||}||}tp;t|} d} t|j D]P\} } |rN||f}t g}|j o_| dko_||j jk}|rd| r|jru|j ru|| j||| |}n | ||| || d}|dd\}} |rd}|r||df}qC|r||f}|stdd |||fDSt|||d S) Nr!rErrrrNNNcs|] }|dur|VqdSrPr!rvr!r!r" z'WavLMEncoder.forward..last_hidden_stater% attentions)rorprrrrJrNrr enumeraterrarandrrA layerdropr_gradient_checkpointing_func__call__tupler rr%rfrhoutput_hidden_states return_dictZall_hidden_statesZall_self_attentionsZexpand_attention_maskZposition_embeddingsZ synced_gpusrgrlayerZdropout_probabilityZskip_the_layerZ layer_outputsr!r!r"r&s^          zWavLMEncoder.forwardNFFTr'r!r!rr"rws rcr) WavLMEncoderStableLayerNormcr)NrFcrr)rrrr!r"rsz8WavLMEncoderStableLayerNorm.__init__..FrrOrrr"rs    z$WavLMEncoderStableLayerNorm.__init__NFTcCs^|rdnd}|r dnd}|dur"|ddd|jd}d||<||} || }||}tp6t|} d} t|jD]O\} } |rI||f}t g}|j oZ| dkoZ||j j k}|r_| r|jrp|j rp|| j||| |}n| |||| d}|dd\}} |rd}|r||df}q>||}|r||f}|stdd |||fDSt|||d S) Nr!rErrr)rfrhrgrcsrrPr!rr!r!r"rrz6WavLMEncoderStableLayerNorm.forward..r)rorprrrrNrr rrrarrrArrrrrJrr rr!r!r"r&sX          z#WavLMEncoderStableLayerNorm.forwardrr'r!r!rr"rs rcs4eZdZdZfddZeddZddZZS)WavLMGumbelVectorQuantizerz Vector quantization using gumbel softmax. See [CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX](https://arxiv.org/pdf/1611.01144.pdf) for more information. cst|j|_|j|_|j|jdkr"td|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.rrEr)rrZnum_codevector_groups num_groupsZnum_codevectors_per_groupnum_varsZcodevector_dimr[r5r`rar codevectorsrKrH weight_proj temperaturerOrr!r"r$s   z#WavLMGumbelVectorQuantizer.__init__cCs8|jdd}ttj|t|ddd }|S)NrrkgHz>rE)meanraexprtr)ZprobsZmarginal_probs perplexityr!r!r"_compute_perplexity9s (z.WavLMGumbelVectorQuantizer._compute_perplexityc Cs|j\}}}||}||||jd}|jrAtjj||j dd}| |}t j ||||jddd}| |}n$|jdd}|j|jd|ddd}||||jd}| |}|||d}|d|j} | |||j|jd} | d||d} | |fS)NrET)tauZhardrkrrl)rrrrqrrr5 functionalZgumbel_softmaxrrZtype_asrasoftmaxrargmaxZ new_zerosZscatter_rorrrt) rr% batch_sizesequence_lengthr7Zcodevector_probsZcodevector_soft_distrZcodevector_idxZcodevectors_per_grouprr!r!r"r&?s*       z"WavLMGumbelVectorQuantizer.forward) r(r)r*rr staticmethodrr&r+r!r!rr"rs   rc@sheZdZeZdZdZdZdZdZ ddZ dde e j efd eefd d Z dd ed e j fddZdS)WavLMPreTrainedModelwavlm input_valuesTFc Cst|tr|jjjjddd|jjjtj |j dSt|t rItj j|j jddtd|j jd|j jdtj |j jddSt|trqtd|jj}tj j |jj| |dtj j |jj| |ddSt|tjr|jjjd|jjd|jdur|jjdSdSt|tjtjfr|jj|jjddSt|tjrtj |j|jdurt|j|j|jd}tj j |j| |ddSdSdS) zInitialize the weightsrRr)rstdrr)abNrl)rrrr1dataZnormal_rZzero_r5inituniform_rr,r8rsqrtr-Z in_channelsZ constant_rDrLZ in_featuresrKrAZinitializer_rangerG GroupNormZfill_r6Zkaiming_normal_r/)rmodulekr!r!r" _init_weightsms<         z"WavLMPreTrainedModel._init_weightsN input_lengths add_adaptercCsn|dur|jjn|}dd}t|jj|jjD] \}}||||}q|r5t|jjD] }||d|jj}q*|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)rArzip conv_kernel conv_stridernum_adapter_layersadapter_stride)rrrrr-rrxr!r!r" _get_feat_extract_output_lengthssz5WavLMPreTrainedModel._get_feat_extract_output_lengthsfeature_vector_lengthrfcCs|jdddddf}|j||d}|tj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dgd dg }|S)NrErkrr)rrr)r) Zcumsumr rrarrrzerosrrrflipr)rr rfrZnon_padded_lengthsZoutput_lengthsrr!r!r""_get_feature_vector_attention_masks  "z7WavLMPreTrainedModel._get_feature_vector_attention_maskrP)r(r)r*rZ config_classZbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_supports_flash_attn_2Z_supports_sdparrrarrrrr rr!r!r!r"rds("  rc&eZdZdfdd ZddZZS)WavLMNoLayerNormConvLayerrcsjt|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ dS)Nrrr-rr)rrrH in_conv_dim out_conv_dimr5r6rr conv_biasr8rr?r@rrAlayer_idrr!r"rs  z"WavLMNoLayerNormConvLayer.__init__cCs||}||}|SrP)r8r@r$r!r!r"r&s  z!WavLMNoLayerNormConvLayer.forwardrr'r!r!rr"rsrcr)WavLMLayerNormConvLayerrcs|t|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)NrrrT)Zelementwise_affine)rrrHrrr5r6rrrr8rGrJrr?r@rrr!r"rs  z WavLMLayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)NrrE)r8rCrJr@r$r!r!r"r&s     zWavLMLayerNormConvLayer.forwardrr'r!r!rr"rsrcr)WavLMGroupNormConvLayerrcst|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)NrrrT)rZ num_channelsZaffine)rrrHrrr5r6rrrr8rr?r@rrJrrr!r"rs   z WavLMGroupNormConvLayer.__init__cCs"||}||}||}|SrP)r8rJr@r$r!r!r"r&s   zWavLMGroupNormConvLayer.forwardrr'r!r!rr"rsrcs0eZdZdZfddZddZddZZS)WavLMFeatureEncoderz.Construct the features from raw audio waveformcstjdkr tddgfddtjdD}njdkr2fddtjD}n td jd t||_ d |_ d |_ dS) Ngrouprrcsg|] }t|ddqS)rr)rrrr!r"r sz0WavLMFeatureEncoder.__init__..rrcsg|]}t|dqS)r)rrrr!r"r sz`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) rrZfeat_extract_normrrZnum_feat_extract_layersr[r5r conv_layersr_requires_grad)rrArrrr"rs       zWavLMFeatureEncoder.__init__cCs|D]}d|_qd|_dS)NF) parameters requires_gradr rparamr!r!r"_freeze_parameterss  z&WavLMFeatureEncoder._freeze_parameterscCs\|dddf}|jr|jrd|_|jD]}|jr'|jr'|jr'||j|}q||}q|S)NT)r rr"rrrr)rrr%Z conv_layerr!r!r"r&s   zWavLMFeatureEncoder.forward)r(r)r*rrr%r&r+r!r!rr"rs  rcr)WavLMAdapterLayercs0ttj|jd|j|j|jdd|_dS)Nrr)rr.)rrr5r6output_hidden_sizeZadapter_kernel_sizer r8rOrr!r"r/s  zWavLMAdapterLayer.__init__cCs||}tjj|dd}|S)Nrrk)r8r5rZglur$r!r!r"r&9s zWavLMAdapterLayer.forwardr'r!r!rr"r&.s  r&cr) WavLMAdaptercsptjjkrtjj|_tj|_nd|_|_t fddt j D|_ j |_ dS)Nc3s|]}tVqdSrP)r&rrxrr!r"rKsz(WavLMAdapter.__init__..)rrr'r7r5rKprojrGproj_layer_normrrr rrrOrrr"rAs     zWavLMAdapter.__init__cCsr|jdur|jdur||}||}|dd}|jD]}tj}|jr,||jkr0||}q|dd}|SrB)r*r+rCrnprandomrr)rr%rZlayerdrop_probr!r!r"r&Ns      zWavLMAdapter.forwardr'r!r!rr"r(@rr(rr mask_prob mask_lengthrf min_masksrics|\}dkr tdkrtdddtjdfdd}|dur:|d n fd d t|D}tj |ft d }g} |} | d krZ|S|D];} || } tjj t | d| dd} t | d kr}d}n| d }t| tj| | tjd |g} | | q\t| } t| dddddf|| f} | || } t ddddf}t||| f|| }| |} | dkrd| | dk<t|| dd |S)af Computes random mask spans for a given shape. Used to implement [SpecAugment: A Simple Data Augmentation Method for ASR](https://arxiv.org/abs/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|}|kr}|d|kr*t|dd}|S)z;Given input length, compute how many spans should be maskedrr)rmax)rnum_masked_spanepsilonr/r.r0rr!r"compute_num_masked_spans  z6_compute_mask_indices..compute_num_masked_spanNrEcsg|]}qSr!r!r))rr!r"rsz)_compute_mask_indices..rrF)replace)r[r,r-ritemdetachrttolistrrrchoicerlenZ concatenaterbZint32appendarrayrZreshaper2Zput_along_axis)rrr.r/rfr0rr6rZ spec_aug_maskZspec_aug_mask_idxsZmax_num_masked_spanrr3Zspec_aug_mask_idxZdummy_mask_idxoffsetsr!r4r"_compute_mask_indices_s\    r@cseZdZdeffdd ZddZddZ  dd ejd e ejd e ej fd d Z e     dde ej d e ej d e ejde ede ede edeeeffddZZS) WavLMModelrAcst|||_t||_t||_|jdks|jdkr)t t |j |_|jr2t||_nt||_|jr>t|nd|_|dS)NrR)rrrArfeature_extractorrDfeature_projectionmask_time_probmask_feature_probr5r`rarr7rmasked_spec_embedZdo_stable_layer_normrencoderrrr(adapter post_initrOrr!r"rs      zWavLMModel.__init__cCtdt|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_extractor  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)rBr%rRr!r!r"rQsz!WavLMModel.freeze_feature_encoderNr%mask_time_indicesrfcCst|jdds |S|\}}}|dur|j|j||<n-|jjdkrK|jrKt||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://arxiv.org/abs/1904.08779). Zapply_spec_augmentTNr)r.r/rfr0)rr)r.r/r0rE)getattrrArmrFrrrDrr@Zmask_time_lengthZmask_time_min_masksraZtensorrrrEZmask_feature_lengthZmask_feature_min_masksexpand)rr%rWrfrrr7Zmask_feature_indicesr!r!r"_mask_hidden_statess4 zWavLMModel._mask_hidden_statesrrhrrric Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}||}|dd}|dur8|j|jd|dd}||\}}|j |||d}|j |||||d} | d}|j dur_| |}|sk||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 )rWrfrfrhrrr)rextract_featuresr%r)rArhruse_return_dictrBrCrrrrCrZrGrHWavLMBaseModelOutputr%r) rrrfrWrhrrr\r%Zencoder_outputsr!r!r"r&0s@     zWavLMModel.forward)NNNNNNN)r(r)r*rrrSrQrarrrrZrrrrrr^r&r+r!r!rr"rAsD   . rArzm WavLM Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). )Z custom_introcseZdZddeeffdd ZddZddZd d Zd d Z e     dd ee j dee j dee dee dee dee j deeeffddZZS) WavLMForCTCN target_langcs~t|t||_t|j|_||_|j dur#t 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)rrrArr5rMZ final_dropoutrNra vocab_sizer[r r:rr'r7rKlm_headrI)rrArar'rr!r"rts     zWavLMForCTC.__init__cCsv|j}|durt|jdddurtd|d|dur,t|jdddur,tddS|dur9|j|dddSdS)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)rarXrAr[loggerinfoZ load_adapter)rrar!r!r" tie_weightss zWavLMForCTC.tie_weightscCrJrVrLNrMrRr!r!r"rSrTz$WavLMForCTC.freeze_feature_extractorcC|jjdSrUrrBr%rRr!r!r"rQz"WavLMForCTC.freeze_feature_encodercC|jD]}d|_qdSz 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. FNrr!r"r#r!r!r"freeze_base_modelzWavLMForCTC.freeze_base_modelrrfrhrrlabelsric Cs||dur|n|jj}|dur||jjkrtd|jj|j|||||d}|d}||}||} d} |dur|durC|ntj |tj 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: r[rrrE)r3rrF)enabled)blankZ reductionZ zero_infinitylosslogitsr%r)rAr]r2rbr[rrNrcraZ ones_likerr rtrZ masked_selectr5rZ log_softmaxZfloat32rCbackendsZcudnnflagsZctc_lossZ pad_token_idZctc_loss_reductionZctc_zero_infinity_HIDDEN_STATES_START_POSITIONr r%r)rrrfrhrrrprr%rurtrZ labels_maskZtarget_lengthsZflattened_targetsZ log_probsoutputr!r!r"r&sN      zWavLMForCTC.forwardrPr_)r(r)r*rrrrfrSrQrnrrarrrrr r&r+r!r!rr"r`ns6  r`z WavLM Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like SUPERB Keyword Spotting. cseZdZfddZddZddZddZe dd ee j d ee j d ee d ee dee dee j de e effddZZS)WavLMForSequenceClassificationcst|t|dr|jrtdt||_|jd}|jr*t t |||_ t |j|j|_t |j|j|_|dS)Nrz\Sequence classification does not support the use of WavLM adapters (config.add_adapter=True)r)rrr:rr[rArruse_weighted_layer_sumr5r`rarb layer_weightsrKr7Zclassifier_proj_size projector num_labels classifierrIrrA num_layersrr!r"rs    z'WavLMForSequenceClassification.__init__cCrJrKrMrRr!r!r"rS!rTz7WavLMForSequenceClassification.freeze_feature_extractorcCrhrUrirRr!r!r"rQ-rjz5WavLMForSequenceClassification.freeze_feature_encodercCrkrlrmr#r!r!r"rn4roz0WavLMForSequenceClassification.freeze_base_modelNrrfrhrrrpricCsz|dur|n|jj}|jjrdn|}|j|||||d}|jjrB|t}tj|dd}tjj |j dd} || dddj dd}n|d}| |}|durV|jdd} n+||jd|} | ddd|jd} d || <|j dd| j dd dd} || } d}|durt}|| d|jj| d}|s| f|td}|dur|f|S|St|| |j|jd S)  input_values (`torch.FloatTensor` of shape `(batch_size, sequence_length)`): Float values of input raw speech waveform. Values can be obtained by loading a `.flac` or `.wav` audio file into an array of type `List[float]` or a `numpy.ndarray`, *e.g.* via the soundfile library (`pip install soundfile`). To prepare the array into `input_values`, the [`AutoProcessor`] should be used for padding and conversion into a tensor of type `torch.FloatTensor`. See [`WavLMProcessor.__call__`] for details. labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): Labels for computing the sequence classification/regression loss. Indices should be in `[0, ..., config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If `config.num_labels > 1` a classification loss is computed (Cross-Entropy). NTr[rrkrErrrRrs)rAr]r{rrxrastackr5rrr|rqrtr}rrrrrorprrr~r r%r)rrrfrhrrrprr% norm_weightsZ pooled_outputZ padding_maskZexpand_padding_maskrurtloss_fctryr!r!r"r&<sH    z&WavLMForSequenceClassification.forwardr_)r(r)r*rrSrQrnrrrarrrrr r&r+r!r!rr"rz s4   rzcseZdZfddZddZddZddZe dd ee j d ee j d ee j d ee dee dee de e effddZZS) WavLMForAudioFrameClassificationcszt|t|dr|jrtdt||_|jd}|jr*t t |||_ t |j|j|_|j|_|dS)Nrz_Audio frame classification does not support the use of WavLM adapters (config.add_adapter=True)r)rrr:rr[rArrr{r5r`rarbr|rKr7r~r init_weightsrrr!r"rs    z)WavLMForAudioFrameClassification.__init__cCrJrgrMrRr!r!r"rSrTz9WavLMForAudioFrameClassification.freeze_feature_extractorcCrhrUrirRr!r!r"rQrjz7WavLMForAudioFrameClassification.freeze_feature_encodercCrkrlrmr#r!r!r"rnroz2WavLMForAudioFrameClassification.freeze_base_modelNrrfrprhrrricCs|dur|n|jj}|jjrdn|}|j|||||d}|jjrB|t}tj|dd}tjj |j dd} || dddj dd}n|d}| |} d} |durht} | | d|jtj| d|jdd} |su| f|td} | St| | |j|jd S) rNTr[rrkrEr)Zaxisrs)rAr]r{rrxrarr5rrr|rqrtrrr~rr r%r)rrrfrprhrrrr%rrurtrryr!r!r"r&s: (z(WavLMForAudioFrameClassification.forwardr_)r(r)r*rrSrQrnrrrarrrrr r&r+r!r!rr"rs4   rcs&eZdZdfdd ZddZZS) AMSoftmaxLoss>@皙?csFtt|||_||_||_tjt ||dd|_ t |_ dS)NT)r") rrrscalemarginr~r5r`raZrandnr1rrt)rZ input_dimr~rrrr!r"rs zAMSoftmaxLoss.__init__c Csx|}tjj|jdd}tjj|dd}t||}||j}tj||j }|j t | ||}| ||}|S)Nrrkr)flattenr5r normalizer1rammrZone_hotr~rrrrt) rr%rpr1Z cos_thetapsiZonehotrurtr!r!r"r&s   zAMSoftmaxLoss.forward)rrr'r!r!rr"rsrcs4eZdZdfdd ZdejdejfddZZS) TDNNLayerrcsvt|dkr|j|dn|j||_|j||_|j||_|j||_t |j|j|j|_ t |_ dS)Nrr)rrtdnn_dimrr tdnn_kernelr-Z tdnn_dilationdilationr5rKkernelZReLUr@rrr!r"rs "   zTDNNLayer.__init__r%ricCstr ddlm}trt|j|rtd|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.lorarrrrNrOrCr1rqrr-rr5rZconv1drrr@)rr%rr1r!r!r"r& s      zTDNNLayer.forwardr)r(r)r*rrarr&r+r!r!rr"rs rzi WavLM Model with an XVector feature extraction head on top for tasks like Speaker Verification. cseZdZfddZddZddZddZd eej e ffd d Z e dd e ejde ejde ede ede ede ejdeeeffddZZS)WavLMForXVectorcstt|_jd}jrtt |||_ t j j d|_fddttj D}t||_t j ddj|_t jj|_tjj|_|dS)Nrrcsg|]}t|qSr!)rrrr!r"r0sz,WavLMForXVector.__init__..rEr)rrrArrr{r5r`rarbr|rKr7rr}rr<rtdnnZxvector_output_dimrBrrr~ objectiver)rrArZ tdnn_layersrrr"r's     zWavLMForXVector.__init__cCrJrgrMrRr!r!r"rS:rTz(WavLMForXVector.freeze_feature_extractorcCrhrUrirRr!r!r"rQFrjz&WavLMForXVector.freeze_feature_encodercCrkrlrmr#r!r!r"rnMroz!WavLMForXVector.freeze_base_modelrcCs&dd}|jjD]}|||d}q|S)z? Computes the output length of the TDNN layers cSs|||dS)Nrr!rr!r!r"rZszBWavLMForXVector._get_tdnn_output_lengths.._conv_out_lengthr)rAr)rrrr-r!r!r"_get_tdnn_output_lengthsUs z(WavLMForXVector._get_tdnn_output_lengthsNrrfrhrrrpricCs|dur|n|jj}|jjrdn|}|j|||||d}|jjrB|t}tj|dd}tjj |j dd} || dddj dd}n|d}| |}|jD]} | |}qN|durf|jdd} |jdd} nC||j dd} || }g} g} t|D]"\}}| ||d|fjdd| ||d|fjddq|t| } t| } tj| | gdd}||}||}d}|dur|||}|s||f|td}|dur|f|S|St||||j|jdS) rNTr[rrkrEr)rtruZ embeddingsr%r)rAr]r{rrxrarr5rrr|rqrtr}rrrr rrr=rrBrrrr%r)rrrfrhrrrprr%rZ tdnn_layerZ mean_featuresZ std_featuresZfeat_extract_output_lengthsZtdnn_output_lengthsrlengthZstatistic_poolingZoutput_embeddingsrurtryr!r!r"r&ds\           zWavLMForXVector.forwardr_)r(r)r*rrSrQrnrrarrrrrrrrrr&r+r!r!rr"r!s6   r)rr`rzrrArr#)HrrNtypingrrrnumpyr,raZtorch.nnr5Ztorch.nn.functionalrr}rZ activationsrZintegrations.deepspeedrZintegrations.fsdpr Zmodeling_outputsr r r r rrZmodeling_utilsrr9rrrZconfiguration_wavlmrZ get_loggerr(rdModulerr,rDrQrrrrrrrrrrrr&r(rrrZndarrayr@r^rArxr`rzrrrr__all__r!r!r!r"s          -')%STFU,#  wrh