`iuddlZddlZddlmZddlmZmZmZddlZ 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)j.e/Z0ee'dGddeZ1Gdde j2Z3Gdde j2Z4GddeZ5GddeZ6Gdd eZ7Gd!d"e j2Z8Gd#d$e j2Z9 dTd&e j2d'e j:d(e j:d)e j:d*ee j:d+ee;d,e;d-ee j:fd.Z<Gd/d0e j2Z=Gd1d2e j2Z>Gd3d4eZ?Gd5d6e j2Z@Gd7d8e j2ZAGd9d:eZBGd;de j2ZDe'Gd?d@e#ZE dUdAeFeGeGfdBe;dCeGd*ee jHdDeGdEe jIf dFZJe ZKe'GdGdHeEZLe'dIGdJdKeEZMdLZNe'dMGdNdOeEZOe'dPGdQdReEZPgdSZQdS)VN) 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_introceZdZUdZdZeejed<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__r!rtorch FloatTensor__annotations__r"r#r$r%tupler&/home/jaya/work/projects/VOICE-AGENT/VIET/agent-env/lib/python3.11/site-packages/transformers/models/unispeech/modeling_unispeech.pyr r :s  )-D(5$ %,,,48hu01888>B): ;BBB9=8E$56===8!#C#Caar0cJ|jdkr|ddddd|j f}|SNr)r8r9r%s r1forwardzUniSpeechSamePadLayer.forward\s;   " ")!!!QQQ0F43F2F0F*FGMr0r'r(r)r7r? __classcell__r;s@r1r3r3WsLKKKKKr0r3c$eZdZfdZdZxZS) UniSpeechPositionalConvEmbeddingcttj|j|j|j|jdz|j|_tjj }ttjj drtjj j }trddl }|j|jjd5||jdd|_dddn #1swxYwYt|jdr-|jj jj}|jj jj}n|jj}|jj}|j|||j||n||jdd|_t-|j|_t0|j|_dS) Nr5) kernel_sizepaddinggroups weight_normr) modifier_rankweight)namedimparametrizations)r6r7nnConv1d hidden_sizer:num_conv_pos_embedding_groupsconvutilsrIhasattrrNr deepspeedzeroGatheredParametersrK original0 original1weight_gweight_vregister_external_parameterr3rGr feat_extract_activation activation)r9configrIrVr[r\r;s r1r7z)UniSpeechPositionalConvEmbedding.__init__cs I    62a77     h* 28,m < < @(3?K % ' ' E    22493CST2UU I I'K aHHH  I I I I I I I I I I I I I I Ity"455 .95<F95<F9-9- N 6 6tX F F F N 6 6tX F F F F# DIH!DDDDI,V-KLL  !?@sC??DDc|dd}||}||}||}|dd}|S)Nrr5) transposerSrGr_r>s r1r?z(UniSpeechPositionalConvEmbedding.forwardse%//155  -00  ]33  66 %//155 r0r@rBs@r1rDrDbsMAAAAABr0rDc&eZdZdfd ZdZxZS)UniSpeechNoLayerNormConvLayerrcZt|dkr|j|dz nd|_|j||_t j|j|j|j||j||j |_ t|j |_ dS)NrrrFstridebias)r6r7conv_dim in_conv_dim out_conv_dimrOrP conv_kernel conv_stride conv_biasrSr r^r_r9r`layer_idr;s r1r7z&UniSpeechNoLayerNormConvLayer.__init__s s r1r?z%UniSpeechNoLayerNormConvLayer.forwards* -00  66 r0rr@rBs@r1rdrdsR A A A A A Ar0rdc&eZdZdfd ZdZxZS)UniSpeechLayerNormConvLayerrct|dkr|j|dz nd|_|j||_t j|j|j|j||j||j |_ t j |jd|_ t|j|_dS)NrrrfT)elementwise_affine)r6r7rirjrkrOrPrlrmrnrS LayerNorm layer_normr r^r_ros r1r7z$UniSpeechLayerNormConvLayer.__init__s s r1r?z#UniSpeechLayerNormConvLayer.forwardse -00 %//B77  66 %//B77  66 r0rsr@rBs@r1rurusR A A A A A Ar0ruc&eZdZdfd ZdZxZS)UniSpeechGroupNormConvLayerrct|dkr|j|dz nd|_|j||_t j|j|j|j||j||j |_ t|j |_ t j|j|jd|_dS)NrrrfT) num_groups num_channelsaffine)r6r7rirjrkrOrPrlrmrnrSr r^r_ GroupNormryros r1r7z$UniSpeechGroupNormConvLayer.__init__s s r1r?z#UniSpeechGroupNormConvLayer.forwards; -00  66  66 r0rsr@rBs@r1r~r~sRrrrrrr r0r~c.eZdZdZfdZdZdZxZS)UniSpeechFeatureEncoderz.Construct the features from raw audio waveformctjdkr7tdgfdt jdz Dz}nDjdkr!fdt jD}nt djd tj||_ d |_ d |_ dS) Ngrouprrpc8g|]}t|dzS)rr)rd.0ir`s r1 z4UniSpeechFeatureEncoder.__init__..s@OOO.fq1uEEEOOOr0rlayerc2g|]}t|S)r)rurs r1rz4UniSpeechFeatureEncoder.__init__..s4DE+FQ???r0z`config.feat_extract_norm` is z), but has to be one of ['group', 'layer']FT) r6r7feat_extract_normr~rangenum_feat_extract_layers ValueErrorrO ModuleList conv_layersgradient_checkpointing_requires_grad)r9r`rr;s ` r1r7z UniSpeechFeatureEncoder.__init__s   #w . .6vJJJKOOOOv=ABBOOOKK % 0 0INvOmInInKKt1Ittt =55&+#"r0cP|D] }d|_ d|_dSNF) parameters requires_gradrr9params r1_freeze_parametersz*UniSpeechFeatureEncoder._freeze_parameterss4__&& ( (E"'E  #r0cr|dddf}|jr|jrd|_|jD] }||}|S)NT)rtrainingrr)r9 input_valuesr% conv_layers r1r?zUniSpeechFeatureEncoder.forwards[$QQQW-    /4= /*.M '* 6 6J&J}55MMr0)r'r(r)r*r7rr?rArBs@r1rrs\88#####($$$       r0rc$eZdZfdZdZxZS)UniSpeechFeatureProjectionc.ttj|jd|j|_tj|jd|j|_ tj |j |_ dS)Nr|eps) r6r7rOrxrilayer_norm_epsryLinearrQ projectionDropoutfeat_proj_dropoutdropoutr9r`r;s r1r7z#UniSpeechFeatureProjection.__init__sn ,vr':@UVVV)FOB$79KLLz&":;; r0c||}||}||}||fSrr)ryrr)r9r%norm_hidden_statess r1r?z"UniSpeechFeatureProjection.forwardsC!__];;(:;;  ]33 000r0r@rBs@r1rrsG<<<<< 1111111r0rmodulequerykeyvalueattention_maskscalingr head_maskc ||ddz}tj||dd|z} || |z} tj| d} || |ddddz} tj| ||j } tj| |} | dd } | | fS)Nr|r5rrMr)pr) sizer+matmulrbrO functionalsoftmaxviewrr contiguous) rrrrrrrrkwargs attn_weights attn_outputs r1eager_attention_forwardrs**R..D(<s}}Q':':;;gEL!#n4 =((2(>>L#innQAq&A&AA =((6?([[L,|U33K''1--88::K  $$r0c>eZdZdZ ddededed ed ed ed eeffd Z dde j dee j dee j dee j deede e dee j ee j eee j ffdZxZS)UniSpeechAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN embed_dim num_headsr is_decoderrh is_causalr`c t||_||_||_||z|_||_|j|z|jkrtd|jd|d|jdz|_||_ ||_ tj ||||_ tj ||||_tj ||||_tj ||||_dS)Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).r)rh)r6r7rrrhead_dimr`rrrrrOrk_projv_projq_projout_proj) r9rrrrrhrr`r;s r1r7zUniSpeechAttention.__init__0s "" !Y.  MI %$. 8 83dn33%.333 }d* $"i 94@@@ i 94@@@ i 94@@@  )YTBBB r0r%key_value_statesrlayer_head_maskoutput_attentionsrreturnc |du}|jdd\}} |r |jdn| } || d|jf} || d|jf} ||j| dd} |r|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|rr5eagerr)rrrr)shaperrrrbrrrr`_attn_implementationrrrrreshaperr)r9r%rrrrris_cross_attentionbsztgt_lensrc_len q_input_shapekv_input_shape query_statescurrent_states key_states value_statesattention_interfacerrs r1r?zUniSpeechAttention.forwardOs.T9%*3B3/ W/AN"(++wgr4=9 wDM:7t{{=116 FPPQRTUVV -?R))]5T[[005~FPPQRTUVV 7t{{>227HRRSTVWXX (? ; +w 6 6"9$+:Z"[ $7$7      % $}>CC$,L/% % %  % % ! \"))#w;;FFHH mmK00 L$..r0)rFTFN)NNNF)r'r(r)r*intfloatboolrrr7r+Tensorrrr.r?rArBs@r1rr-s^GG  ,0CCCC C  C  CC)CCCCCCD481526,1 3/3/|3/#5<03/!. 3/ "%,/ 3/ $D> 3/-.3/ u|Xel3XeEL>Q5RR S3/3/3/3/3/3/3/3/r0rc$eZdZfdZdZxZS)UniSpeechFeedForwardcttj|j|_tj|j|j|_ t|j trt|j |_n |j |_tj|j|j|_tj|j|_dSrr)r6r7rOractivation_dropoutintermediate_dropoutrrQintermediate_sizeintermediate_dense isinstance hidden_actstrr intermediate_act_fn output_densehidden_dropoutoutput_dropoutrs r1r7zUniSpeechFeedForward.__init__s $&Jv/H$I$I!"$)F,>@X"Y"Y f' - - 9'-f.?'@D $ $'-'8D $If&>@RSS j)>??r0c||}||}||}||}||}|Srr)rrrrrr>s r1r?zUniSpeechFeedForward.forwardsg// >> 00?? 11-@@ ))-88 ++M:: r0r@rBs@r1rrsL @ @ @ @ @r0rc&eZdZfdZddZxZS)UniSpeechEncoderLayerctt|j|j|jd||_tj|j |_ tj |j|j |_ t||_tj |j|j |_dS)NFrrrrr`r)r6r7rrQnum_attention_headsattention_dropout attentionrOrrrrxrryr feed_forwardfinal_layer_normrs r1r7zUniSpeechEncoderLayer.__init__s +(0,    z&"788 ,v'9v?TUUU088 " V-?VEZ [ [ [r0NFc|}||||\}}}||}||z}||}|||z}||}|f}|r||fz }|SNrr)rrryrrr9r%rr attn_residualr_outputss r1r?zUniSpeechEncoderLayer.forwards% )- .L]*8* * & |Q ]33 % 5  66 %(9(9-(H(HH --m<<  "  '   &Gr0rr@rBs@r1rrsQ \ \ \ \ \r0rc eZdZfdZ d dejdeejdeded ef d Z de ejdfd ejfd Z xZ S)UniSpeechEncoderct|_t|_t jjj|_ t j j |_ t j fdtjD|_d|_dS)Nrc.g|]}tSr/)rrr r`s r1rz-UniSpeechEncoder.__init__..s"$l$l$lq%:6%B%B$l$l$lr0Fr6r7r`rDpos_conv_embedrOrxrQrryrrrrrnum_hidden_layerslayersrrs `r1r7zUniSpeechEncoder.__init__s  >vFF,v'9v?TUUUz&"788 m$l$l$l$lERXRjLkLk$l$l$lmm &+###r0NFTr%rroutput_hidden_states return_dictc|rdnd}|rdnd}|;|ddd|jd}d||<|||}||} || z}||}||}tpt|} |j D]a} |r||fz}tj g} |j o| |j jk} | r| r| |||}|d}| rd}|r ||dfz}b|r||fz}|std|||fDSt!||| S) Nr/r|rr5rrNNc3K|]}||V dSrrr/rvs r1 z+UniSpeechEncoder.forward..(mmq_`_l_l_l_l_lmmr0last_hidden_stater%r&) unsqueezerepeatr_update_full_maskrryrr r rr+randrr` layerdropr.rr9r%rrrrall_hidden_statesall_self_attentionsexpand_attention_maskposition_embeddings synced_gpusrdropout_probabilityskip_the_layer layer_outputss r1r?zUniSpeechEncoder.forwards#7@BBD$5?bb4  %$2$<$BB gii  $.$/BB r0r%c||}||}||}||}|Srr)r?r@rBrCr>s r1r?z!UniSpeechAttnAdapterLayer.forward-sL -00  m44  M22  m44 r0)r'r(r)r7r+r,r?rArBs@r1r:r:s[ C C C C CU%6r0r:cXeZdZfdZ ddejdeejdefdZxZ S) $UniSpeechEncoderLayerStableLayerNormctt|j|j|jd||_tj|j |_ tj |j|j |_ t||_tj |j|j |_t#|ddt%||_dSd|_dS)NFrrr<)r6r7rrQrrrrOrrrrxrryrrrgetattrr: adapter_layerrs r1r7z-UniSpeechEncoderLayerStableLayerNorm.__init__8s +(0,    z&"788 ,v'9v?TUUU088 " V-?VEZ [ [ [ 6-t 4 4 @!:6!B!BD   !%D   r0NFr%rrcJ|}||}||||\}}}||}||z}||||z}|j|||z}|f}|r||fz }|Sr)ryrrrrrIrs r1r?z,UniSpeechEncoderLayerStableLayerNorm.forwardKs &  66 )- .L]*8* * & |Q ]33 % 5 %(9(9$:O:OP]:^:^(_(__   ))D,>,>},M,MMM "  '   &Gr0r) r'r(r)r7r+rrrr?rArBs@r1rFrF7s~&&&&&,26"' |!. r0rFcbeZdZfdZ d dZdeejdfdejfdZxZ S) UniSpeechEncoderStableLayerNormct|_t|_t jjj|_ t j j |_ t j fdtjD|_d|_dS)Nrc.g|]}tSr/)rFrs r1rz.ms" c c ca 1& 9 9 c c cr0Frrs `r1r7z(UniSpeechEncoderStableLayerNorm.__init__fs  >vFF,v'9v?TUUUz&"788 m c c c c5IaCbCb c c c   ',###r0NFTc|rdnd}|rdnd}|;|ddd|jd}d||<|||}||} || z}||}t pt|} |jD]a} |r||fz}tj g} |j o| |j j k} | r| r| |||}|d}| rd}|r ||dfz}b||}|r||fz}|std|||fDSt!||| S) Nr/r|rr5rrrc3K|]}||V dSrrr/rs r1rz:UniSpeechEncoderStableLayerNorm.forward..rr0r)rr rr!rrr r rr+r"rr`r#ryr.rr$s r1r?z'UniSpeechEncoderStableLayerNorm.forwardqs#7@BBD$5?bb4  %$2$<$eZdZdZfdZedZdZxZS)UniSpeechGumbelVectorQuantizerz Vector quantization using gumbel softmax. See `[CATEGORICAL REPARAMETERIZATION WITH GUMBEL-SOFTMAX](https://huggingface.co/papers/1611.01144) for more information. ct|j|_|j|_|j|jzdkr td|jd|jdtj tj d|j|jz|j|jz|_ tj |jd|j|jz|_d|_dS)Nrz`config.codevector_dim z5 must be divisible by `config.num_codevector_groups` z for concatenationrr|r5)r6r7num_codevector_groupsrnum_codevectors_per_groupnum_varscodevector_dimrrO Parameterr+r, codevectorsrri weight_proj temperaturers r1r7z'UniSpeechGumbelVectorQuantizer.__init__s  68  4? 2a 7 7Y&*?YY59_YYY  <  a4=!@&BW[_[jBj k k  9V_R%8$/DM:YZZr0c |d}tjtj|tj|dzzd }|S)NrrgHz>r|)meanr+expsumlog)probsmarginal_probs perplexitys r1_compute_perplexityz2UniSpeechGumbelVectorQuantizer._compute_perplexitys^**Y .59^VZEZ;[;[*[ac d d ddeeiikk r0c|j\}}}||}|||z|jzd}|jrt j||j d |}tj |||z|jdd}| |}n|d}|j|jd|ddd}|||z|jd}| |}|||zd}|d|jz} | ||z|j|jd} | d||d} | |fS)Nr|T)tauhardrr?r{)rr[rrrrOrgumbel_softmaxrr\type_asr+rreargmax new_zerosscatter_rrZrWr`) r9r% batch_sizesequence_lengthrQcodevector_probscodevector_soft_distrdcodevector_idxcodevectors_per_grouprZs r1r?z&UniSpeechGumbelVectorQuantizer.forwards3@3F0 O[((77 %**:+G$/+Y[]^^ = D!};;##%%4+;$ <  gm$$  $)="":#?RTUU[[]]ce$$$ 112FGGJJ+11b199N6}6 8KLUUN''A..    044Z/5QSWSbdfgg 112BCCJ+00o1MrRR 0 : :2 > >AQ Q+00o1Mt`d`moqrr !oob))..z?BOO J&&r0) r'r(r)r*r7 staticmethodrer?rArBs@r1rSrSsl (\ #'#'#'#'#'#'#'r0rScxeZdZUeed<dZdZdZdZdZ dZ dZ de e jeffdZded e jfd Zd S) UniSpeechPreTrainedModelr` unispeechrTc \t|tro|jjjdd|jjjtj |j dSt|trtj |j jddtjd|j jd|j jzz ztj |j jddSt|t&r}tjd|jjz }tj |jj| |tj |jj| |dSt|tjrT|jjd|jj|j |jjdSdSt|tjtjfr?|jj|jjddSt|tjrtj |j|j[tj|j|j|jdzz }tj |j| |dSdSdS) zInitialize the weightsrr)r^stdrr5)abNri)rrSr[rKdatanormal_rhzero_rOinituniform_rZrDrSmathsqrtrF in_channels constant_rr in_featuresrr`initializer_rangerxrfill_rPkaiming_normal_rH)r9rks r1 _init_weightsz&UniSpeechPreTrainedModel._init_weightss f< = = 9   % * 2 2 2 C C C   # ( . . 0 0 0 G  V/ 0 0 0 0 0  @ A A 9 GOO " !v{'>q'AFKD['["\]]]     G  fk. 2 2 2 2 2  : ; ; 9 !f/;;<._conv_out_length<s&9\K7wWWWZ[[ [r0)zipr`rlrm)r9rrrFrgs r1 _get_feat_extract_output_lengthsz9UniSpeechPreTrainedModel._get_feat_extract_output_lengths7s\  \ \ \ $'t{'> @W#X#X Q Q K,,]KPPMMr0feature_vector_lengthrc|ddddf}||tj}|jd}tj||f|j|j}d|tj |jd|j|dz f<| dgd dg }|S)Nr|rr)r4devicer)r) cumsumrtor+longrzerosr4rarangeflipr)r9rrnon_padded_lengthsoutput_lengthsros r1"_get_feature_vector_attention_maskz;UniSpeechPreTrainedModel._get_feature_vector_attention_maskFs,22r2::111b5A>>?QRRUUV[V`aa#)!,  . /~7KTbTi   uv ^%9!%<^EZ[[[]kno]opq',,bT2299"==BBB4HHMMOOr0N)r'r(r)rr-base_model_prefixmain_input_namesupports_gradient_checkpointing_supports_flash_attn_supports_sdpa_supports_flex_attnrrr+ LongTensorrrrr/r0r1rwrw s #$O&*#N999B eEDTVYDY>Z      ]b]m      r0rwr mask_prob mask_length min_masksrc@|\}dkrtdkrtdddtjdfd}|9|dnfd t|D}tj |ft }g} |} | d kr|S|D]} || } tj tj | dz z | d } t| d krdz }n| d }tj| tj| | z tj |zg} | | tj| } tj| dddddf|| f} | || z} tj ddddf}tj||| f|| z}| |z} | dz krdz | | dz k<tj|| 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`: `ct|zz z}t|}|zkrz}|dz z |krt|dz z d}|S)z;Given input length, compute how many spans should be maskedrr)rmax)rnum_masked_spanepsilonrrrrps r1compute_num_masked_spanz6_compute_mask_indices..compute_num_masked_span|s~i,6DwNOOoy99 [ (? : :-.s 9 9 9!o 9 9 9r0r4rF)replace)rnprandomr"itemdetachr`tolistrrrchoicerlen concatenateonesint32appendarray broadcast_torrput_along_axis)rrrrrrorr spec_aug_maskspec_aug_mask_idxsmax_num_masked_spanrrspec_aug_mask_idxdummy_mask_idxoffsetsrrps `` ` @@r1_compute_mask_indicesrVsP0#(JQABBB_$$ :^i : :'6 : : :   innQ$$&&G         $  % ##B''..000 9 9 9 9uZ'8'8 9 9 9Hj/:$GGGM11/BBa%55 11,??I,, IlkAo6 7 7RW-    ! !Q & &-q0NN.q1NN (;o(MUWU] ^ ^ ^ao o p   !!"34444"455111aaa:&5H+(V,33J@SVa@abbi $$T4]3Gog 4G'UVV^^'+5G,g5/A"555GVYZGZ-!0CCDm%7B??? r0c*eZdZdeffd Z ddejdeejdeejfdZ e ddeej deej deejd ee d ee d ee d e eeffd ZxZS)UniSpeechModelr`ct|||_t||_t ||_|jdks |jdkrBtj tj |j |_|jrt#||_nt'||_|dS)Nr)r6r7r`rfeature_extractorrfeature_projectionmask_time_probmask_feature_probrOrYr+rrQrmasked_spec_embeddo_stable_layer_normrLencoderr  post_initrs r1r7zUniSpeechModel.__init__s     !8!@!@"}G[chcm n n n /3/E/H/HI\/]/]M+ , ; (1 , , ,#8[)+7 K;+< $$$ $)<0D]Mainis#t#t#t #74#@#G#GO]_#`#` 23M. /r0rrrrrc||n |jj}||n |jj}||n |jj}||}|dd}|!||jd|}||\}}| |||}| |||||} | d}|s||f| ddzSt||| j | j 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. Nrr5)rrrrrrr)rextract_featuresr%r&)r`rruse_return_dictrrbrrrrrUniSpeechBaseModelOutputr%r&) r9rrrrrrrr%encoder_outputss r1r?zUniSpeechModel.forwardsW2C1N--TXT_Tq$8$D $+Jj &1%~1  ,, )/!5# '  (*  K!#34qrr7JJ J'+-)7&1     r0rNNNNN)r'r(r)rr7r+r,rrrrrrrr.rr?rArBs@r1rrs<(:>59 ,,(,$E$56,!!12 ,,,,\269=,0/3&*2 2 u|,2 !.2 $E$56 2 $D> 2 'tn 2 d^2  u.. /2 2 2 ^2 2 2 2 2 r0rzZ UniSpeech Model with a vector-quantization module and ctc loss for pre-training. c eZdZdeffd ZdefdZdZdZe dde j d e j d e j defd Z e dd ee jdee jdeedeedeedeeeff dZxZS)UniSpeechForPreTrainingr`ct|t||_t j|j|_t||_ t j |j |j |_ t j |j |j|_t j |j|j|_t j|j|_|dSrr)r6r7rrxrOrfeat_quantizer_dropoutdropout_featuresrS quantizerrrXproj_codevector_dim project_qrQ project_hidnum_ctc_classesctc_proj final_dropoutrrrs r1r7z UniSpeechForPreTraining.__init__Ms    '// " 6+H I I7??6#8&:TUU9V%?ASTT &"4f6LMM z&"677  r0r\c||j_dS)zb Set the Gumbel softmax temperature to a given value. Only necessary for training N)rr\)r9r\s r1set_gumbel_temperaturez.UniSpeechForPreTraining.set_gumbel_temperature\s&1"""r0cbtjdt|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_encoderr9s r1freeze_feature_extractorz0UniSpeechForPreTraining.freeze_feature_extractorb;   Q    ##%%%%%r0cB|jjdS Calling this function will disable the gradient computation for the feature encoder so that its parameter will not be updated during training. Nrxrrrs r1rz.UniSpeechForPreTraining.freeze_feature_encodern! (;;=====r0rtarget_featuresnegative_featurespredicted_featuresctj||gd}tj||d}||}||z }|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. rrr|)r+catcosine_similarityrrk)rrrr\logitss r1compute_contrastive_logitsz2UniSpeechForPreTraining.compute_contrastive_logitsusq )_6G$HaPPP();)A)A)C)C_EZEZE\E\bdeee00+% r0NrrrrrrcN||n |jj}||||||}|d}||d}||\} } || |jjj} | | } tj | d| d |jj} | dd} tj| |j} | dd} | d} || d| | dz} || } || } d}|s#|||| | f|ddzS|| | f|ddzSt/||| | |j|jS) 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 ```Nrrrr|rr5)r!r"r#r$r%r&)r`rrxrrrrrKr4rr+emptyrr replace_probrb bernoullirrr masked_fillrrr r%r&)r9rrrrrr transformer_featuresrquantized_featuresr$prob_replace_matrixsampled_replace_matrixr r!s r1r?zUniSpeechForPreTraining.forwardsW*&1%,A,A$.BWB],^,^__!--.@AA#k*>*C*CA*F*FH\HaHabcHdHdeekk K $  2;;AqAA!&1D!E!E!J!J!L!L!O!OPdPk!l!l!7!A!A!Q!G!G!7!A!A"!E!E%112H#NN  * *,B+BC H H  f%%v&& c24FH]^ahijikikalll(*<>STW^_`_a_aWbb b,1'9"7!/)     r0)r)NNNN)r'r(r)rr7rrrrrur+r,r rrrrrr.r r?rArBs@r1rrGsu       1#1111 & & &>>>  * ,"- \&26,0/3&* D D u|,D !.D $D> D 'tn D d^ D  u33 4D D D ^D D D D D r0rr5zq UniSpeech Model with a `language modeling` head on top for Connectionist Temporal Classification (CTC). ceZdZddeeffd ZdZdZdZdZ e ddee j d ee j d ee d ee d ee d ee j deeeffdZxZS)UniSpeechForCTCN target_langct|t||_t j|j|_||_|j td|j dt|dr|j r|jn|j}t j||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)r6r7rrxrOrrrr vocab_sizerr;rUroutput_hidden_sizerQrlm_headr)r9r`rrr;s r1r7zUniSpeechForCTC.__init__s    '//z&"677 &   $HHHH *1)G)G vFL^ vF % %djdv y!3V5FGG  r0c|j}|)t|jddtd|d|2t|jddtddS|||ddSdS)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. Nr<zCannot pass `target_lang`: z- if `config.adapter_attn_dim` is not defined.z)By default `target_lang` is set to 'eng'.T) force_load)rrHr`rloggerinfo load_adapter)r9rs r1 tie_weightszUniSpeechForCTC.tie_weightss&  "wt{||jjkrt d|jj||||||}|d}||}||} d} |Z||ntj |tj }| | d tj } |dk} | d} || }tj| dtjdd}tjjd 5tj||| | |jj|jj|jj } dddn #1swxYwY|s| f|t6dz}| | f|zn|St9| | |j|j 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|)rMr4rF)enabled)blank reduction zero_infinityr!r r%r&)r`rrrrrxrrr+ ones_likerrr`r masked_selectrOr log_softmaxfloat32rbbackendscudnnflagsctc_loss pad_token_idctc_loss_reductionctc_zero_infinity_HIDDEN_STATES_START_POSITIONrr%r&)r9rrrrrr+r r%r r!r labels_masktarget_lengthsflattened_targets log_probsoutputs r1r?zUniSpeechForCTC.forward's"&1% HSM:<<<* & & &>>>(((26,0/3&*)-D D u|,D !.D $D> D 'tn D d^ D &D  un$ %D D D ^D D D D D r0rz UniSpeech Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like SUPERB Keyword Spotting. ceZdZfdZdZdZdZe ddee j dee j dee d ee d ee d ee j d e e effd ZxZS)"UniSpeechForSequenceClassificationct|t|dr|jrt dt ||_|jdz}|jr.tj tj ||z |_ tj|j|j|_tj|j|j|_|dS)Nrz`Sequence classification does not support the use of UniSpeech adapters (config.add_adapter=True)r)r6r7rUrrrrxruse_weighted_layer_sumrOrYr+r layer_weightsrrQclassifier_proj_size projector num_labels classifierr)r9r` num_layersr;s r1r7z+UniSpeechForSequenceClassification.__init__vs    6= ) ) f.@ r (//-1  ( S!#ej.D.Dz.Q!R!RD 6#5v7RSS)F$?ARSS r0cbtjdt|dSrrrs r1rz;UniSpeechForSequenceClassification.freeze_feature_extractorrr0cB|jjdSrrrs r1rz9UniSpeechForSequenceClassification.freeze_feature_encoderrr0cL|jD] }d|_ dSr'r(rs r1r)z4UniSpeechForSequenceClassification.freeze_base_modelr*r0Nrrrrrr+rcd||n |jj}|jjrdn|}||||||}|jjrx|t}t j|d}tj |j d} || dddz d}n|d}| |}||d} n||jd|} | ddd|jd} d || <| d| d ddz } || } d}|Kt)}|| d|jj| d}|s| f|tdz}||f|zn|St-|| |j|j 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). NTrrrr|rr5rr1)r`rrFrxr=r+stackrOrrrGrr`rIr^rrrr rKrrJrr%r&)r9rrrrrr+r r% norm_weights pooled_output padding_maskexpand_padding_maskr r!loss_fctrBs r1r?z*UniSpeechForSequenceClassification.forwardsW.&1%->r1a-H-HHMMRSMTTMM#AJM}55  !)..1.55MMBB=CVWXCY[ijjL"."8"8"<"<"C"CAq-J]^_J`"a"a 25M.. /)--!-44|7G7GA7G7N7N7S7STVXY7Z7ZZM//  '))H8FKKDK,BCCV[[QS__UUD FY)F)G)G!HHF)-)9TGf$$v E'!/)     r0r)r'r(r)r7rrr)rrr+rrrr.rr?rArBs@r1rDrDos" & & &>>>(((26,0/3&*)-B B u|,B !.B $D> B 'tn B d^ B &B  u.. /B B B ^B B B B B r0rD)rrrDrrw)NrNr=)Rrr dataclassesrtypingrrrnumpyrr+torch.nnrOr activationsr integrations.deepspeedr integrations.fsdpr modeling_attn_mask_utilsr r modeling_flash_attention_utilsrmodeling_layersrmodeling_outputsrrrrrmodeling_utilsrrprocessing_utilsrrTrrrconfiguration_unispeechrintegrations.flex_attentionr get_loggerr'r r Moduler3rDrdrur~rrrrrrrrr r:rFrLrSrwr.rrndarrayrrrrr=rrD__all__r/r0r1rjs, !!!!!!,,,,,,,,,, %%%%%%!!!!!!@@@@@@777777ggggggggBBBBBB999999GFFFFFFF&&&&&&JJJJJJJJJJ444444 !!KJJJJJJ  H % %  :::::K::   :.BI*****ry***Z$>*"<6"<0&&&&&bi&&&R 1 1 1 1 1 1 1 1* $(,%% I% <% % < % U\* % e_ %% %%%%%