o i[@sddlZddlZddlmZmZmZddlZddlmZddlm Z ddl m Z ddl m Z ddlmZdd lmZmZmZdd lmZmZdd lmZdd lmZmZmZmZdd lmZddl m!Z!m"Z"ddl#m$Z$m%Z%e&e'Z(   d-dej)dej*dej*dej*deej*dee+de+deej*fddZ,Gdddej)Z-Gdd d eZ.eGd!d"d"eZ/ed#d$Gd%d&d&e/Z0Gd'd(d(ej)Z1ed)d$Gd*d+d+e/e Z2gd,Z3dS).N)CallableOptionalUnion)nn)ACT2FN)Cache)GenerationMixin)GradientCheckpointingLayer)BaseModelOutputBaseModelOutputWithPastCausalLMOutputWithPast)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack)TransformersKwargsauto_docstringcan_return_tuplelogging)check_model_inputs) AutoModelAutoModelForCausalLM) VoxtralConfigVoxtralEncoderConfigmodulequerykeyvalueattention_maskscalingdropout head_maskc Ks|dur |dd}t||dd|} |dur5|jdkr5| |ddddddd|jdf} tjj| dd} |durK| | dddd} tjj | ||j d } t| |} | dd } | | fS) Nrr)dimrptraining) sizetorchmatmul transposendimshaper functionalsoftmaxviewr#r, contiguous) rrrr r!r"r#r$kwargs attn_weights attn_outputr:`/home/ubuntu/.local/lib/python3.10/site-packages/transformers/models/voxtral/modeling_voxtral.pyeager_attention_forward-s * r<cseZdZdZ      ddededed ed ed ed eed eeffdd Z de j dedefddZ   dde j dee j dee j dede e j ee j ee e j ff ddZZS)VoxtralAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrFTN embed_dim num_headsr# is_decoderbias is_causal layer_idxconfigc st||_||_||_|||_||_|j||jkr*td|jd|d|jd|_||_ ||_ |durG|rGt 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.FrA)super__init__r>r?r#head_dimrD ValueErrorr"r@rBlogger warning_once __class____name__rCrLineark_projv_projq_projout_proj) selfr>r?r#r@rArBrCrDrLr:r;rGNs0     zVoxtralAttention.__init__tensorseq_lenbszcCs ||||j|jddS)Nrr)r5r?rHr0r6)rSrUrVrWr:r:r;_shapevs zVoxtralAttention._shape hidden_statesr!layer_head_maskoutput_attentionsreturncKs|\}}}||||j||} |||d|} |||d|} t} |jjdkr6t |jj} | || | | |f|j sBdn|j d||d|\} }| ||d } || } | |fS)z#Input shape: Batch x Time x Channelr%eagerr?)r#r"r[r$)r-rXrQr"rOrPr<rD_attn_implementationrr,r#reshaper6rR)rSrYr!rZr[r7rWtgt_len_ query_states key_states value_statesattention_interfacer9r8r:r:r;forwardys0     zVoxtralAttention.forward)rFTFNN)NNF)rM __module__ __qualname____doc__intfloatboolrrrGr.TensorrXtuplerg __classcell__r:r:rTr;r=KsP (r=c sLeZdZdeffdd Z d dejdejdejded ejf d d ZZ S) VoxtralEncoderLayerrDcst|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#rD)rFrGd_modelr>r=encoder_attention_headsattention_dropout self_attnr LayerNormself_attn_layer_normr#ractivation_function activation_fnactivation_dropoutrNencoder_ffn_dimfc1fc2final_layer_normrSrDrTr:r;rGs  zVoxtralEncoderLayer.__init__FrYr!rZr[r\cCs|}||}|j||||d\}}tjj||j|jd}||}|}||}|||}tjj||j |jd}| |}tjj||j|jd}||}|j t j krgt |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. )rYr!rZr[r*i)minmax)rwrurr3r#r,r~ryr|rzr}dtyper.float16finforclamp)rSrYr!rZr[residualr8 clamp_valuer:r:r;rgs*     zVoxtralEncoderLayer.forward)F) rMrhrirrGr.rnrmrgrpr:r:rTr;rqsrqc@sFeZdZUeed<dZdZdZdZdZ dZ dZ dZ dZ dZddZdS)VoxtralPreTrainedModelrDmodelTNpast_key_valuescCst|jdr |jjn|jjj}t|tjtjfr0|jj j d|d|j dur.|j j dSdSt|tj rE|jj d|j j dSt|tjrd|jj j d|d|jdurf|jj |j dSdSdS)Ninitializer_ranger)meanstdr^)hasattrrDr audio_config isinstancerrNConv1dweightdatanormal_rAzero_rvfill_ Embedding padding_idx)rSrrr:r:r; _init_weightss$      z$VoxtralPreTrainedModel._init_weights)rMrhrir__annotations__base_model_prefixsupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa_supports_flex_attn_supports_cache_class_supports_attention_backend_can_compile_fullgraphrr:r:r:r;rs  rz: The Voxtral encoder, which is a Whisper encoder. ) custom_introcseZdZUdZeed<dZdgZee dZ deffdd Z dd Z d e jfd d Zd e jfddZe ddeefddZdejfddZZS)VoxtralEncoderz Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a [`VoxtralEncoderLayer`]. Args: config: VoxtralEncoderConfig rDinput_featuresrq) attentionsrYcstj|_j|_j}j|_j|_j |_ j r%t |nd|_ tj|j|ddd|_tj||dddd|_t|j ||_|jdtfdd tjD|_tj|_tjddd |_d|_|dS) Nr^rr) kernel_sizepaddingr)rstriderFcsg|]}tqSr:)rq).0rbrDr:r; 0sz+VoxtralEncoder.__init__..)r)rFrGr#encoder_layerdrop layerdroprr num_mel_bins pad_token_idrmax_source_positionsscale_embeddingmathsqrt embed_scalerrconv1conv2rembed_positionsrequires_grad_ ModuleListrangeencoder_layerslayersrv layer_norm AvgPool1d avg_poolergradient_checkpointing post_init)rSrDr>rTrr;rGs"    zVoxtralEncoder.__init__cCs|D]}d|_qd|_dS)NF) parameters requires_grad_requires_grad)rSparamr:r:r;_freeze_parameters9s  z!VoxtralEncoder._freeze_parametersr\cCs|jSNrrSr:r:r;get_input_embeddings>sz#VoxtralEncoder.get_input_embeddingsr cCs ||_dSrrrSr r:r:r;set_input_embeddingsA z#VoxtralEncoder.set_input_embeddingsNr7c Ks|jj|jjd|jjd}|jd|kr(td|d|jdd|d|j|jjj |jjj d}t j ||}t j ||}|ddd }|jj}|||j }t j j||j|jd }t|jD]\}} | ||d d } | d}qj||}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 .)rdevicerrr*N)r!rZ)last_hidden_state)rDrrrrr2rItorrrrr3gelupermuterr#r, enumeraterrr ) rSrr!r7expected_seq_length inputs_embeds embed_posrYidx encoder_layer layer_outputsr:r:r;rgDs.   zVoxtralEncoder.forward input_lengthscCs(|ddd}|ddd}||fS)zs Computes the output length of the convolutional layers and the output length of the audio encoder rrr:)rSroutput_lengthsr:r:r; _get_feat_extract_output_lengthsusz/VoxtralEncoder._get_feat_extract_output_lengthsr)rMrhrirjrrmain_input_namerr=rq_can_record_outputsrGrrModulerrrrrrgr. LongTensorrrpr:r:rTr;rs$   0rcs*eZdZdeffdd ZddZZS)VoxtralMultiModalProjectorrDcsNttj|jj|jjdd|_t |j |_ tj|jj|jjdd|_ dS)NFrE) rFrGrrNrintermediate_size text_config hidden_sizelinear_1rprojector_hidden_actactlinear_2rrTr:r;rGs  z#VoxtralMultiModalProjector.__init__cCs"||}||}||}|Sr)rrr)rSaudio_featuresrYr:r:r;rgs   z"VoxtralMultiModalProjector.forward)rMrhrirrGrgrpr:r:rTr;r~srzs The Voxtral model, which consists of Whisper encoder, a multi-modal projector and a LLama language model. cs4eZdZdgZddiZddgdgfiZdgZfddZd d Zd d Z d dZ ddZ ddZ ddZ dejfddZdejfddZee          d+deejdeejdeejdeejdeed eejd!eejd"eed#eejd$eeejfd%eed&efd'd(Zfd)d*ZZ S),VoxtralForConditionalGenerationzlm_head.weightlm_head colwise_reprYlogitsrcsHt||jj|_t|j|_t|j|_ t ||_ | dSr) rFrGr vocab_sizer from_configr audio_towerrlanguage_modelrmulti_modal_projectorrrrTr:r;rGs    z(VoxtralForConditionalGeneration.__init__cC |jSr)rrrr:r:r;rrz4VoxtralForConditionalGeneration.get_input_embeddingscC|j|dSr)rrrr:r:r;rz4VoxtralForConditionalGeneration.set_input_embeddingscCrr)rget_output_embeddingsrr:r:r;rrz5VoxtralForConditionalGeneration.get_output_embeddingscCrr)rset_output_embeddings)rSnew_embeddingsr:r:r;rrz5VoxtralForConditionalGeneration.set_output_embeddingscCrr)r set_decoder)rSdecoderr:r:r;rrz+VoxtralForConditionalGeneration.set_decodercCrr)r get_decoderrr:r:r;rrz+VoxtralForConditionalGeneration.get_decoderrcCs0||}|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`rDrrr)rSr audio_outputsaudio_hidden_states audio_embedsr:r:r;get_audio_featuress  z2VoxtralForConditionalGeneration.get_audio_featurescCstdt||S)NzUThe method `get_audio_embeds` is deprecated. Please use `get_audio_features` instead.)warningswarn FutureWarningr)rSrr:r:r;get_audio_embedss z0VoxtralForConditionalGeneration.get_audio_embedsNr input_idsr! position_idsrrlabels use_cachecache_positionlogits_to_keepr7r\c  Ks|dur ||}|dur.|dur.||} ||jjkd} || |j| |j}|jd||||||| | 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."] ```Nr%)r!r rrr r r r r:) rrrDaudio_token_id unsqueezemasked_scatterrrr)rSrrr!r rrr r r r r7raudio_token_maskoutputsr:r:r;rgs*1    z'VoxtralForConditionalGeneration.forwardcsH|dd}|d}tj|i|}|dur"|ddkr"||d<|S)Nrr r)popgetrFprepare_inputs_for_generation)rSargsr7rr  model_inputsrTr:r;rs  z=VoxtralForConditionalGeneration.prepare_inputs_for_generation) NNNNNNNNNr)!rMrhri_tied_weights_keys_tp_plan_pp_plan_keep_in_fp32_modules_strictrGrrrrrrr. FloatTensorrrrrrrrnrrmrrkrrr rgrrpr:r:rTr;rsh       Hr)rrr)NrN)4rrtypingrrrr.r activationsr cache_utilsr generationr modeling_layersr modeling_outputsr r r modeling_utilsrrprocessing_utilsrutilsrrrr utils.genericrautorrconfiguration_voxtralrr get_loggerrMrJrrnrlr<r=rqrrrr__all__r:r:r:r;sh          Z?#q