a h5Z @sddlZddlmZmZmZddlZddlmZddlmZddl m Z ddl m Z ddl mZdd lmZmZmZdd lmZmZdd lmZdd lmZmZmZmZdd lmZddlm Z m!Z!ddl"m#Z#m$Z$e%e&Z'd&ej(ej)ej)ej)eej)ee*e*eej)dddZ+Gdddej(Z,GdddeZ-eGdddeZ.eddGddde.Z/Gd d!d!ej(Z0ed"dGd#d$d$e.e Z1gd%Z2dS)'N)CallableOptionalUnion)nn)ACT2FN)Cache)GenerationMixin)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPastCausalLMOutputWithPast)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)TransformersKwargsauto_docstringcan_return_tuplelogging)check_model_inputs) AutoModelAutoModelForCausalLM) VoxtralConfigVoxtralEncoderConfig)modulequerykeyvalueattention_maskscalingdropout head_maskc Ks|dur|dd}t||dd|} |durj|jdkrj| |ddddddd|jdf} tjj| dd} |dur| | dddd} tjj | ||j d } t| |} | dd } | | fS) Nrr)dimrptraining) sizetorchmatmul transposendimshaper functionalZsoftmaxviewr#r, contiguous) rrrr r!r"r#r$kwargs attn_weights attn_outputr9h/var/www/html/assistant/venv/lib/python3.9/site-packages/transformers/models/voxtral/modeling_voxtral.pyeager_attention_forward,s * r;c seZdZdZdeeeeeeeeeedfdd Z e j eed d d Z de j ee j ee j ee e j ee j ee e j fd d dZZS)VoxtralAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN) embed_dim num_headsr# is_decoderbias is_causal layer_idxconfigc st||_||_||_|||_||_|j||jkrTtd|jd|d|jd|_||_ ||_ |dur|rt d|j jd||_tj||dd|_tj|||d|_tj|||d|_tj|||d|_dS) Nz;embed_dim must be divisible by num_heads (got `embed_dim`: z and `num_heads`: z).r&zInstantiating a decoder z without passing `layer_idx` is not recommended and will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` when creating this class.Fr@)super__init__r=r>r#head_dimrC ValueErrorr"r?rAloggerZ warning_once __class____name__rBrLineark_projv_projq_projout_proj) selfr=r>r#r?r@rArBrCrJr9r:rFMs0     zVoxtralAttention.__init__)tensorseq_lenbszcCs ||||j|jddS)Nrr)r4r>rGr0r5)rQrSrTrUr9r9r:_shapeuszVoxtralAttention._shape hidden_statesr!layer_head_maskoutput_attentionsreturncKs|\}}}||||j||} |||d|} |||d|} t} |jjdkrlt |jj} | || | | |f|j sdn|j d||d|\} }| ||d } || } | |fS)z#Input shape: Batch x Time x Channelr%eagerr?)r#r"rZr$)r-rVrOr"rMrNr;rCZ_attn_implementationrr,r#reshaper5rP)rQrXr!rYrZr6rUZtgt_len_Z query_statesZ key_statesZ value_statesZattention_interfacer8r7r9r9r:forwardxs0      zVoxtralAttention.forward)rFTFNN)NNF)rK __module__ __qualname____doc__intfloatboolrrrFr.TensorrVtupler` __classcell__r9r9rRr:r<Js8(r<csBeZdZedfdd ZdejejejeejdddZZ S) VoxtralEncoderLayerrCcst|j|_t|j|j|j|d|_t |j|_ |j |_ t |j |_|j|_t|j|j|_t|j|j|_t |j|_dS)N)r=r>r#rC)rErFd_modelr=r<Zencoder_attention_headsZattention_dropout self_attnr LayerNormself_attn_layer_normr#rZactivation_function activation_fnactivation_dropoutrLZencoder_ffn_dimfc1fc2final_layer_normrQrCrRr9r:rFs  zVoxtralEncoderLayer.__init__FrWcCs|}||}|j||||d\}}tjj||j|jd}||}|}||}|||}tjj||j |jd}| |}tjj||j|jd}||}|j t j krt |j jd}t j|| |d}||fS)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. )rXr!rYrZr*i)minmax)rormrr3r#r,rtrprrrqrsdtyper.Zfloat16Zfinforwclamp)rQrXr!rYrZZresidualr7Z clamp_valuer9r9r:r`s*     zVoxtralEncoderLayer.forward)F) rKrarbrrFr.rgrfr`rir9r9rRr:rjsrjc@sFeZdZUeed<dZdZdZdZdZ dZ dZ dZ dZ dZddZdS)VoxtralPreTrainedModelrCmodelTNpast_key_valuescCst|jdr|jjn|jjj}t|tjtjfrZ|jj j d|d|j dur|j j nbt|tj r|jj d|j j n:t|tjr|jj j d|d|jdur|jj |j dS)Ninitializer_ranger)meanstdr])hasattrrCr} audio_config isinstancerrLConv1dweightdataZnormal_r@Zzero_rnZfill_ Embedding padding_idx)rQrrr9r9r: _init_weightss     z$VoxtralPreTrainedModel._init_weights)rKrarbr__annotations__Zbase_model_prefixZsupports_gradient_checkpointing_no_split_modulesZ_skip_keys_device_placementZ_supports_flash_attnZ_supports_sdpaZ_supports_flex_attnZ_supports_cache_classZ_supports_attention_backendZ_can_compile_fullgraphrr9r9r9r:rzs rzz: The Voxtral encoder, which is a Whisper encoder. )Z custom_introcseZdZUdZeed<dZdgZee dZ edfdd Z d d Z e jd d d Ze jdddZedeedddZejdddZZS)VoxtralEncoderz Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`VoxtralEncoderLayer`]. Args: config: VoxtralEncoderConfig rCinput_featuresrj)Z attentionsrXrkcstj|_j|_j}j|_j|_j |_ j rJt |nd|_ tj|j|ddd|_tj||dddd|_t|j ||_|jdtfdd tjD|_tj|_tjddd |_d|_|dS) Nr]rr) kernel_sizepaddingr)rstriderFcsg|] }tqSr9)rj).0r_rkr9r: /z+VoxtralEncoder.__init__..)r)rErFr#Zencoder_layerdropZ layerdroprlZ num_mel_binsZ pad_token_idrmax_source_positionsZscale_embeddingmathsqrtZ embed_scalerrconv1conv2rembed_positionsZrequires_grad_Z ModuleListrangeZencoder_layerslayersrn layer_normZ AvgPool1dZ avg_poolerZgradient_checkpointing post_init)rQrCr=rRrkr:rFs"   zVoxtralEncoder.__init__cCs|D] }d|_qd|_dS)NF) parametersZ requires_gradZ_requires_grad)rQparamr9r9r:_freeze_parameters8s z!VoxtralEncoder._freeze_parameters)r[cCs|jSNrrQr9r9r:get_input_embeddings=sz#VoxtralEncoder.get_input_embeddings)r cCs ||_dSrrrQr r9r9r:set_input_embeddings@sz#VoxtralEncoder.set_input_embeddingsN)r6c Ks|jj|jjd|jjd}|jd|krPtd|d|jdd|d|j|jjj |jjj d}t j ||}t j ||}|ddd }|jj}|||j }t j j||j|jd }t|jD]\}} | ||d d } | d}q||}t|d S)a 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*): Voxtral 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. rr%z:Qwen2Audio expects the mel input features to be of length z , but found z-. Make sure to pad the input mel features to .)rxdevicerrr*N)r!rY)last_hidden_state)rCrrrrr2rHtorrxrrr3ZgeluZpermuterr#r, enumeraterrr ) rQrr!r6Zexpected_seq_length inputs_embedsZ embed_posrXidxZ encoder_layerZ layer_outputsr9r9r:r`Cs.   zVoxtralEncoder.forward) input_lengthscCs(|ddd}|ddd}||fS)zs Computes the output length of the convolutional layers and the output length of the audio encoder rrr9)rQrZoutput_lengthsr9r9r: _get_feat_extract_output_lengthstsz/VoxtralEncoder._get_feat_extract_output_lengths)N)rKrarbrcrrZmain_input_namerr<rjZ_can_record_outputsrFrrModulerrrrrr`r. LongTensorrrir9r9rRr:rs"  0rcs*eZdZedfdd ZddZZS)VoxtralMultiModalProjectorrkcsNttj|jj|jjdd|_t |j |_ tj|jj|jjdd|_ dS)NFrD) rErFrrLrintermediate_size text_configZ hidden_sizelinear_1rZprojector_hidden_actactlinear_2rurRr9r:rF~s  z#VoxtralMultiModalProjector.__init__cCs"||}||}||}|Sr)rrr)rQZaudio_featuresrXr9r9r:r`s   z"VoxtralMultiModalProjector.forward)rKrarbrrFr`rir9r9rRr:r}srzs The Voxtral model, which consists of Whisper encoder, a multi-modal projector and a LLama language model. cseZdZdgZddiZddgdgfiZdgZfddZd d Zd d Z d dZ ddZ ddZ ddZ ejdddZeedeejeejeejeejeeeejeejeeeejeeejfeeed ddZfddZZS) VoxtralForConditionalGenerationzlm_head.weightZlm_headZ colwise_reprXZlogitsrcsHt||jj|_t|j|_t|j|_ t ||_ | dSr) rErFrZ vocab_sizer from_configr audio_towerrlanguage_modelrmulti_modal_projectorrrurRr9r:rFs    z(VoxtralForConditionalGeneration.__init__cCs |jSr)rrrr9r9r:rsz4VoxtralForConditionalGeneration.get_input_embeddingscCs|j|dSr)rrrr9r9r:rsz4VoxtralForConditionalGeneration.set_input_embeddingscCs |jSr)rget_output_embeddingsrr9r9r:rsz5VoxtralForConditionalGeneration.get_output_embeddingscCs|j|dSr)rset_output_embeddings)rQZnew_embeddingsr9r9r:rsz5VoxtralForConditionalGeneration.set_output_embeddingscCs|j|dSr)r set_decoder)rQdecoderr9r9r:rsz+VoxtralForConditionalGeneration.set_decodercCs |jSr)r get_decoderrr9r9r:rsz+VoxtralForConditionalGeneration.get_decoder)rcCs0||}|j}|d|jjj}||}|S)a This method is used to get the audio embeddings from input features (a log mel spectrogram), meaning inferring the audio encoder and the multi-modal projector. Args: input_features (`torch.FloatTensor`): 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__`] Returns: `torch.FloatTensor`: The audio embeddings. r%)rrr^rCrrr)rQrZ audio_outputsZaudio_hidden_states audio_embedsr9r9r:get_audio_embedss   z0VoxtralForConditionalGeneration.get_audio_embedsNr) input_idsrr! position_idsr|rlabels use_cachecache_positionlogits_to_keepr6r[c  Ks`|dur||}|dur:||} ||jjk} | || <|jf||||||| | d| }|S)aj Example: ```python >>> from transformers import VoxtralForConditionalGeneration, AutoProcessor >>> import torch >>> device = "cuda" if torch.cuda.is_available() else "cpu" >>> repo_id = "mistralai/Voxtral-Mini-3B-2507" >>> processor = AutoProcessor.from_pretrained(repo_id) >>> model = VoxtralForConditionalGeneration.from_pretrained(repo_id, dtype=torch.bfloat16, device_map=device) >>> conversation = [ { "role": "user", "content": [ { "type": "audio", "url": "https://huggingface.co/datasets/hf-internal-testing/dummy-audio-samples/resolve/main/dude_where_is_my_car.wav", }, {"type": "text", "text": "What can you tell me about this audio?"}, ], } ] >>> inputs = processor.apply_chat_template(conversation) >>> inputs = inputs.to(device, dtype=torch.bfloat16) >>> outputs = model.generate(**inputs, max_new_tokens=30) >>> processor.batch_decode(outputs[:, inputs.input_ids.shape[1]:], skip_special_tokens=True) ["This audio is a humorous conversation between two friends, likely in English, where one of them is trying to figure out what the other's tattoo says."] ```N)r!rr|rrrrr)rrrCZaudio_token_idr)rQrrr!rr|rrrrrr6rZaudio_token_maskoutputsr9r9r:r`s&1     z'VoxtralForConditionalGeneration.forwardcsH|dd}|d}tj|i|}|durD|ddkrD||d<|S)Nrrr)popgetrEprepare_inputs_for_generation)rQargsr6rrZ model_inputsrRr9r:rs   z=VoxtralForConditionalGeneration.prepare_inputs_for_generation) NNNNNNNNNr) rKrarbZ_tied_weights_keysZ_tp_planZ_pp_planZ_keep_in_fp32_modules_strictrFrrrrrrr.Z FloatTensorrrrrrrgrrfrrdrrr r`rrir9r9rRr:rsN  Fr)rzrr)NrN)3rtypingrrrr.rZ activationsrZ cache_utilsrZ generationr Zmodeling_layersr Zmodeling_outputsr r r Zmodeling_utilsrrZprocessing_utilsrutilsrrrrZ utils.genericrautorrZconfiguration_voxtralrrZ get_loggerrKrIrrgrer;r<rjrzrrr__all__r9r9r9r:sV         Z?#q