o ۷iQ@sddlZddlZddlmZmZmZddlZddlZddlm 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.   dEd e j*d!ej/d"ej/d#ej/d$eej/d%ee0d&e0d'eej/fd(d)Z1Gd*d+d+e j*Z2Gd,d-d-e j*Z3Gd.d/d/eZ4Gd0d1d1e j*Z5e Gd2d3d3eZ6  dFd4e7e8e8fd5e0d6e8d$eej9d7e8d8ej:f d9d:Z;e Gd;d<dGd?d@d@e6Z>e dAd>GdBdCdCe6Z?gdDZ@dS)GN)CallableOptionalUnion)nn)CrossEntropyLoss)ACT2FN)is_deepspeed_zero3_enabled)is_fsdp_managed_module)FlashAttentionKwargs)GradientCheckpointingLayer)BaseModelOutputCausalLMOutputSequenceClassifierOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringlogging) SEWConfigc&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__Z/home/ubuntu/vllm_env/lib/python3.10/site-packages/transformers/models/sew/modeling_sew.pyr/s  z SEWNoLayerNormConvLayer.__init__cCs||}||}|SN)r&r(r* hidden_statesr/r/r0forward=s  zSEWNoLayerNormConvLayer.forwardr__name__ __module__ __qualname__rr4 __classcell__r/r/r-r0r.srcr)SEWLayerNormConvLayerrcs|t|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)NrrrT)elementwise_affine)rrrr r!rr"r#r$r%r& LayerNorm layer_normrr'r(r)r-r/r0rDs  zSEWLayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)N)r& transposer>r(r2r/r/r0r4Ss     zSEWLayerNormConvLayer.forwardr5r6r/r/r-r0r;Csr;cr)SEWGroupNormConvLayerrcst|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)NrrrT) num_groups num_channelsaffine)rrrr r!rr"r#r$r%r&rr'r( GroupNormr>r)r-r/r0r_s   zSEWGroupNormConvLayer.__init__cC"||}||}||}|Sr1)r&r>r(r2r/r/r0r4os   zSEWGroupNormConvLayer.forwardr5r6r/r/r-r0rB^srBc$eZdZfddZddZZS)SEWPositionalConvEmbeddingcs(ttj|j|j|j|jd|j|jd|_tj j }t tj j dr)tj j j }t r}ddl}|jj|jjdd||jddd|_Wdn1sNwYt |jdrf|jj jj}|jj jj}n|jj}|jj}|j|||j||n ||jddd|_t|j|_t|j|_dS) N)rpaddinggroupsr weight_normr modifier_rankweight)namedimparametrizations)rrrr" hidden_sizenum_conv_pos_embeddingsnum_conv_pos_embedding_groupssqueeze_factorr&utilsrMhasattrrSr deepspeedzeroGatheredParametersrP original0 original1weight_gweight_vregister_external_parameterSEWSamePadLayerrKrr'r()r*r+rMrZr_r`r-r/r0rws6      z#SEWPositionalConvEmbedding.__init__cCrGr1)r&rKr(r2r/r/r0r4s   z"SEWPositionalConvEmbedding.forwardr6r/r/r-r0rIvs "rIcrH)rbcs*t|ddkrd|_dSd|_dS)NrJrr)rrnum_pad_remove)r*rUr-r/r0rs  zSEWSamePadLayer.__init__cCs,|jdkr|ddddd|j f}|SNr)rcr2r/r/r0r4s zSEWSamePadLayer.forwardr6r/r/r-r0rbs rbcrH) SEWUpsamplingcs:tt|j|j|j|_t|j|_ |j|_dSr1) rrrLinearrTrW projectionrr'r(r*r+r-r/r0rs   zSEWUpsampling.__init__cCsd||}||}|jdkr0|\}}}||j}||j}||||j|}||||}|S)Nr)rgr(rWsizereshape)r*r3bszsrc_len src_embed_dimtgt_len tgt_embed_dimr/r/r0r4s     zSEWUpsampling.forwardr6r/r/r-r0res recs0eZdZdZfddZddZddZZS)SEWFeatureEncoderz.Construct the features from raw audio waveformcstjdkr tddgfddtjdD}njdkr2fddtjD}n td jd t||_ d |_ d |_ dS) Ngrouprr,csg|] }t|ddqS)rrr)r.0ir+r/r0 sz.SEWFeatureEncoder.__init__..rlayercsg|]}t|dqS)rr)r;rsrvr/r0rwsz`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) rrfeat_extract_normrBrangenum_feat_extract_layers ValueErrorr ModuleList conv_layersgradient_checkpointing_requires_grad)r*r+r~r-rvr0rs       zSEWFeatureEncoder.__init__cCs|D]}d|_qd|_dSNF) parameters requires_gradrr*paramr/r/r0_freeze_parameterss  z$SEWFeatureEncoder._freeze_parameterscCs:|dddf}|jr|jrd|_|jD]}||}q|S)NT)rtrainingrr~)r* input_valuesr3 conv_layerr/r/r0r4s    zSEWFeatureEncoder.forward)r7r8r9__doc__rrr4r:r/r/r-r0rps  rpmodulequerykeyvalueattention_maskscalingdropout head_maskc Ks|dur |dd}t||dd|} |dur| |} tjj| dd} |dur5| |dddd} tjj| ||j d} t| |} | dd } | | fS)Nr@rJrrRr)pr) ritorchmatmulrAr functionalsoftmaxviewrr contiguous) rrrrrrrrkwargs attn_weights attn_outputr/r/r0eager_attention_forwards  rcseZdZdZ     ddededed ed ed ed eeffd d Z    dde j dee j dee j dee j deede e dee j ee j eee j ffddZZS) SEWAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN embed_dim num_headsr is_decoderr is_causalr+cst||_||_||_|||_||_|j||jkr*td|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+r|rrrrrfk_projv_projq_projout_proj)r*rrrrrrr+r-r/r0r s&    zSEWAttention.__init__r3key_value_statesrlayer_head_maskoutput_attentionsrreturncKs|du}|jdd\}} |r|jdn| } || d|jf} || d|jf} ||j| dd} |r4|n|}||j| dd}||j| dd}t}|jj dkr\t |jj }||| |||f|j shdn|j |j ||d|\}}||| d}||}||dfS)z#Input shape: Batch x Time x ChannelNr@rrJeagerr)rrrr)shaperrrrArrrr+_attn_implementationrrrrrjrr)r*r3rrrrris_cross_attentionrkrnrl q_input_shapekv_input_shape query_statescurrent_states key_states value_statesattention_interfacerrr/r/r0r4(s:      zSEWAttention.forward)rFTFN)NNNF)r7r8r9rintfloatboolrrrrTensorrr tupler4r:r/r/r-r0rsR"  rcrH)SEWFeedForwardcsptt|j|_t|j|j|_ t |j t r"t |j |_n|j |_t|j|j|_t|j|_dSr1)rrrDropoutactivation_dropoutintermediate_dropoutrfrTintermediate_sizeintermediate_dense isinstance hidden_actstrrintermediate_act_fn output_densehidden_dropoutoutput_dropoutrhr-r/r0r_s  zSEWFeedForward.__init__cCs6||}||}||}||}||}|Sr1)rrrrrr2r/r/r0r4ls     zSEWFeedForward.forwardr6r/r/r-r0r^s  rcs&eZdZfddZdddZZS)SEWEncoderLayercshtt|j|j|jd|d|_t|j |_ tj |j|j d|_ t||_tj |j|j d|_dS)NF)rrrrr+eps)rrrrTnum_attention_headsattention_dropout attentionrrrrr=layer_norm_epsr>r feed_forwardfinal_layer_normrhr-r/r0rws  zSEWEncoderLayer.__init__NFcCsf|}|j|||d\}}}||}||}||}|||}||}|f}|r1||f7}|S)Nrr)rrr>rr)r*r3rr attn_residualr_outputsr/r/r0r4s    zSEWEncoderLayer.forwardrr6r/r/r-r0rvs rcs.eZdZfddZ    dddZZS) SEWEncodercst|_t|_tjj|_tj j j d|_ t j|_tfddtjD|_t|_d|_dS)Nrcsg|]}tqSr/)rrtrrvr/r0rwsz'SEWEncoder.__init__..F)rrr+rIpos_conv_embedr AvgPool1drWpoolr=rTrr>rrrr}rznum_hidden_layerslayersreupsamplerrhr-rvr0rs     zSEWEncoder.__init__NFTc Cs|rdnd}|r dnd}|dur|ddd|jd}|jjdkr5d||<|dur2d|vr2|nd}nfd||<|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]7}|r||f}tg}|jo||jjk}|r|r ||||d }|d}|rd }|r||df}q|r%||f}||}|jd| krBtj |ddd| |jdf}|sQt!dd|||fDSt"|||dS)Nr/r@rrJflash_attention_2rrdevice?dtype.rNNcss|] }|dur|VqdSr1r/)rtvr/r/r0 sz%SEWEncoder.forward..last_hidden_stater3 attentions)# unsqueezerepeatrr+rlongsumrWrarangerrexpandtorfinfominrArrrir>rr r rrandr layerdroprrrpadrr )r*r3rroutput_hidden_states return_dictall_hidden_statesall_self_attentionsexpand_attention_mask input_lengthsoutput_lengthsmax_encoder_length attention_idsn_input_timestepsposition_embeddingspooled_hidden_states min_length synced_gpusrxdropout_probabilityskip_the_layer layer_outputsr/r/r0r4st      &               zSEWEncoder.forward)NFFTr6r/r/r-r0rs rc@sbeZdZUeed<dZdZdZdZdZ dZ ddZ de e jeffd d Zd ed e jfd dZdS)SEWPreTrainedModelr+sewrTFc Cst|tr)tjj|jjddtd|jj d|jj dtj |jj dnt|tj r;|jjjd|jjdn}t|tjtjfrR|j j|jjdnft|tjrtrddl}t|drt|d r|jj|j|jgdd tj|jjWdn1swYn*|jj|jdd tj|jjWdn1swYntj|jjt|tj tjfr|j dur|j jdSdSdS) zInitialize the weightsrrJr)meanstdrrNr`r_rN)rrIrinitnormal_r&rPmathsqrtr in_channels constant_rrfdatar+initializer_ranger=rFzero_fill_r"r rZrYr[r\r`r_kaiming_normal_)r*rrZr/r/r0 _init_weights s8     z SEWPreTrainedModel._init_weightsrcCs4dd}t|jj|jjD] \}}||||}q |S)zH Computes the output length of the convolutional layers cSstj|||dddS)Nfloor) rounding_moder)rdiv) input_lengthrrr/r/r0_conv_out_length/szMSEWPreTrainedModel._get_feat_extract_output_lengths.._conv_out_length)zipr+r#r$)r*rrrrr/r/r0 _get_feat_extract_output_lengths*sz3SEWPreTrainedModel._get_feat_extract_output_lengthsfeature_vector_lengthrcCs~||dtj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dg d dg }|S)Nr@r)rrrr) r rrrrrzerosrrrflipcumsumr)r*r!rr batch_sizer/r/r0"_get_feature_vector_attention_mask9s "z5SEWPreTrainedModel._get_feature_vector_attention_maskN)r7r8r9r__annotations__base_model_prefixmain_input_namesupports_gradient_checkpointing_supports_flash_attn_supports_sdpa_supports_flex_attnrrr LongTensorrr r&r/r/r/r0r s  r r mask_prob mask_length min_masksrcs|\}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)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|}|kr}|d|kr*t|dd}|S)z;Given input length, compute how many spans should be maskedrr)rmax)rnum_masked_spanepsilonr0r/r1sequence_lengthr/r0compute_num_masked_spanls  z6_compute_mask_indices..compute_num_masked_spanNr@csg|]}qSr/r/r)r7r/r0rwsz)_compute_mask_indices..rrF)replace)r|nprandomritemdetachrtolistrzr"rchoicerlen concatenateonesint32appendarray broadcast_torjr3put_along_axis)rr/r0rr1r%r8r spec_aug_maskspec_aug_mask_idxsmax_num_masked_spanrr4spec_aug_mask_idxdummy_mask_idxoffsetsr/r5r0_compute_mask_indicesFs\    rNcseZdZdeffdd Z  ddejdeejdeejfdd Z e     dd eej deej deejd ee d ee d ee de eeffddZZS)SEWModelr+cst|||_t||_tj|jd|jd|_ |jd|j k|_ |j r1t |jd|j |_ t|j|_|jdksB|jdkrNtt|j |_t||_|dS)Nr@rr)rrr+rpfeature_extractorrr=rrr>rTproject_featuresrffeature_projectionrfeat_proj_dropoutfeature_dropoutmask_time_probmask_feature_prob Parameterrruniform_masked_spec_embedrencoder post_initrhr-r/r0rs    zSEWModel.__init__Nr3mask_time_indicesrcCst|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://huggingface.co/papers/1904.08779). apply_spec_augmentTNr)r/r0rr1)rr)r/r0r1r@)getattrr+rirYrrrUrrNmask_time_lengthmask_time_min_masksrtensorrrrVmask_feature_lengthmask_feature_min_masksr)r*r3r\rr%r7rTmask_feature_indicesr/r/r0_mask_hidden_statess4 zSEWModel._mask_hidden_statesrrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}||}|dd}||}|jr6||}| |}|durH| |j d|}|j ||d}|j |||||d} | d}|sh|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. NrrJ)r\rrrrrr)r+rruse_return_dictrPrAr>rQrRrTr&rrerZr r3r) r*rrr\rrrextract_featuresr3encoder_outputsr/r/r0r4s8     zSEWModel.forwardrNNNNN)r7r8r9rrr FloatTensorrr.rerrrrrr r4r:r/r/r-r0rOs@ . rOzk SEW Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). ) custom_introcseZdZddeeffdd ZddZddZd d Zd d Z e     dd ee j dee j dee dee dee dee j deeeffddZZS) SEWForCTCN 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 [`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)rrrOr rr final_dropoutrrn vocab_sizer|r.rYrooutput_hidden_sizerTrflm_headr[)r*r+rnrrr-r/r0rAs     zSEWForCTC.__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. Nadapter_attn_dimzCannot pass `target_lang`: z- if `config.adapter_attn_dim` is not defined.z)By default `target_lang` is set to 'eng'.T) force_load)rnr^r+r|loggerinfo load_adapter)r*rnr/r/r0 tie_weights^s zSEWForCTC.tie_weightscCtdt|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/r0freeze_feature_extractors  z"SEWForCTC.freeze_feature_extractorcC|jjdSr{Nr rPrrr/r/r0rz SEWForCTC.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. FNr rrrr/r/r0freeze_base_modelzSEWForCTC.freeze_base_modelrrrrrlabelsrc 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: rfrrr@)rRrrF)enabled)blank reduction zero_infinitylosslogitsr3r)r+rgr3rqr|r rrsr ones_likerr rr masked_selectrr log_softmaxfloat32rAbackendscudnnflagsctc_loss pad_token_idctc_loss_reductionctc_zero_infinity_HIDDEN_STATES_START_POSITIONrr3r)r*rrrrrrrr3rrr labels_masktarget_lengthsflattened_targets log_probsoutputr/r/r0r4sN      zSEWForCTC.forwardr1rj)r7r8r9rrrryrrrrrrrrrrr4r:r/r/r-r0rm;s6  rmz SEW 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)SEWForSequenceClassificationcst|t|dr|jrtdt||_|jd}|jr*t t |||_ t |j|j|_t |j|j|_|dS)NrozZSequence classification does not support the use of SEW adapters (config.add_adapter=True)r)rrrYror|rOr ruse_weighted_layer_sumrrWrrB layer_weightsrfrTclassifier_proj_size projector num_labels classifierr[)r*r+ num_layersr-r/r0rs    z%SEWForSequenceClassification.__init__cCrz)z Calling this function will disable the gradient computation for the feature encoder so that its parameters will not be updated during training. r|Nr}rr/r/r0rrz5SEWForSequenceClassification.freeze_feature_extractorcCrrrrr/r/r0rrz3SEWForSequenceClassification.freeze_feature_encodercCrrrrr/r/r0rrz.SEWForSequenceClassification.freeze_base_modelNrrrrrrrcCsz|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) a 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]`, a `numpy.ndarray` or a `torch.Tensor`, *e.g.* via the torchcodec library (`pip install torchcodec`) or 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 [`SEWProcessor.__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). NTrfrrr@rrJrr)r+rgrr rrstackrrrrrrrr r&rrrrrrrr3r)r*rrrrrrrr3 norm_weights pooled_output padding_maskexpand_padding_maskrrloss_fctrr/r/r0r4 sH    z$SEWForSequenceClassification.forwardrj)r7r8r9rrrrrrrrrrrrr4r:r/r/r-r0rs4   r)rmrrOr )NrNrd)Arr~typingrrrnumpyr:rrtorch.nnr activationsrintegrations.deepspeedr integrations.fsdpr modeling_flash_attention_utilsr modeling_layersr modeling_outputsr rrmodeling_utilsrrprocessing_utilsrrXrrconfiguration_sewr get_loggerr7rvrr;rBModulerIrbrerprrrrrrrr rrr.ndarrayrNrOrrmr__all__r/r/r/r0s          + , X$fI  wz s