U ,-et@sNdZddlZddlmZmZmZddlZddlZddl Zddlm Z ddl m Z ddl mZddlmZdd lmZmZmZmZmZdd lmZdd lmZmZmZmZd d lmZd dl m!Z!m"Z"e#e$Z%dZ&dZ'dgZ(ej)e*e*dddZ+d>ej,ej-ej.e*dddZ/d?ej)ej-ee*dddZ0d@ee*e*fe1e*eej2e*ej3dddZ4ej)e*ej)ddd Z5ej3d!d"d#Z6Gd$d%d%e j7Z8Gd&d'd'e j9Z:Gd(d)d)e j9Z;Gd*d+d+e j9Zd/Z?d0Z@Gd1d2d2e=ZAGd3d4d4e=ZBed5e>Gd6d7d7e=ZCed8e>Gd9d:d:e=ZDed;e@Gd.compute_num_masked_spanNrcsg|]}qSr r .0_)rHr r! sz)_compute_mask_indices..r$rF)replace)rnprandomranditemsumdetachtolistranger2r9choicer-lenZ concatenateonesZint32appendarrayZ broadcast_toreshaperCZput_along_axis)rr<r=r>r? batch_sizerI input_lengthsZ spec_aug_maskZspec_aug_mask_idxsZmax_num_masked_spanrDrEZspec_aug_mask_idxZdummy_mask_idxoffsetsr rFr!_compute_mask_indicesms`    ra)inputs filter_widthr@cCsr|dks|ddkrtd|d}|jd|kr6|Stjj|||ddfdd}|d|ddd|f}|S) z Applies a median filter of width `filter_width` along the last dimension of the input. The `inputs` tensor is assumed to be 3- or 4-dimensional. rrz&`filter_width` should be an odd numberrZreflect)mode.)rrr functionalpadZunfoldsort)rbrcZ pad_widthresultr r r!_median_filtersrj)matrixcCs6|j\}}tj|d|dftjdtj}tj|d|dftjd }d|d<td|dD]}td|dD]}||d|df}||d|f}|||df} ||kr|| kr|d} } n&||kr|| kr|d} } n | d} } ||d|df| |||f<| |||f<qrq`|jdd}|jdd}d|dddf<d|dddf<g} g} |dkst|dkr| |d| |d|||fdkr|d8}|d8}nN|||fdkr|d8}n2|||fdkr|d8}ntd|d|d q`t| ddd } t| ddd } | | fS) z Measures similarity between two temporal sequences: the input audio and the output tokens. Used to generate token-level timestamps. rrNr)rrrdNz9Internal error in dynamic time warping. Unexpected trace[z, z]. Please file a bug report.r) rrPrZfloat32infrWr[ RuntimeErrorr\)rkZ output_lengthrDZcosttracejiZc0c1c2ct text_indices time_indicesr r r!_dynamic_time_warpingsL "       rxcs6eZdZdeeeedfdd Zd ddZZS) WhisperPositionalEmbeddingN) num_positions embedding_dim padding_idxcst||dSN)super__init__)selfrzr{r| __class__r r!r0sz#WhisperPositionalEmbedding.__init__rcCs|j|||jdSNr)weightr)rrr&r r r!forward3sz"WhisperPositionalEmbedding.forward)N)r)__name__ __module__ __qualname__rBrrr __classcell__r r rr!ry/sryc seZdZdZdeeeeedfdd Zej eedd d Z dej e ej e e ej e ej e ej ee ej e ej e e ej fd d dZ ZS)WhisperAttentionz=Multi-headed attention from 'Attention Is All You Need' paperFT) embed_dim num_headsdropout is_decoderbiascst||_||_||_|||_|j||jkrNtd|jd|d|jd|_||_t j ||dd|_ 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).gFr)r~rrrrhead_dimrscalingrrLineark_projv_projq_projout_proj)rrrrrrrr r!r:s   zWhisperAttention.__init__)tensorseq_lenr4cCs ||||j|jddS)Nrrd)r/rr transpose contiguous)rrrr4r r r!_shapeVszWhisperAttention._shapeN) hidden_stateskey_value_statespast_key_valuer>layer_head_maskoutput_attentionsr@cCsx|dk }|\}} } |||j} |r\|dk r\|djd|jdkr\|d} |d} n|r|||d|} |||d|} n|dk r|||d|} |||d|} tj|d| gdd} tj|d| gdd} n(|||d|} |||d|} |j r | | f}||j d|j f}|| | |j |} | j |} | j |} | d}t| | dd}|||j | |fkrtd||j | |fd||dk r ||d| |fkrtd |d| |fd|| ||j | ||}| ||j | |}tjj|dd}|dk r||j fkrhtd |j fd|| dddd| ||j | |}| ||j | |}|r| ||j | |}| ||j | |}nd}tjj||j|jd }t|| }|||j | |j fkr4td ||j | |j fd|| ||j | |j }|dd}| || |j}||}|||fS) z#Input shape: Batch x Time x ChannelNrrdrrr'z$Attention weights should be of size z , but is z!Attention mask should be of size z/Head mask for a single layer should be of size ptrainingz `attn_output` should be of size )r.rrrrrrr)r1rrrr/r]Zbmmrrrrfsoftmaxrrrr)rrrrr>rrZis_cross_attentionr4r5rLZ query_statesZ key_statesZ value_statesZ proj_shaper: attn_weightsZattn_weights_reshapedZ attn_probsZ attn_outputr r r!rZs~       "    zWhisperAttention.forward)rFT)NNNNF)rrr__doc__rBfloatr9rr)Tensorrrrrrr r rr!r7s4 rcsBeZdZedfdd ZdejejejeejdddZZ S) WhisperEncoderLayerconfigcst|j|_t|j|j|jd|_t |j|_ |j |_ t |j |_|j|_t|j|j|_t|j|j|_t |j|_dS)N)rrr)r~rd_modelrrZencoder_attention_headsattention_dropout self_attnr LayerNormself_attn_layer_normrractivation_function activation_fnactivation_dropoutrZencoder_ffn_dimfc1fc2final_layer_normrrrr r!rs  zWhisperEncoderLayer.__init__F)rr>rrr@c Cs|}||}|j||||d\}}}tjj||j|jd}||}|}||}|||}tjj||j |jd}| |}tjj||j|jd}||}|j t j krt |st |rt |j jd}t j|| |d}|f} |r| |f7} | S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. )rr>rrri)r,rC)rrrrfrrrrrrrr$r)Zfloat16isinfanyisnanr+rCclamp) rrr>rrresidualrrLZ clamp_valueoutputsr r r!rs8        zWhisperEncoderLayer.forward)F) rrrrrr)rr9rrr r rr!rsrc sxeZdZedfdd Zd ejeejeejeejeejeejeeejee ee ejd dd Z Z S) WhisperDecoderLayerrcst|j|_t|j|j|jdd|_|j|_t |j |_ |j |_ t |j|_t|j|j|jdd|_t |j|_t |j|j|_t |j|j|_t |j|_dS)NT)rrrr)rr)r~rrrrZdecoder_attention_headsrrrrrrrrrr encoder_attnencoder_attn_layer_normrZdecoder_ffn_dimrrrrrr r!rs,  zWhisperDecoderLayer.__init__NFT) rr>encoder_hidden_statesencoder_attention_maskrcross_attn_layer_head_maskrr use_cacher@c Cs^|} ||}|dk r"|ddnd} |j|| |||d\}} } tjj||j|jd}| |}d}d}|dk r|} ||}|dk r|ddnd}|j||||||d\}}}tjj||j|jd}| |}| |} |} ||}| | |}tjj||j |jd}| |}tjj||j|jd}| |}|f}|rJ|| |f7}| rZ|| f7}|S)a Args: hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)` attention_mask (`torch.FloatTensor`): attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. encoder_hidden_states (`torch.FloatTensor`): cross attention input to the layer of shape `(batch, seq_len, embed_dim)` encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values. layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size `(encoder_attention_heads,)`. cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of size `(decoder_attention_heads,)`. past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. Nrd)rrr>rrr)rrr>rrr) rrrrfrrrrrrrrr)rrr>rrrrrrrrZself_attn_past_key_valueZself_attn_weightsZpresent_key_valueZcross_attn_present_key_valueZcross_attn_weightsZcross_attn_past_key_valuerr r r!r4sT        zWhisperDecoderLayer.forward)NNNNNNFT) rrrrrr)rrrr9rrr r rr!rs* rc@sFeZdZeZdZdZdZddgZddZ dd d Z e j d d d Z dS)WhisperPreTrainedModelmodelinput_featuresTrrcCs|jj}t|tjtjfrD|jjjd|d|j dk r~|j j n:t|tj r~|jjjd|d|j dk r~|jj|j  dS)Nr)meanstd) rZinit_std isinstancerrConv1drdataZnormal_rZzero_ Embeddingr|)rmodulerr r r! _init_weightss   z$WhisperPreTrainedModel._init_weightsFcCst|ttfr||_dSr})rWhisperDecoderWhisperEncodergradient_checkpointing)rrvaluer r r!_set_gradient_checkpointingsz2WhisperPreTrainedModel._set_gradient_checkpointing)r_cCs|ddd}|S)zH Computes the output length of the convolutional layers rrdr )rr_r r r! _get_feat_extract_output_lengthssz7WhisperPreTrainedModel._get_feat_extract_output_lengthsN)F)rrrr config_classbase_model_prefixZmain_input_nameZsupports_gradient_checkpointingZ_no_split_modulesrrr) LongTensorrr r r r!rs raL This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads etc.) This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and behavior. Parameters: config ([`WhisperConfig`]): Model configuration class with all the parameters of the model. Initializing with a config file does not load the weights associated with the model, only the configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. aN Args: input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, sequence_length)`): Float values mel features extracted from the raw speech waveform. Raw speech waveform 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_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] attention_mask (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing *SpecAugment* data augmentation on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are decoder input IDs?](../glossary#decoder-input-ids) Whisper uses the `decoder_start_token_id` as the starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last `decoder_input_ids` have to be input (see `past_key_values`). decoder_attention_mask (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*): Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also be used by default. If you want to change padding behavior, you should read [`modeling_whisper._prepare_decoder_attention_mask`] and modify to your needs. See diagram 1 in [the BART paper](https://arxiv.org/abs/1910.13461) for more information on the default strategy. head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules in the encoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. decoder_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules in the decoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the cross-attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`) `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. decoder_inputs_embeds (`torch.FloatTensor` of shape `(batch_size, target_sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `decoder_input_ids` you can choose to directly pass an embedded representation. If `past_key_values` is used, optionally only the last `decoder_inputs_embeds` have to be input (see `past_key_values`). This is useful if you want more control over how to convert `decoder_input_ids` indices into associated vectors than the model's internal embedding lookup matrix. use_cache (`bool`, *optional*): If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see `past_key_values`). output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. aY Args: input_features (`torch.FloatTensor` of shape `(batch_size, feature_size, sequence_length)`): Float values mel features extracted from the raw speech waveform. Raw speech waveform 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_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules in the encoder. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. encoder_outputs (`tuple(tuple(torch.FloatTensor)`, *optional*): Tuple consists of (`last_hidden_state`, *optional*: `hidden_states`, *optional*: `attentions`) `last_hidden_state` of shape `(batch_size, sequence_length, hidden_size)`, *optional*) is a sequence of hidden-states at the output of the last layer of the encoder. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. csXeZdZdZedfdd ZddZejddd Z ejd d d Z dddZ Z S)rz Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`WhisperEncoderLayer`]. Args: config: WhisperConfig rcstj|_j|_j}j|_j|_j |_ j rJt |nd|_ tj|j|ddd|_tj||dddd|_t|j ||_tfddtjD|_tj|_d |_|dS) Nr8rr) kernel_sizepaddingrd)rZstridercsg|] }tqSr )rrJrr r!rMGsz+WhisperEncoder.__init__..F)r~rrZencoder_layerdrop layerdroprZ num_mel_binsrr|max_source_positionsscale_embeddingmathsqrt embed_scalerrconv1conv2rembed_positions ModuleListrWZencoder_layerslayersr layer_normr post_init)rrrrrr!r7s  zWhisperEncoder.__init__cCs|D] }d|_qd|_dS)NF) parametersZ requires_gradZ_requires_grad)rparamr r r!_freeze_parametersNs z!WhisperEncoder._freeze_parametersr@cCs|jSr}rrr r r!get_input_embeddingsSsz#WhisperEncoder.get_input_embeddingsrcCs ||_dSr}rrrr r r!set_input_embeddingsVsz#WhisperEncoder.set_input_embeddingsNcsdk r n|jj|dk r |n|jj}|dk r4|n|jj}tj||}tj||}| ddd}|j j }||} tjj | |j |j d} |rdnd} rdnd} |dk r|dt|jkstdt|jd|dd t|jD]\} } |r| | f} d }|j r6tg}||jkr6d }|rBd }nr|jr|j rfd d}tjj|| | d|dk r|| nd}n"| | d|dk r|| ndd}|d} r| |df} q|| } |r| | f} |stdd| | | fDSt| | | dS)aK Args: input_features (`torch.LongTensor` of shape `(batch_size, feature_size, sequence_length)`): Float values of mel features extracted from the raw speech waveform. Raw speech waveform 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_features`, the [`AutoFeatureExtractor`] should be used for extracting the mel features, padding and conversion into a tensor of type `torch.FloatTensor`. See [`~WhisperFeatureExtractor.__call__`] attention_mask (`torch.Tensor`)`, *optional*): Whisper does not support masking of the `input_features`, this argument is preserved for compatibility, but it is not used. By default the silence in the input log mel spectrogram are ignored. head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. Nrrdrrr z&The head_mask should be specified for  layers, but it is for .FT)NNcsfdd}|S)Ncs|fSr}r rb)rrr r!custom_forwardszMWhisperEncoder.forward..create_custom_forward..custom_forwardr rrrrr!create_custom_forwardsz5WhisperEncoder.forward..create_custom_forward)rrcss|]}|dk r|VqdSr}r rKvr r r! sz)WhisperEncoder.forward..last_hidden_stater attentions)rroutput_hidden_statesuse_return_dictrrfZgelurrpermuterrrrr.rYrAssertionError enumerater)rRrrutils checkpointrtupler )rrr> head_maskrr return_dict inputs_embedsZ embed_posrZencoder_statesZall_attentionsidxZ encoder_layerZto_dropdropout_probability layer_outputsrr rr!rYsn"         zWhisperEncoder.forward)NNNNN) rrrrrrrrModulerrrrr r rr!r.src sHeZdZdZedfdd ZddZddZd d Zdd d Z Z S)rz Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`WhisperDecoderLayer`] Args: config: WhisperConfig rcstj|_j|_j|_j|_j|_j rFt j nd|_ tjj |j|_t|jj |_tfddtjD|_tj |_d|_|dS)Nr8csg|] }tqSr )rrJrr r!rMsz+WhisperDecoder.__init__..F)r~rrZdecoder_layerdroprrr|max_target_positionsrrrrrrrr vocab_size embed_tokensryrrrWdecoder_layersrrrrrrrrr!rs  zWhisperDecoder.__init__cCs|jSr}rrr r r!rsz#WhisperDecoder.get_input_embeddingscCs ||_dSr}rrr r r!rsz#WhisperDecoder.set_input_embeddingscCsXd}|ddkr$t||j|j|d}|dk rTt||j|dd}|dkrL|n||}|S)Nrr)r%r&)r5)r7r$r%r;)rr> input_shaper r&Zcombined_attention_maskZexpanded_attn_maskr r r!_prepare_decoder_attention_masks z.WhisperDecoder._prepare_decoder_attention_maskNc  sdk r n|jj| dk r | n|jj} dk r4n|jj| dk rH| n|jj} |dk rj|dk rjtdnD|dk r|} |d| d}n"|dk r|dd} ntd|dk r|ddjdnd} |dkr| |}| || || }|dk r|j || d}n|j || d}||}t j j||j|jd}|jrX|jrXrXtd d | rbd nd}rpd nd}r|dk rd nd}rd nd}t||gd d gD]V\}}|dk r|dt|jkstd|dt|jd|ddqt|jD]D\}}| r(||f7}|jrJtg}||jkrJq|dk r\||nd}|jr|jrćfdd}tjj|||||d|dk r||nd|dk r||ndd}n<|||||dk r||nd|dk r||nd|d}|d}r&||rdndf7}r||df7}|dk r||df7}q||}| rp||f7}rz|nd}| stdd|||||fDSt|||||dS)aA Args: input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`): Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide it. Indices can be obtained using [`WhisperTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*): Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`: - 1 for tokens that are **not masked**, - 0 for tokens that are **masked**. [What are attention masks?](../glossary#attention-mask) encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*): Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention of the decoder. head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*): Mask to nullify selected heads of the attention modules in encoder to avoid performing cross-attention on hidden heads. Mask values selected in `[0, 1]`: - 1 indicates the head is **not masked**, - 0 indicates the head is **masked**. past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`. Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `decoder_input_ids` of shape `(batch_size, sequence_length)`. inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*): Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This is useful if you want more control over how to convert `input_ids` indices into associated vectors than the model's internal embedding lookup matrix. output_attentions (`bool`, *optional*): Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned tensors for more detail. output_hidden_states (`bool`, *optional*): Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for more detail. return_dict (`bool`, *optional*): Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. NzTYou cannot specify both decoder_input_ids and decoder_inputs_embeds at the same timerzEYou have to specify either decoder_input_ids or decoder_inputs_embedsrrd)r&rz^`use_cache = True` is incompatible with gradient checkpointing. Setting `use_cache = False`...Fr r cross_attn_head_maskThe `z` should be specified for rrcsfdd}|S)Ncs|fSr}r r)rrrr r!rszMWhisperDecoder.forward..create_custom_forward..custom_forwardr rrrrr!rsz5WhisperDecoder.forward..create_custom_forward)r>rrrrrrrrcss|]}|dk r|VqdSr}r rr r r!rsz)WhisperDecoder.forward..)rpast_key_valuesrrcross_attentions) rrrrrrr.r/rrrrrrfrrrloggerZ warning_onceziprYrrrr)rRrrr rr r )rrr>rr rrr rrrr rr&Z positionsrZall_hidden_statesZall_self_attnsZall_cross_attentionsZnext_decoder_cacheZ attn_maskZ mask_namerZ decoder_layerrrrrZ next_cacher rr!rsG    $         zWhisperDecoder.forward) NNNNNNNNNNN) rrrrrrrrrrrr r rr!rs"rzUThe bare Whisper Model outputting raw hidden-states without any specific head on top.cseZdZedfdd ZddZddZdd Zd d Zd d Z de j e e j dddZeeeeedde e j e e j e e j e e j e e je e je e je eee j e eee j e ee j e ee ee ee eeee jefdddZZS) WhisperModelrcs,t|t||_t||_|dSr})r~rrencoderrdecoderrrrr r!rs   zWhisperModel.__init__cCs|jjSr}r"rrr r r!rsz!WhisperModel.get_input_embeddingscCs ||j_dSr}r#rr r r!rsz!WhisperModel.set_input_embeddingscCs|jSr})r!rr r r! get_encoderszWhisperModel.get_encodercCs|jSr})r"rr r r! get_decoderszWhisperModel.get_decodercCs|jdSz Calling this function will disable the gradient computation for the Whisper encoder so that its parameters will not be updated during training. Nr!rrr r r!freeze_encoderszWhisperModel.freeze_encoderN)rr>cCst|jdds|S|\}}}|jjdkr|jrt||f|jj|jj||jjd}tj ||j tj d}|dddf d|d}d||<|jj dkr|jrt||f|jj |jj|jjd}tj ||j tj 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_augmentTr)r<r=r>r?)r%r$Nr)r<r=r?)getattrrr.Zmask_time_probrraZmask_time_lengthZmask_time_min_masksr)rr%r9r3Zmask_feature_probZmask_feature_lengthZmask_feature_min_masks)rrr>r^ hidden_sizerHZmask_time_indicesZmask_feature_indicesr r r!_mask_input_featuress0 z!WhisperModel._mask_input_features output_typer)rr>decoder_input_idsdecoder_attention_maskr decoder_head_maskrencoder_outputsrdecoder_inputs_embedsrrrr r@c Cs| dk r | n|jj} | dk r | n|jj} | dk r4| n|jj} |dk rH|n|jj}|dkr||j||d}|j||| | |d}nH|rt|tst|dt |dkr|dndt |dkr|dndd}|j |||d||| | | | | |d }|s||St |j |j |j|j|j|j |j|jd S) a{ Returns: Example: ```python >>> import torch >>> from transformers import AutoFeatureExtractor, WhisperModel >>> from datasets import load_dataset >>> model = WhisperModel.from_pretrained("openai/whisper-base") >>> feature_extractor = AutoFeatureExtractor.from_pretrained("openai/whisper-base") >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> inputs = feature_extractor(ds[0]["audio"]["array"], return_tensors="pt") >>> input_features = inputs.input_features >>> decoder_input_ids = torch.tensor([[1, 1]]) * model.config.decoder_start_token_id >>> last_hidden_state = model(input_features, decoder_input_ids=decoder_input_ids).last_hidden_state >>> list(last_hidden_state.shape) [1, 2, 512] ```N)r>r rrr rrrdr) rr>rr rrr rrrr )rrdecoder_hidden_statesdecoder_attentionsrencoder_last_hidden_staterencoder_attentions)rrrrrr+r!rr rYr"r rrrrr)rrr>r.r/r r0rr1rr2rrrr Zdecoder_outputsr r r!rsZ&zWhisperModel.forward)N)NNNNNNNNNNNNNN)rrrrrrrr$r%r(r) FloatTensorrrr+rWHISPER_INPUTS_DOCSTRINGrr _CONFIG_FOR_DOCrrr9rrrr r rr!r sX  +  r z^The Whisper Model with a language modeling head. Can be used for automatic speech recognition.cs\eZdZdZdgZedfdd ZddZdd Zd d Z d d Z e j dddZ ddZeeeeedd#eejeejeejeejeejeejeejeeeejeeeejeeejeejeeeeeeeeeeejefdddZd$eejeejdfdd Zd%ddZeddZ d&d!d"Z!Z"S)'WhisperForConditionalGenerationrzproj_out.weightrcs8t|t||_tj|j|jdd|_| dS)NFr) r~rr rrrrrproj_outrrrr r!rws  z(WhisperForConditionalGeneration.__init__cCs |jSr})rr$rr r r!r$sz+WhisperForConditionalGeneration.get_encodercCs |jSr})rr%rr r r!r%sz+WhisperForConditionalGeneration.get_decodercCs|jSr}r<rr r r!get_output_embeddingssz5WhisperForConditionalGeneration.get_output_embeddingscCs ||_dSr}r=)rZnew_embeddingsr r r!set_output_embeddingssz5WhisperForConditionalGeneration.set_output_embeddingsrcCs |jSr})rrrr r r!rsz4WhisperForConditionalGeneration.get_input_embeddingscCs|jjdSr&)rr!rrr r r!r(sz.WhisperForConditionalGeneration.freeze_encoderr,N)rr>r.r/r r0rr1rr2labelsrrrr r@cCs|dk r |n|jj}| dk r@|dkr@| dkr@t| |jj|jj}|j||||||||| | | | ||d}||d}d}| dk rt}| |j } || d|jj | d}|s|f|dd}|dk r|f|S|St |||j|j|j|j|j|j|jd S)a5 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the language modeling loss. Indices should either be in `[0, ..., config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`. Returns: Example: ```python >>> import torch >>> from transformers import AutoProcessor, WhisperForConditionalGeneration >>> from datasets import load_dataset >>> processor = AutoProcessor.from_pretrained("openai/whisper-tiny.en") >>> model = WhisperForConditionalGeneration.from_pretrained("openai/whisper-tiny.en") >>> ds = load_dataset("hf-internal-testing/librispeech_asr_dummy", "clean", split="validation") >>> inputs = processor(ds[0]["audio"]["array"], return_tensors="pt") >>> input_features = inputs.input_features >>> generated_ids = model.generate(inputs=input_features) >>> transcription = processor.batch_decode(generated_ids, skip_special_tokens=True)[0] >>> transcription ' Mr. Quilter is the apostle of the middle classes, and we are glad to welcome his gospel.' ```N) r>r.r1r/r r0rrr2rrrr rrr) losslogitsrr4r5rr6rr7)rrr"rrrr<rr0r%r/rr]r rr4r5rr6rr7)rrr>r.r/r r0rr1rr2r@rrrr rZ lm_logitsrAloss_fctoutputr r r!rsX1 z'WhisperForConditionalGeneration.forwardF)rb prompt_idsc  s\|dkr|j}|dk r0t|ds(td||_nd|_| dk r^t|dsPtd| } | |_|dk r~t|dsxtd||_d}t|jd r|jjdk r|jj}n.t|jd r|jjdk r|jj}n | d d}|dk s| dk s|dkrl| dk rlg}t|d r|j|j kr |j}n|jt krDd t |jd }n`|jt krdd |jd }n@t|jd k}td|jd|rtt n tt d|d|j |fn |dt|dr|jtkr|d |j|jfntd|jdtdn t|dr.|d |jdft|drl|jsl|rX|dddnd}|||jf|dk r|||_| dk r| ddk rtd| } | ^}}||jj d dd}| d|i| dddk rR| dt|7<| d|jjkrRtdt|d| dt|d| dd |jjd!|jjd | d dpd|j}||jfd"d#|D}d$d#t|D}||_|jrt|g}| r d%| d&<d%| d'<t|ddd(krtd)t|d*std+| d,dk r | d,|_tj ||||||f| }| rXt|d*rXt|d,d}|j!||j"|d-|d.<|S)/a Generates sequences of token ids for models with a language modeling head. Most generation-controlling parameters are set in `generation_config` which, if not passed, will be set to the model's default generation configuration. You can override any `generation_config` by passing the corresponding parameters to generate(), e.g. `.generate(inputs, num_beams=4, do_sample=True)`. For an overview of generation strategies and code examples, check out the [following guide](./generation_strategies). Parameters: inputs (`torch.Tensor` of varying shape depending on the modality, *optional*): The sequence used as a prompt for the generation or as model inputs to the encoder. If `None` the method initializes it with `bos_token_id` and a batch size of 1. For decoder-only models `inputs` should of in the format of `input_ids`. For encoder-decoder models *inputs* can represent any of `input_ids`, `input_values`, `input_features`, or `pixel_values`. generation_config (`~generation.GenerationConfig`, *optional*): The generation configuration to be used as base parametrization for the generation call. `**kwargs` passed to generate matching the attributes of `generation_config` will override them. If `generation_config` is not provided, the default will be used, which had the following loading priority: 1) from the `generation_config.json` model file, if it exists; 2) from the model configuration. Please note that unspecified parameters will inherit [`~generation.GenerationConfig`]'s default values, whose documentation should be checked to parameterize generation. logits_processor (`LogitsProcessorList`, *optional*): Custom logits processors that complement the default logits processors built from arguments and generation config. If a logit processor is passed that is already created with the arguments or a generation config an error is thrown. This feature is intended for advanced users. stopping_criteria (`StoppingCriteriaList`, *optional*): Custom stopping criteria that complement the default stopping criteria built from arguments and a generation config. If a stopping criteria is passed that is already created with the arguments or a generation config an error is thrown. This feature is intended for advanced users. prefix_allowed_tokens_fn (`Callable[[int, torch.Tensor], List[int]]`, *optional*): If provided, this function constraints the beam search to allowed tokens only at each step. If not provided no constraint is applied. This function takes 2 arguments: the batch ID `batch_id` and `input_ids`. It has to return a list with the allowed tokens for the next generation step conditioned on the batch ID `batch_id` and the previously generated tokens `inputs_ids`. This argument is useful for constrained generation conditioned on the prefix, as described in [Autoregressive Entity Retrieval](https://arxiv.org/abs/2010.00904). synced_gpus (`bool`, *optional*, defaults to `False`): Whether to continue running the while loop until max_length (needed for ZeRO stage 3) return_timestamps (`bool`, *optional*): Whether to return the timestamps with the text. This enables the `WhisperTimestampsLogitsProcessor`. task (`str`, *optional*): Task to use for generation, either "translate" or "transcribe". The `model.config.forced_decoder_ids` will be updated accordingly. language (`str`, *optional*): Language token to use for generation, can be either in the form of `<|en|>`, `en` or `english`. You can find all the possible language tokens in the `model.generation_config.lang_to_id` dictionary. is_multilingual (`bool`, *optional*): Whether or not the model is multilingual. prompt_ids (`torch.Tensor`, *optional*): Rank-1 tensor of token IDs created by passing text to [`~WhisperProcessor.get_prompt_ids`] that is provided as a prompt to each chunk. This can be used to provide or "prompt-engineer" a context for transcription, e.g. custom vocabularies or proper nouns to make it more likely to predict those words correctly. It cannot be used in conjunction with `decoder_start_token_id` as it overwrites this value. return_token_timestamps (`bool`, *optional*): Whether to return token-level timestamps with the text. This can be used with or without the `return_timestamps` option. To get word-level timestamps, use the tokenizer to group the tokens into words. kwargs (`Dict[str, Any]`, *optional*): Ad hoc parametrization of `generate_config` and/or additional model-specific kwargs that will be forwarded to the `forward` function of the model. If the model is an encoder-decoder model, encoder specific kwargs should not be prefixed and decoder specific kwargs should be prefixed with *decoder_*. Return: [`~utils.ModelOutput`] or `torch.LongTensor`: A [`~utils.ModelOutput`] (if `return_dict_in_generate=True` or when `config.return_dict_in_generate=True`) or a `torch.FloatTensor`. If the model is *not* an encoder-decoder model (`model.config.is_encoder_decoder=False`), the possible [`~utils.ModelOutput`] types are: - [`~generation.GreedySearchDecoderOnlyOutput`], - [`~generation.SampleDecoderOnlyOutput`], - [`~generation.BeamSearchDecoderOnlyOutput`], - [`~generation.BeamSampleDecoderOnlyOutput`] If the model is an encoder-decoder model (`model.config.is_encoder_decoder=True`), the possible [`~utils.ModelOutput`] types are: - [`~generation.GreedySearchEncoderDecoderOutput`], - [`~generation.SampleEncoderDecoderOutput`], - [`~generation.BeamSearchEncoderDecoderOutput`], - [`~generation.BeamSampleEncoderDecoderOutput`] Nno_timestamps_token_idabYou are trying to return timestamps, but the generation config is not properly set.Make sure to initialize the generation config with the correct attributes that are needed such as `no_timestamps_token_id`.For more details on how to generate the approtiate config, refer to https://github.com/huggingface/transformers/issues/21878#issuecomment-1451902363F lang_to_ida<The generation config is outdated and is thus not compatible with the `language` argumentto `generate`. Either set the language using the `forced_decoder_ids` in the model config, or update the generation config as per the instructions https://github.com/huggingface/transformers/issues/25084#issuecomment-1664398224 task_to_ida4The generation config is outdated and is thus not compatible with the `task` argumentto `generate`. Either set the task using the `forced_decoder_ids` in the model config, or update the generation config as per the instructions https://github.com/huggingface/transformers/issues/25084#issuecomment-1664398224forced_decoder_idslanguagez<|z|>rdzUnsupported language: z. Language should be one of: rr)rNtaskrz3`task is not supported. The task should be one of `rAZ transcriberrrzfWhen specifying `prompt_ids`, you cannot also specify `decoder_start_token_id` as it gets overwritten.Zmax_new_tokensz)The length of the sliced `prompt_ids` is z, and the `max_new_tokens` zP. Thus, the combined length of the sliced `prompt_ids` and `max_new_tokens` is: z@. This exceeds the `max_target_positions` of the Whisper model: z. You should either reduce the length of your prompt, or reduce the value of `max_new_tokens`, so that their combined length is less that cSsg|] \}}|qSr r )rKZ_ranktokenr r r!rMsz.cSsg|]\}}|d|fqS)rr )rKZrankrLr r r!rMsTrZreturn_dict_in_generate translatez@Token-level timestamps may not be reliable for task 'translate'.alignment_headszModel generation config has no `alignment_heads`, token-level timestamps not available. See https://gist.github.com/hollance/42e32852f24243b748ae6bc1f985b13a on how to add this property to the generation config. num_frames)rOZtoken_timestamps)#generation_confighasattrrreturn_timestampslowerrJrKrrIgetrGkeysrvaluesrYlistr[rrHrFrVrupdatepoprrr r)rwarningrOr~generate_extract_token_timestampsrN)rrbrPZlogits_processorZstopping_criteriaZprefix_allowed_tokens_fnZ synced_gpusrRrKrJZis_multilingualrEZreturn_token_timestampskwargsrIZlanguage_tokenZis_language_coderrZtext_prompt_idsZnon_prompt_forced_decoder_idsrrOrr r!r[si        $ ,      >         z(WhisperForConditionalGeneration.generatecKs,|dk r|ddddf}||||ddS)Nr)r1rr.rr/r )rr.rrr1r>r]r r r!prepare_inputs_for_generations z=WhisperForConditionalGeneration.prepare_inputs_for_generationcs.d}|D] }|tfdd|Df7}q|S)Nr c3s"|]}|d|jVqdS)rN)Z index_selectr0r%)rKZ past_statebeam_idxr r!rszAWhisperForConditionalGeneration._reorder_cache..)r )rr`Zreordered_pastZ layer_pastr r_r!_reorder_caches z.WhisperForConditionalGeneration._reorder_cache{Gz?csBgt|jjD](tjfdd|jDddqtfdd|D}|ddddg}|d k r|d d |df}tj |d d d d\}}|||}t ||jj }|j dd}tj |jtjd} t| jdD]`} t||  \} } tjt| dddt} | | |}t|| | dd f<q| S)a Calculates token-level timestamps using the encoder-decoder cross-attentions and dynamic time-warping (DTW) to map each output token to a position in the input audio. If `num_frames` is specified, the encoder-decoder cross-attentions will be cropped before applying DTW. Returns: tensor containing the timestamps in seconds for each predicted token csg|] }|qSr r )rKx)rqr r!rM&szMWhisperForConditionalGeneration._extract_token_timestamps..rdr'cs$g|]\}}|dd|fqSr}r )rKlh)rr r!rM*srrrN.rTF)r(ZkeepdimZunbiasedrN)rr)Zconstant_values)rWrrr[r)r1rstackrZstd_meanrjZmedian_filter_widthrZ zeros_like sequencesrlrrxdoublecpunumpyrPrgdiffZastyper9r)rZgenerate_outputsrNZtime_precisionrOweightsrrrkZ timestampsZ batch_idxrvrwZjumpsZ jump_timesr )rrqr!r\s$ &   z9WhisperForConditionalGeneration._extract_token_timestamps)NNNNNNNNNNNNNNN) NNNNNFNNNNNN)NNNN)rbN)#rrrrZ_tied_weights_keysrrr$r%r>r?rrrr(rr9rr r:rr)r8rrrr9rrr[r^ staticmethodrar\rr r rr!r;os  b    r;z Whisper Encoder Model with a sequence classification head on top (a linear layer over the pooled output) for tasks like SUPERB Keyword Spotting. c seZdZfddZddZejdddZejdd d Ze e e e e d deejeejeeeejeejeeeeeeeeeje fd ddZZS)WhisperForAudioClassificationcslt|t||_|jd}|jr 1` a classification loss is computed (Cross-Entropy). Returns: Example: ```python >>> import torch >>> from transformers import AutoFeatureExtractor, WhisperForAudioClassification >>> from datasets import load_dataset >>> feature_extractor = AutoFeatureExtractor.from_pretrained("sanchit-gandhi/whisper-medium-fleurs-lang-id") >>> model = WhisperForAudioClassification.from_pretrained("sanchit-gandhi/whisper-medium-fleurs-lang-id") >>> ds = load_dataset("google/fleurs", "all", split="validation", streaming=True) >>> sample = next(iter(ds)) >>> inputs = feature_extractor( ... sample["audio"]["array"], sampling_rate=sample["audio"]["sampling_rate"], return_tensors="pt" ... ) >>> input_features = inputs.input_features >>> with torch.no_grad(): ... logits = model(input_features).logits >>> predicted_class_ids = torch.argmax(logits).item() >>> predicted_label = model.config.id2label[predicted_class_ids] >>> predicted_label 'Afrikaans' ```Nr3rr'rr)rArBrr)rrrrr!ror)rfrrfrrqr/rTrrrrtrr0r%rsrrr)rrr r1r@rrr rZ norm_weightsZ pooled_outputrBrArCrDr r r!resD/    z%WhisperForAudioClassification.forward)NNNNNNN)rrrrr(rrrrr WHISPER_ENCODER_INPUTS_DOCSTRINGrrr:rr)rrrr8r9rrrr r rr!rnCs.  rn)r)N)Nr)FrrtypingrrrrjrPr)Ztorch.utils.checkpointrZtorch.nnrZ activationsrZgeneration.logits_processr Zmodeling_outputsr r r r rZmodeling_utilsrrrrrrZconfiguration_whisperrZtokenization_whisperrrZ get_loggerrrr:Z_CHECKPOINT_FOR_DOCZ%WHISPER_PRETRAINED_MODEL_ARCHIVE_LISTrrBr"Sizer$r%r7r;rrZndarrayrarjrxrryrrrrrZWHISPER_START_DOCSTRINGr9rurrr r;rnr r r r!s          w6EuU%S