o ۷iu@sFddlZddlZddlmZddlmZmZmZddlZ ddl Z ddl m Z ddl m Z ddlmZddlmZddlmZdd lmZmZdd lmZdd lmZdd lmZmZmZmZm Z dd l!m"Z"m#Z#ddl$m%Z%ddl&m'Z'm(Z(m)Z)ddl*m+Z+e(rddl,m-Z-e).e/Z0ee'ddGdddeZ1Gddde j2Z3Gddde j2Z4GdddeZ5GdddeZ6Gdd d eZ7Gd!d"d"e j2Z8Gd#d$d$e j2Z9  % dVd&e j2d'e j:d(e j:d)e j:d*ee j:d+ee;d,e;d-ee j:fd.d/ZGd4d5d5eZ?Gd6d7d7e j2Z@Gd8d9d9e j2ZAGd:d;d;eZBGdd?d?e j2ZDe'Gd@dAdAe#ZE  dWdBeFeGeGfdCe;dDeGd*ee jHdEeGdFe jIf dGdHZJe ZKe'GdIdJdJeEZLe'dKdGdLdMdMeEZMdNZNe'dOdGdPdQdQeEZOe'dRdGdSdTdTeEZPgdUZQdS)XN) dataclass)CallableOptionalUnion)CrossEntropyLoss)ACT2FN)is_deepspeed_zero3_enabled)is_fsdp_managed_module)_prepare_4d_attention_mask#_prepare_4d_attention_mask_for_sdpa)FlashAttentionKwargs)GradientCheckpointingLayer)BaseModelOutputCausalLMOutput ModelOutputSequenceClassifierOutputWav2Vec2BaseModelOutput)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)auto_docstringis_torch_flex_attn_availablelogging)UniSpeechConfig)make_flex_block_causal_maskzh Output type of [`UniSpeechForPreTrainingOutput`], with potential hidden states and attentions. ) custom_introc@seZdZUdZdZeejed<dZ eejed<dZ eejed<dZ eejed<dZ ee ejed<dZee ejed<dS) UniSpeechForPreTrainingOutputa loss (*optional*, returned when model is in train mode, `torch.FloatTensor` of shape `(1,)`): Total loss as the sum of the contrastive loss (L_m) and the diversity loss (L_d) as stated in the [official paper](https://huggingface.co/papers/2006.11477). projected_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`): Hidden-states of the model projected to *config.proj_codevector_dim* that can be used to predict the masked projected quantized states. projected_quantized_states (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.proj_codevector_dim)`): Quantized extracted feature vectors projected to *config.proj_codevector_dim* representing the positive target vectors for contrastive loss. codevector_perplexity (`torch.FloatTensor` of shape `(1,)`): The perplexity of the codevector distribution, used to measure the diversity of the codebook. Nlossprojected_statesprojected_quantized_statescodevector_perplexity hidden_states attentions)__name__ __module__ __qualname____doc__rrtorch FloatTensor__annotations__r r!r"r#tupler$r-r-f/home/ubuntu/vllm_env/lib/python3.10/site-packages/transformers/models/unispeech/modeling_unispeech.pyr:s rc$eZdZfddZddZZS)UniSpeechSamePadLayercs*t|ddkrd|_dSd|_dS)Nrr)super__init__num_pad_remove)selfnum_conv_pos_embeddings __class__r-r.r3Xs  zUniSpeechSamePadLayer.__init__cCs,|jdkr|ddddd|j f}|SNr)r4r5r#r-r-r.forward\s zUniSpeechSamePadLayer.forwardr%r&r'r3r; __classcell__r-r-r7r.r0Ws r0cr/) UniSpeechPositionalConvEmbeddingcs$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) Nr1) kernel_sizepaddinggroups weight_normr) modifier_rankweight)namedimparametrizations)r2r3nnConv1d hidden_sizer6num_conv_pos_embedding_groupsconvutilsrBhasattrrGr deepspeedzeroGatheredParametersrD original0 original1weight_gweight_vregister_external_parameterr0r@rfeat_extract_activation activation)r5configrBrOrTrUr7r-r.r3cs4     z)UniSpeechPositionalConvEmbedding.__init__cCs:|dd}||}||}||}|dd}|S)Nrr1) transposerLr@rXr:r-r-r.r;s     z(UniSpeechPositionalConvEmbedding.forwardr<r-r-r7r.r>bs !r>c&eZdZdfdd ZddZZS)UniSpeechNoLayerNormConvLayerrcsjt|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ dS)Nrrr?stridebias)r2r3conv_dim in_conv_dim out_conv_dimrHrI conv_kernel conv_stride conv_biasrLrrWrXr5rYlayer_idr7r-r.r3s  z&UniSpeechNoLayerNormConvLayer.__init__cCs||}||}|SN)rLrXr:r-r-r.r;s  z%UniSpeechNoLayerNormConvLayer.forwardrr<r-r-r7r.r\sr\cr[)UniSpeechLayerNormConvLayerrcs|t|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ tj |jdd|_ t |j|_dS)Nrrr]T)elementwise_affine)r2r3r`rarbrHrIrcrdrerL LayerNorm layer_normrrWrXrfr7r-r.r3s  z$UniSpeechLayerNormConvLayer.__init__cCs:||}|dd}||}|dd}||}|S)N)rLrZrmrXr:r-r-r.r;s     z#UniSpeechLayerNormConvLayer.forwardrir<r-r-r7r.rjsrjcr[)UniSpeechGroupNormConvLayerrcst|dkr|j|dnd|_|j||_tj|j|j|j||j||j d|_ t |j |_ tj|j|jdd|_dS)Nrrr]T) num_groups num_channelsaffine)r2r3r`rarbrHrIrcrdrerLrrWrX GroupNormrmrfr7r-r.r3s   z$UniSpeechGroupNormConvLayer.__init__cCs"||}||}||}|Srh)rLrmrXr:r-r-r.r;s   z#UniSpeechGroupNormConvLayer.forwardrir<r-r-r7r.rpsrpcs0eZdZdZfddZddZddZZS)UniSpeechFeatureEncoderz.Construct the features from raw audio waveformcstjdkr tddgfddtjdD}njdkr2fddtjD}n td jd t||_ d |_ d |_ dS) Ngrouprrgcsg|] }t|ddqS)rrw)r\.0irYr-r. sz4UniSpeechFeatureEncoder.__init__..rlayercsg|]}t|dqS)rw)rjrxr{r-r.r|s z`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) r2r3feat_extract_normrprangenum_feat_extract_layers ValueErrorrH ModuleList conv_layersgradient_checkpointing_requires_grad)r5rYrr7r{r.r3s        z UniSpeechFeatureEncoder.__init__cCs|D]}d|_qd|_dSNF) parameters requires_gradrr5paramr-r-r._freeze_parameterss  z*UniSpeechFeatureEncoder._freeze_parameterscCs:|dddf}|jr|jrd|_|jD]}||}q|S)NT)rtrainingrr)r5 input_valuesr# conv_layerr-r-r.r;s    zUniSpeechFeatureEncoder.forward)r%r&r'r(r3rr;r=r-r-r7r.rus  rucr/)UniSpeechFeatureProjectioncsJttj|jd|jd|_t|jd|j|_ t |j |_ dS)Nroeps) r2r3rHrlr`layer_norm_epsrmLinearrJ projectionDropoutfeat_proj_dropoutdropoutr5rYr7r-r.r3s z#UniSpeechFeatureProjection.__init__cCs&||}||}||}||fSrh)rmrr)r5r#norm_hidden_statesr-r-r.r;s   z"UniSpeechFeatureProjection.forwardr<r-r-r7r.rs rmodulequerykeyvalueattention_maskscalingr head_maskc Ks|dur |dd}t||dd|} |dur| |} tjj| dd} |dur5| |dddd} tjj| ||j d} t| |} | dd } | | fS)Nror1rrFr)pr) sizer)matmulrZrH functionalsoftmaxviewrr contiguous) rrrrrrrrkwargs attn_weights attn_outputr-r-r.eager_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)UniSpeechAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN embed_dim num_headsr is_decoderr_ is_causalrYcst||_||_||_|||_||_|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_)r2r3rrrhead_dimrYrrrrrHrk_projv_projq_projout_proj)r5rrrrr_rrYr7r-r.r30s&    zUniSpeechAttention.__init__r#key_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 ChannelNrorr1eagerr)rrrr)shaperrrrZrrrrY_attn_implementationrrrrreshaperr)r5r#rrrrris_cross_attentionbsztgt_lensrc_len q_input_shapekv_input_shape query_statescurrent_states key_states value_statesattention_interfacerrr-r-r.r;Os:      zUniSpeechAttention.forward)rFTFN)NNNF)r%r&r'r(intfloatboolrrr3r)Tensorrr r,r;r=r-r-r7r.r-sR"  rcr/)UniSpeechFeedForwardcsptt|j|_t|j|j|_ t |j t r"t |j |_n|j |_t|j|j|_t|j|_dSrh)r2r3rHractivation_dropoutintermediate_dropoutrrJintermediate_sizeintermediate_dense isinstance hidden_actstrrintermediate_act_fn output_densehidden_dropoutoutput_dropoutrr7r-r.r3s  zUniSpeechFeedForward.__init__cCs6||}||}||}||}||}|Srh)rrrrrr:r-r-r.r;s     zUniSpeechFeedForward.forwardr<r-r-r7r.rs  rcs&eZdZfddZdddZZS)UniSpeechEncoderLayercshtt|j|j|jd|d|_t|j |_ tj |j|j d|_ t||_tj |j|j d|_dS)NFrrrrrYr)r2r3rrJnum_attention_headsattention_dropout attentionrHrrrrlrrmr feed_forwardfinal_layer_normrr7r-r.r3s  zUniSpeechEncoderLayer.__init__NFcCsf|}|j|||d\}}}||}||}||}|||}||}|f}|r1||f7}|SNrr)rrrmrrr5r#rr attn_residualr_outputsr-r-r.r;s    zUniSpeechEncoderLayer.forwardrr<r-r-r7r.rs rc sjeZdZfddZ    ddejdeejded ed ef d d Z de ejdfd ejfddZ Z S)UniSpeechEncodercft|_t|_tjjjd|_ t j |_ t fddtjD|_d|_dS)Nrcg|]}tqSr-)rryrr{r-r.r|z-UniSpeechEncoder.__init__..Fr2r3rYr>pos_conv_embedrHrlrJrrmrrrrrnum_hidden_layerslayersrrr7r{r.r3s    zUniSpeechEncoder.__init__NFTr#rroutput_hidden_states return_dictcCs*|rdnd}|r dnd}|dur"|ddd|jd}d||<|||}||} || }||}||}tpAt|} |j D]3} |rN||f}t g} |j o[| |j jk} | r`| rk| |||d}|d}| rod}|rx||df}qE|r||f}|stdd |||fDSt|||d S) Nr-rorr1rrNNcs|] }|dur|VqdSrhr-ryvr-r-r. z+UniSpeechEncoder.forward..last_hidden_stater#r$) unsqueezerepeatr_update_full_maskrrmrr r rr)randrrY layerdropr,rr5r#rrrrall_hidden_statesall_self_attentionsexpand_attention_maskposition_embeddings synced_gpusr}dropout_probabilityskip_the_layer layer_outputsr-r-r.r;sL          zUniSpeechEncoder.forward inputs_embedscC|dur>|jjdkrd|vr|}|Sd}|S|jjdkr$t||j}|S|jjdkr8t|tjr6t|dd}|St||j}|SNflash_attention_2rsdpaflex_attentionF)r rYrr dtyperr)rrr r5rr r-r-r.r       z"UniSpeechEncoder._update_full_maskNFFT) r%r&r'r3r)tensorrrrr;rrr=r-r-r7r.rs,   < rcs,eZdZfddZdejfddZZS)UniSpeechAttnAdapterLayercsZt|j|_|j|_t|j|_t |j|j|_ t |_ t |j|j|_ dS)z Implements adapter modules directly with 3D tensor weight as parameters and without using ModuleList to speed up training throughput. N)r2r3adapter_attn_dim input_dimrJ hidden_dimrHrlnormrlinear_1ReLUact_fnlinear_2rr7r-r.r3s  z"UniSpeechAttnAdapterLayer.__init__r#cCs,||}||}||}||}|Srh)rrrrr:r-r-r.r;-s    z!UniSpeechAttnAdapterLayer.forward)r%r&r'r3r)r*r;r=r-r-r7r.rs rcs@eZdZfddZ  d dejdeejdefdd ZZ S) $UniSpeechEncoderLayerStableLayerNormcstt|j|j|jd|d|_t|j |_ tj |j|j d|_ t||_tj |j|j d|_t|dddurAt||_dSd|_dS)NFrrr)r2r3rrJrrrrHrrrrlrrmrrrgetattrr adapter_layerrr7r-r.r38s   z-UniSpeechEncoderLayerStableLayerNorm.__init__NFr#rrcCsz|}||}|j|||d\}}}||}||}||||}|jdur1|||}|f}|r;||f7}|Sr)rmrrrrr!rr-r-r.r;Ks    z,UniSpeechEncoderLayerStableLayerNorm.forwardr) r%r&r'r3r)rrrr;r=r-r-r7r.r7s rcsLeZdZfddZ    d ddZdeejdfd ejfd d ZZ S) UniSpeechEncoderStableLayerNormcr)Nrcrr-)rrr{r-r.r|mrz.Frrr7r{r.r3fs   z(UniSpeechEncoderStableLayerNorm.__init__NFTcCs*|rdnd}|r dnd}|dur"|ddd|jd}d||<|||}||} || }||}tp.r)rrrrrrr r rr)rrrYrrmr,rrr-r-r.r;qsL          z'UniSpeechEncoderStableLayerNorm.forwardrr cCr r rrr-r-r.rrz1UniSpeechEncoderStableLayerNorm._update_full_maskr) r%r&r'r3r;rr)rrr=r-r-r7r.r"es  > r"cs4eZdZdZfddZeddZddZZS)UniSpeechGumbelVectorQuantizerz Vector quantization using gumbel softmax. See `[CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX](https://huggingface.co/papers/1611.01144) 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 concatenationrror1)r2r3num_codevector_groupsrqnum_codevectors_per_groupnum_varscodevector_dimrrH Parameterr)r* codevectorsrr` weight_proj temperaturerr7r-r.r3s   z'UniSpeechGumbelVectorQuantizer.__init__cCs8|jdd}ttj|t|ddd }|S)NrrgHz>ro)meanr)expsumlog)probsmarginal_probs perplexityr-r-r._compute_perplexitys (z2UniSpeechGumbelVectorQuantizer._compute_perplexityc Cs|j\}}}||}||||jd}|jr?tjj||j dd |}t j ||||jddd}| |}n$|jdd}|j|jd|ddd}||||jd}| |}|||d}|d|j} | |||j|jd} | d||d} | |fS)NroT)tauhardrr?rn)rr*rrqrrHrgumbel_softmaxrr+type_asr)rr3argmax new_zerosscatter_rr)r&r.) r5r# batch_sizesequence_lengthrJcodevector_probscodevector_soft_distr2codevector_idxcodevectors_per_groupr)r-r-r.r;s0       z&UniSpeechGumbelVectorQuantizer.forward) r%r&r'r(r3 staticmethodr3r;r=r-r-r7r.r#s   r#c@sbeZdZUeed<dZdZdZdZdZ dZ ddZ de e jeffdd Zd ed e jfd d ZdS)UniSpeechPreTrainedModelrY unispeechrTc 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 weightsrr)r,stdrr1)abNr6)rr#r*rDdatanormal_r_zero_rHinituniform_r)r>rLmathsqrtr? in_channels constant_rr in_featuresrrYinitializer_rangerlrtfill_rIkaiming_normal_rA)r5rkr-r-r. _init_weightss<         z&UniSpeechPreTrainedModel._init_weights input_lengthscCs4dd}t|jj|jjD] \}}||||}q |S)zH Computes the output length of the convolutional layers cSstj|||dddS)Nfloor) rounding_moder)r)div) input_lengthr?r^r-r-r._conv_out_length<szSUniSpeechPreTrainedModel._get_feat_extract_output_lengths.._conv_out_length)ziprYrcrd)r5rWr\r?r^r-r-r. _get_feat_extract_output_lengths7sz9UniSpeechPreTrainedModel._get_feat_extract_output_lengthsfeature_vector_lengthrcCs|jdddddf}||tj}|jd}tj||f|j|jd}d|tj |jd|jd|df<| dgd dg }|S)Nrorr)rdevicer)r`) cumsumr^tor)longrzerosrr`arangeflipr)r5r_rnon_padded_lengthsoutput_lengthsr<r-r-r."_get_feature_vector_attention_maskFs "z;UniSpeechPreTrainedModel._get_feature_vector_attention_maskN)r%r&r'rr+base_model_prefixmain_input_namesupports_gradient_checkpointing_supports_flash_attn_supports_sdpa_supports_flex_attnrVrr) LongTensorrr^rir-r-r-r.rC s !rCr 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)r[num_masked_spanepsilonrrrqrsr=r-r.compute_num_masked_span|s  z6_compute_mask_indices..compute_num_masked_spanNrocsg|]}qSr-r-r)r=r-r.r|sz)_compute_mask_indices..rrF)replace)rnprandomritemdetachr.tolistrrdrchoicerelen concatenateonesint32appendarray broadcast_torruput_along_axis)rrqrrrrsr<ryrW spec_aug_maskspec_aug_mask_idxsmax_num_masked_spanr[rvspec_aug_mask_idxdummy_mask_idxoffsetsr-rwr._compute_mask_indicesVs\    rcseZdZdeffdd Z  ddejdeejdeejfdd Z e     dd eej deej deejd ee d ee d ee de eeffddZZS)UniSpeechModelrYcszt|||_t||_t||_|jdks|jdkr)t t |j |_|jr2t||_nt||_|dS)Nr)r2r3rYrufeature_extractorrfeature_projectionmask_time_probmask_feature_probrHr(r)rrJrLmasked_spec_embeddo_stable_layer_normr"encoderr post_initrr7r-r.r3s      zUniSpeechModel.__init__Nr#mask_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)rqrrrrs)r`r)rqrrrsro)r rYrrrbrrrrmask_time_lengthmask_time_min_masksr)rr`rrmask_feature_lengthmask_feature_min_masksexpand)r5r#rrr<r=rJmask_feature_indicesr-r-r._mask_hidden_statess4 z"UniSpeechModel._mask_hidden_statesrrrrrc Cs|dur|n|jj}|dur|n|jj}|dur|n|jj}||}|dd}|dur6||jd|}||\}}|j |||d}|j |||||d} | d}|s_||f| ddSt ||| j | j dS)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. Nrr1)rrrrrrr)rextract_featuresr#r$)rYrruse_return_dictrrZrirrrrUniSpeechBaseModelOutputr#r$) r5rrrrrrrr#encoder_outputsr-r-r.r;s8  zUniSpeechModel.forwardrNNNNN)r%r&r'rr3r)r*rrprrrrrr,rr;r=r-r-r7r.rs@ . rzZ UniSpeech Model with a vector-quantization module and ctc loss for pre-training. cseZdZdeffdd ZdefddZddZd d Ze dd e j d e j de j defddZ e    ddee jdee jdeedeedeedeeeff ddZZS)UniSpeechForPreTrainingrYcs~t|t||_t|j|_t||_ t |j |j |_ t |j |j|_t |j|j|_t|j|_|dSrh)r2r3rrDrHrfeat_quantizer_dropoutdropout_featuresr# quantizerrr'proj_codevector_dim project_qrJ project_hidnum_ctc_classesctc_proj final_dropoutrrrr7r-r.r3Ms    z UniSpeechForPreTraining.__init__r+cCs ||j_dS)zb Set the Gumbel softmax temperature to a given value. Only necessary for training N)rr+)r5r+r-r-r.set_gumbel_temperature\s z.UniSpeechForPreTraining.set_gumbel_temperaturecCtdt|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_encoderr5r-r-r.freeze_feature_extractorb  z0UniSpeechForPreTraining.freeze_feature_extractorcC|jjdS Calling this function will disable the gradient computation for the feature encoder so that its parameter will not be updated during training. NrDrrrr-r-r.rnz.UniSpeechForPreTraining.freeze_feature_encoderrtarget_featuresnegative_featurespredicted_featurescCs@tj||gdd}tj||dd}||}||}|S)z Compute logits for contrastive loss based using cosine similarity as the distance measure between `[positive_feature, negative_features]` and `[predicted_features]`. Additionally, temperature can be applied. rrro)r)catcosine_similarityrr8)rrrr+logitsr-r-r.compute_contrastive_logitsus   z2UniSpeechForPreTraining.compute_contrastive_logitsNrrrrrrcCsJ|dur|n|jj}|j|||||d}|d}||d}||\} } || |jjj} | | } t | d| d |jj} | dd} t | |j} | dd} | d} || d| | d} || } || } d}|s|dur||| | f|ddS|| | f|ddSt||| | |j|jdS) a Example: ```python >>> import torch >>> from transformers import AutoFeatureExtractor, UniSpeechForPreTraining >>> feature_extractor = AutoFeatureExtractor.from_pretrained("microsoft/unispeech-large-1500h-cv") >>> model = UniSpeechForPreTraining.from_pretrained("microsoft/unispeech-large-1500h-cv") >>> # TODO: Add full pretraining example ```Nrrrrorr1)rr r!r"r#r$)rYrrDrrrrbrDrrr)emptyrrS replace_probrZ bernoullirr`r masked_fillrrrr#r$)r5rrrrrrtransformer_featuresrquantized_featuresr"prob_replace_matrixsampled_replace_matrixrrr-r-r.r;sL        zUniSpeechForPreTraining.forward)r)NNNN)r%r&r'rr3rrrrrBr)r*rrrrrrr,rr;r=r-r-r7r.rGsD   rr1zq UniSpeech Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). cseZdZddeeffdd ZddZddZd d Zd d Z e     dd ee j dee j dee dee dee dee j deeeffddZZS)UniSpeechForCTCN target_langcs~t|t||_t|j|_||_|j dur#t d|j dt |dr.|j r.|jn|j}t||j |_|dS)a3 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 [`UniSpeechForCTC`] 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: `UniSpeechForCTC.from_pretrained(..., vocab_size=vocab_size)`. or define `vocab_size` of your model's configuration. add_adapter)r2r3rrDrHrrrr vocab_sizerr8rNroutput_hidden_sizerJrlm_headr)r5rYrrr7r-r.r3s     zUniSpeechForCTC.__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. NrzCannot pass `target_lang`: z- if `config.adapter_attn_dim` is not defined.z)By default `target_lang` is set to 'eng'.T) force_load)rr rYrloggerinfo load_adapter)r5rr-r-r. tie_weightss zUniSpeechForCTC.tie_weightscCr)rrNrrr-r-r.r rz(UniSpeechForCTC.freeze_feature_extractorcCrrrrr-r-r.rrz&UniSpeechForCTC.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. FNrDrrrr-r-r.freeze_base_modelz!UniSpeechForCTC.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: rrrzro)rFrrF)enabled)blank reduction zero_infinityrrr#r$)rYrrurrrDrrr) ones_likercr^r.rb masked_selectrHr log_softmaxfloat32rZbackendscudnnflagsctc_loss pad_token_idctc_loss_reductionctc_zero_infinity_HIDDEN_STATES_START_POSITIONrr#r$)r5rrrrrrrr#rrrW labels_masktarget_lengthsflattened_targets log_probsoutputr-r-r.r;'sN      zUniSpeechForCTC.forwardrhr)r%r&r'rrr3rrrrrr)rrrr,rr;r=r-r-r7r.rs6  rz UniSpeech 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)"UniSpeechForSequenceClassificationcst|t|dr|jrtdt||_|jd}|jr*t t |||_ t |j|j|_t |j|j|_|dS)Nrz`Sequence classification does not support the use of UniSpeech adapters (config.add_adapter=True)r)r2r3rNrrrrDruse_weighted_layer_sumrHr(r)r layer_weightsrrJclassifier_proj_size projector num_labels classifierr)r5rY num_layersr7r-r.r3vs    z+UniSpeechForSequenceClassification.__init__cCrrrrr-r-r.rrz;UniSpeechForSequenceClassification.freeze_feature_extractorcCrrrrr-r-r.rrz9UniSpeechForSequenceClassification.freeze_feature_encodercCrrrrr-r-r.rrz4UniSpeechForSequenceClassification.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 [`UniSpeechProcessor.__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). NTrrrrorr1rr)rYrrrDrr)stackrHrrrrr.rr,rirrrrrrrr#r$)r5rrrrrrrr# norm_weights pooled_output padding_maskexpand_padding_maskrrloss_fctrr-r-r.r;sH    z*UniSpeechForSequenceClassification.forwardr)r%r&r'r3rrrrrr)rrrr,rr;r=r-r-r7r.ros4   r)rrrrrC)NrNr9)RrMr dataclassesrtypingrrrnumpyr|r)torch.nnrHr activationsrintegrations.deepspeedr integrations.fsdpr modeling_attn_mask_utilsr r modeling_flash_attention_utilsr modeling_layersrmodeling_outputsrrrrrmodeling_utilsrrprocessing_utilsrrMrrrconfiguration_unispeechrintegrations.flex_attentionr get_loggerr%rrModuler0r>r\rjrprurrrrrrrrrrr"r#rCr,rrpndarrayrrrrrrr__all__r-r-r-r.s             -) X$].aFM  wv s