o ZhT@sddlmZddlmZmZmZmZmZddlZddl m Z ddl m Z ddl mZmZddlmZddlmZdd lmZdd lmZdd lmZdd lmZmZdd lmZmZddl m!Z!m"Z"ddl#m$Z$ddl%m&Z&m'Z'm(Z(m)Z)m*Z*m+Z+ddl,m-Z-ddl.m/Z/m0Z0ddl1m2Z2e*rddl3m4Z4ddl5m6Z6e+7e8Z9eGddde'Z:e(dde(Gddde"Z;edGdd d e j<Z=Gd!d"d"e j<Z>Gd#d$d$e j<Z?d%d&Z@dMd'd(ZAd)ejBd*eCd+ejBfd,d-ZD .dNd/e jd?d?e j<ZJGd@dAdAee&ZKe(dBdGdCdDdDe;eZLGdEdFdFe j<ZMe(GdGdHdHe;ZNe(dIdGdJdKdKe;e2ZOgdLZPdS)O) dataclass)CallableListOptionalTupleUnionN)ACT2FN)Cache DynamicCache)GenerationMixin)use_kernel_forward_from_hub)AttentionMaskConverter)FlashAttentionKwargs)GradientCheckpointingLayer)BaseModelOutputWithPastCausalLMOutputWithPast)ROPE_INIT_FUNCTIONSdynamic_rope_update)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack) LossKwargs ModelOutputauto_docstringcan_return_tupleis_torch_flex_attn_availablelogging) AutoModel) CsmConfigCsmDepthDecoderConfig)CsmGenerationMixin) BlockMask)make_flex_block_causal_maskc@seZdZUdZdZeejed<dZ ejed<dZ ee e ejed<dZ ee ejdfed<dZ ee ejdfed<dZeejed <dZejed <dZee e ejed <dZee ejdfed <dZee ejdfed <dZeejed<dS)CsmOutputWithPastaf Base class for the model autoregressive outputs. Args: loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction). logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax). 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)`) Contains pre-computed hidden-states (key and values in the self-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding. hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads. depth_decoder_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction) of the depth decoder model. depth_decoder_logits (`torch.FloatTensor` of shape `(batch_size, sequence_length, config.vocab_size)`): Prediction scores of the depth decoder (scores for each vocabulary token before SoftMax). depth_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)`) depth_decoder_hidden_states (`tuple(torch.FloatTensor)`, *optional*, returned when `output_hidden_states=True` is passed or when `config.output_hidden_states=True`): Tuple of `torch.FloatTensor` (one for the output of the embeddings, if the model has an embedding layer, + one for the output of each layer) of shape `(batch_size, sequence_length, hidden_size)`. Hidden-states of the model at the output of each layer plus the optional initial embedding outputs. depth_decoder_attentions (`tuple(torch.FloatTensor)`, *optional*, returned when `output_attentions=True` is passed or when `config.output_attentions=True`): Tuple of `torch.FloatTensor` (one for each layer) of shape `(batch_size, num_heads, sequence_length, sequence_length)`. backbone_loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned when `labels` is provided): Language modeling loss (for next-token prediction) of the backbone model. Nlosslogitspast_key_values. hidden_states attentionsdepth_decoder_lossdepth_decoder_logitsdepth_decoder_past_key_valuesdepth_decoder_hidden_statesdepth_decoder_attentions backbone_loss)__name__ __module__ __qualname____doc__r'rtorch FloatTensor__annotations__r(r)rr*r+r,r-r.r/r0r1r9r9S/var/www/auris/lib/python3.10/site-packages/transformers/models/csm/modeling_csm.pyr&6s ,r&z[ The bare Csm Model outputting raw hidden-states without any specific head on top. )Z custom_introc@sDeZdZeZdZdZdgZdgZdZ dZ dZ dZ dZ dZddZdS)CsmPreTrainedModelmodelTCsmDecoderLayerr)cCs|jj}t|tjr"|jjjd|d|jdur |jj dSdSt|tj rC|jjjd|d|j durA|jj|j  dSdSt|t ra|j }t|dD] }|jj|jd|dqQdSt|tro|jjddSdS)N)meanstdr g?)configZinitializer_range isinstancennLinearweightdataZnormal_biasZzero_ Embedding padding_idxCsmCodebooksHead num_codebooksrange CsmRMSNormZfill_)selfmoduler@rKir9r9r: _init_weightss&      z CsmPreTrainedModel._init_weightsN)r2r3r4r! config_classZbase_model_prefixZsupports_gradient_checkpointingZ_no_split_modulesZ_skip_keys_device_placementZ_supports_flash_attn_2Z_supports_sdpaZ_supports_cache_classZ_supports_quantized_cacheZ_supports_static_cacheZ_supports_attention_backendrQr9r9r9r:r;qs r;ZRMSNormcs.eZdZdfdd ZddZddZZS) rMư>cs&ttt||_||_dS)z9 CsmRMSNorm is equivalent to T5LayerNorm N)super__init__rC Parameterr6onesrEvariance_epsilon)rN hidden_sizeeps __class__r9r:rUs  zCsmRMSNorm.__init__cCsJ|j}|tj}|djddd}|t||j}|j||S)NrT)Zkeepdim) dtypetor6float32powr?ZrsqrtrXrE)rNr*Z input_dtypeZvariancer9r9r:forwards  zCsmRMSNorm.forwardcCst|jjd|jS)Nz, eps=)tuplerEshaperXrNr9r9r: extra_reprszCsmRMSNorm.extra_repr)rS)r2r3r4rUrbrf __classcell__r9r9r[r:rMsrMcs8eZdZddeffdd ZeeddZZ S)CsmRotaryEmbeddingNrAcstt|dr|jdur|jd|jd|_nd|_|j|_|j|_||_ t |j|_ | |j |\}|_ |j d|dd|j|_dS)N rope_scaling rope_typetypedefaultinv_freqF persistent)rTrUhasattrrigetrjZmax_position_embeddingsZmax_seq_len_cachedZoriginal_max_seq_lenrArZ rope_init_fnattention_scalingregister_bufferrmZoriginal_inv_freq)rNrAdevicermr[r9r:rUs   zCsmRotaryEmbedding.__init__c Cs|jddddf|jddd|j}|dddddf}t|jjtr6|jjdkr6|jjnd}t j |dd+|| dd}t j ||fdd }| |j}||j} Wdn1smwY|j|jd | j|jd fS) Nrr]r ZmpscpuF) device_typeenabledrdim)r^)rmfloatexpandrdr_rtrBrkstrr6Zautocast transposecatcosrrsinr^) rNx position_idsZinv_freq_expandedZposition_ids_expandedrvZfreqsZembrrr9r9r:rbs0&zCsmRotaryEmbedding.forwardN) r2r3r4r!rUr6Zno_gradrrbrgr9r9r[r:rhs rhc$eZdZfddZddZZS)CsmMLPcsxt||_|j|_|j|_tj|j|j|jd|_tj|j|j|jd|_ tj|j|j|jd|_ t |j |_ dS)NrG)rTrUrArYZintermediate_sizerCrDZmlp_bias gate_projup_proj down_projr Z hidden_actact_fnrNrAr[r9r:rUs zCsmMLP.__init__cCs$||||||}|Sr)rrrr)rNrrr9r9r:rbs zCsmMLP.forwardr2r3r4rUrbrgr9r9r[r:rs  rcCsH|dd|jddf}|d|jdddf}tj| |fddS)z*Rotates half the hidden dims of the input..Nr]rrx)rdr6r~)rx1Zx2r9r9r: rotate_halfsrcCsD||}||}||t||}||t||}||fS)aApplies Rotary Position Embedding to the query and key tensors. Args: q (`torch.Tensor`): The query tensor. k (`torch.Tensor`): The key tensor. cos (`torch.Tensor`): The cosine part of the rotary embedding. sin (`torch.Tensor`): The sine part of the rotary embedding. position_ids (`torch.Tensor`, *optional*): Deprecated and unused. unsqueeze_dim (`int`, *optional*, defaults to 1): The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. Returns: `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. ) unsqueezer)qkrrrZ unsqueeze_dimZq_embedZk_embedr9r9r:apply_rotary_pos_embs  rr*n_repreturncCs^|j\}}}}|dkr |S|dddddddddf|||||}||||||S)z This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch, num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim) r N)rdr{reshape)r*rbatchnum_key_value_headsslenhead_dimr9r9r: repeat_kvs 0rr>rOquerykeyvalueattention_maskscalingdropoutc Kst||j}t||j} t||dd|} |dur3|ddddddd|jdf} | | } tjj| dtj d |j } tjj | ||j d} t| | } | dd} | | fS)Nrrr])ryr^)ptrainingr )rnum_key_value_groupsr6matmulr}rdrC functionalZsoftmaxr`r_r^rr contiguous) rOrrrrrrkwargs key_states value_states attn_weights causal_mask attn_outputr9r9r:eager_attention_forward s & rcseZdZdZdedeffdd Z  ddejde ejejfd e ejd e e d e ej d e ed e eje eje e ejffddZZS) CsmAttentionz=Multi-headed attention from 'Attention Is All You Need' paperrA layer_idxcst||_||_t|d|j|j|_|j|j|_ |jd|_ |j |_ d|_ t j|j|j|j|jd|_t j|j|j|j|jd|_t j|j|j|j|jd|_t j|j|j|j|jd|_dS)NrgTr)rTrUrArgetattrrYZnum_attention_headsrrrrattention_dropoutZ is_causalrCrDZattention_biasq_projk_projv_projo_projrNrArr[r9r:rU*s(   zCsmAttention.__init__Nr*position_embeddingsrpast_key_valuecache_positionrrcKsH|jdd}g|d|jR}|||dd} |||dd} |||dd} |\} } t| | | | \} } |durW| | |d}|| | |j |\} } t }|j j dkrw|j j dkrq| ddrqtd nt|j j }||| | | |f|jsd n|j|jd |\}}|jg|dR}||}||fS) Nr]r r)rrreagersdpaoutput_attentionsFz`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.r>)rr)rdrrviewr}rrrupdaterrrA_attn_implementationrqlogger warning_oncerrrrrrr)rNr*rrrrrZ input_shapeZ hidden_shapeZ query_statesrrrrZ cache_kwargsZattention_interfacerrr9r9r:rbAs@      zCsmAttention.forward)NN)r2r3r4r5r!intrUr6Tensorrrr LongTensorrrrbrgr9r9r[r:r's(rcseZdZdedeffdd Z       ddejdeejd eej d ee d ee d ee d eej dee ejejfde ede ejee ejejfffddZZS)r=rArcsRt|j|_t||d|_t||_t|j|jd|_ t|j|jd|_ dS)N)rArrZ) rTrUrYr self_attnrmlprM rms_norm_epsinput_layernormpost_attention_layernormrr[r9r:rUus  zCsmDecoderLayer.__init__NFr*rrrr use_cacherrrrc Kst|} ||}|jd||||||||d| \}} | |}|} ||}||}| |}|f} |r8| | f7} | S)N)r*rrrrrrrr9)rrrr) rNr*rrrrrrrrZresidualZself_attn_weightsoutputsr9r9r:rbs.      zCsmDecoderLayer.forward)NNNFFNN)r2r3r4r!rrUr6rrrr boolrrrr7rbrgr9r9r[r:r=ts<    r=cs eZdZeZfddZddZddZee          d!de j d e e j d e e jd e e j d e ed e e j de ede ede ede e j deedeeeffddZ d"d ee jdfde jde jd edef ddZed e jdedede jde jdef dd ZZS)#CsmDepthDecoderModelcstj|_j|_tjjj|_ t fddt j D|_ tjjd|_td|_d|_tjjjdd|_|dS)Ncg|]}t|qSr9r=.0rrAr9r: z1CsmDepthDecoderModel.__init__..rrFr)rTrU pad_token_idrI vocab_sizerCrHrKZbackbone_hidden_size embed_tokens ModuleListrLnum_hidden_layerslayersrMrYrnormrh rotary_embgradient_checkpointingrDinputs_embeds_projector post_initrr[rr:rUs   zCsmDepthDecoderModel.__init__cC|jSrrrer9r9r:get_input_embeddingsz)CsmDepthDecoderModel.get_input_embeddingscC ||_dSrrrNrr9r9r:set_input_embeddings z)CsmDepthDecoderModel.set_input_embeddingsN input_idsbackbone_last_hidden_staterrr) inputs_embedsrroutput_hidden_statesrflash_attn_kwargsrc  KsX|durtjstdd}|dur|n|jj}| dur | n|jj} |dur*|n|jj}|du|duAr:t d|j rI|j rI|rItdd}|rR|durRt }| dur|dur^| nd} |duri|jdn|jd} |duru|jn|j}tj| | | |d} |durtj| ddd }||j}|||}| ddk}|dur||dddf<n tjs|rtd ||}|||| ||}|}| d}|||}| rd nd}|rd nd}|jd|jjD]'}| r||f7}||f|||||| |d | }|d}|r||df7}q||}| r||f7}t||r&|nd||d S)aJ backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*): The last hidden state of the backbone model. Such input is required when the first codebook token (the one generated by the backbone model) is provided in the `input_ids` argument. NzCustom `position_ids` were provided but will be ignored. CSM depth decoder automatically determines position_ids from `cache_position` and as it requires them to be identical across the batch, the provided position_ids will be ignored.z;You must specify exactly one of input_ids or inputs_embeds.X`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.Frr rt)minzvWhen the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference.r9rrrrrrrZlast_hidden_stater)r*r+)r6compilerZ is_compilingrrrArrr ValueErrorrrr get_seq_lengthrdrtarangeclamprrwarningr_update_causal_maskrrrrrr)rNrrrrr)rrrrrrpast_seen_tokensZinputs_seq_lengthrt codebook_idxsoffsetZinput_ids_are_first_codebookrr*rall_hidden_statesall_self_attns decoder_layer layer_outputsr9r9r:rbs               zCsmDepthDecoderModel.forwardFr$ input_tensorc C:|jjdkr|dur|dkr|SdS|jjdkr&t|tjr$t|}|S|dur.|nd}|dur7|jnd}|jjdkrO|sO|sOt j ||||j drOdS|j }|j d} |r^|} nt|tjri|j d n|| d} |j|| | |||j dd } |jjdkr|dur|jjd vr|st|j} t | | } | S NZflash_attention_2r>Zflex_attentionrFr)rZpast_key_values_lengthZ is_trainingr r])sequence_length target_lengthr^r batch_size)cudaZxpuZnpurAranyrBr6rr%rZis_compileablerZ_ignore_causal_mask_sdparr^rdZget_max_cache_shape5_prepare_4d_causal_attention_mask_with_cache_positionrtrkfinforZ_unmask_unattended rNrrrr)rrZusing_compilable_cacher^rrr min_dtyper9r9r:r7T           z(CsmDepthDecoderModel._update_causal_maskrrr^rc KD|dur|dkr|}|St|j}tj||f|||jd}|dkr+tj|dd}|tj||jd|ddk9}|ddddddf |ddd}|dur| }|j d} |ddddddd| f|ddddddf |j} | dk} |ddddddd| f | ||ddddddd| f<|S aM Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing. Args: attention_mask (`torch.Tensor`): A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape `(batch_size, 1, query_length, key_value_length)`. sequence_length (`int`): The sequence length being processed. target_length (`int`): The target length: when generating with static cache, the mask should be as long as the static cache, to account for the 0 padding, the part of the cache that is not filled yet. dtype (`torch.dtype`): The dtype to use for the 4D attention mask. cache_position (`torch.Tensor`): Indices depicting the position of the input sequence tokens in the sequence. batch_size (`torch.Tensor`): Batch size. N)Z fill_valuer^rtr )Zdiagonalrr]rryr6r rfullrtZtriurrr{clonerdr_Z masked_fill rrrr^rrrrrZ mask_lengthZ padding_maskr9r9r:r {,  $ 6  zJCsmDepthDecoderModel._prepare_4d_causal_attention_mask_with_cache_position) NNNNNNNNNNF)r2r3r4r"rRrUrrrrr6rrr7rr rrrrrrrbr staticmethodrr^r rgr9r9r[r:rs      x  Drcs&eZdZfddZdddZZS)rJcs0t||_tt|jd|||_dSNr )rTrUrKrCrVr6emptyrE)rNrYrKrr[r9r:rUs  zCsmCodebooksHead.__init__Ncsf|durjd}|jt|n |d}|j|fddtjdDtjddS)Nr c s2g|]}tjdd|ddf|jqSr)rCrZlinearT)rZ codebook_idxZcodebook_weightr*r9r:rs$z,CsmCodebooksHead.forward..rrx)rdrEr6rrLstack)rNr*rZ seq_lengthrr9rr:rbs    zCsmCodebooksHead.forwardrrr9r9r[r:rJs rJc@s eZdZdS)KwargsForCausalLMN)r2r3r4r9r9r9r:rs ra$ The CsmDepthDecoder Model transformer, with a [`CsmCodebooksHead`] on top, which can be seen a position-specific language modeling head, allowing to use a different linear layer for each codebook (e.g. position 0 is the first codebook and uses the first codebook head, etc.) c!sJeZdZdZdZdZfddZddZddZdd Z d d Z e e            dd e jdee jdee jdee jdeeeee jfdee jdee jdeedeedeedee jdeee jfdeedeeeffddZ    d d e jdeedee jdee jdee jf fdd ZZS)!CsmDepthDecoderForCausalLMNcs>t|t||_|j|_t|j|j|j|_| dSr) rTrUrr<rrJrYrKcodebooks_headrrr[r9r:rUs   z#CsmDepthDecoderForCausalLM.__init__cC|jjSrr<rrer9r9r:rz/CsmDepthDecoderForCausalLM.get_input_embeddingscC ||j_dSrr#rr9r9r:r z/CsmDepthDecoderForCausalLM.set_input_embeddingscCrrr<)rNdecoderr9r9r: set_decoderrz&CsmDepthDecoderForCausalLM.set_decodercCrrr'rer9r9r: get_decoderrz&CsmDepthDecoderForCausalLM.get_decoderrrrrrr)rlabelsrrrrlogits_to_keeprrc  Ks| dur| n|jj} | dur| n|jj} |jd|||||||| | | d | }|d}t| trA| dkr:tdd}n t| d}n| }||dd|ddf| durW| |nd}|}d}|dur}|dddf}|j d|d|jj |d| }t |||j |j |jdS) a backbone_last_hidden_state (`torch.FloatTensor` of shape `(batch_size, backbone_hidden_size)`, *optional*): The last hidden state of the backbone model. Such input is required when the first codebook token (the one generated by the backbone model) is provided in the `input_ids` argument. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked 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]`. N) rrrrr)rrrrrrr .)r(r+r shift_labels)r'r(r)r*r+r9)rArrr<rBrslicer!r loss_functionrrr)r*r+)rNrrrrr)rr+rrrrr,rrr* slice_indicesr(r'r-r9r9r:rbsT   & z"CsmDepthDecoderForCausalLM.forwardc  sHtj|||||fi|}|dddk}|s|d|d|S)Nrrrr)rTprepare_inputs_for_generationpop) rNrr)rrrr model_inputsZis_first_generation_stepr[r9r:r1<s   z8CsmDepthDecoderForCausalLM.prepare_inputs_for_generation) NNNNNNNNNNNrNNNN)r2r3r4_tied_weights_keysZ_tp_planZ_pp_planrUrrr)r*rrr6rrr7rrr rrrrrrrrbr1rgr9r9r[r:r s        Or cr)CsmBackboneModelEmbeddingscsDtt|j|j|j|_|jdt |j|jdddS)Naudio_tokens_offsetsFrn) rTrUrCrHrKrrYembed_audio_tokensrsr6rrr[r9r:rUTs  z#CsmBackboneModelEmbeddings.__init__cCs |||j}|jdd}|S)Nrrx)r8r7sum)rNrZ input_embedsr9r9r:rb[s z"CsmBackboneModelEmbeddings.forwardrr9r9r[r:r6Ss r6cs eZdZfddZddZddZee         d dee j d ee j d ee j d ee d ee j d eedeedeedee j deedefddZ d!d ee j dfde j de j d e def ddZed e j dedede jde j def ddZZS)"CsmBackboneModelcsvtj|_j|_t|_tfddt j D|_ t j jd|_td|_d|_|dS)Ncrr9rrrr9r:rirz-CsmBackboneModel.__init__..rrF)rTrUrrIrr6rrCrrLrrrMrYrrrhrrrrr[rr:rUcs    zCsmBackboneModel.__init__cCrrrrer9r9r:rrrz%CsmBackboneModel.get_input_embeddingscCrrrrr9r9r:rurz%CsmBackboneModel.set_input_embeddingsNrrrr)rrrrrrrc  Ks|dur|n|jj}|dur|n|jj}|dur|n|jj}|du|duAr*td|jr9|jr9|r9tdd}t |t dt fsFtd|durO| |}|rX|durXt }| durt|durd|nd} tj| | |jd|jd} |dur}| d}|||| ||} |} || |}|rd nd}|rd nd}|jd|jjD]&}|r|| f7}|| f| ||||| |d | }|d} |r||df7}q|| } |r|| f7}t| |r|nd||d S) a& input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks) or (batch_size, sequence_length)`): 1. (batch_size, sequence_length): corresponds to the input sequence prepared with the processor from the text prompt. Such input requires `input_values` to be provided so that audio can be encoded in codebook tokens and then merged with the text tokens. 2. (batch_size, sequence_length, num_codebooks): codebook tokens generated during the autoregressive decoding. Such input is not meant to be used by end users. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) Nz:You must specify exactly one of input_ids or inputs_embedsrFzBThe `past_key_values` should be either a `Cache` object or `None`.rr rr9rr)rArrrrrrrrrBrkr rr rr6rrdrtrrrrrrr)rNrrrr)rrrrrrrrr*rrrrrr9r9r:rbxsx             zCsmBackboneModel.forwardFr$rc Crrr r r9r9r:rrz$CsmBackboneModel._update_causal_maskrrr^rc Krrrrr9r9r:r (rzFCsmBackboneModel._prepare_4d_causal_attention_mask_with_cache_position) NNNNNNNNNr)r2r3r4rUrrrrrr6rrr r7rrrrrbrrrrr^r rgr9r9r[r:r:as     p  Dr:z The Csm model consists of two llama-like auto-regressive transformer models: a backbone model that predicts the first codebook token and a depth decoder that predicts the other codebook tokens. c#seZdZddgZfddZddZddZd d Zd d Zd dZ e fddZ fddZ    d0de ejde ejde ejde ejde ejf ddZ    d0dejde ede ejde ejde ejf fdd Zee             !d1dejde ejde ejde ejd"e ejde eeeejfde ejde ejd#e ed$e ed%e ede ejd&eeejfd'eedeeeffd(d)Ze dejd*ed+ed,ej!dejd-ef d.d/Z"Z#S)2CsmForConditionalGenerationz5backbone_model.embed_tokens.embed_audio_tokens.weightz'depth_decoder.model.embed_tokens.weightcspt||j|_tj|j|jdd|_t|j|j|_ t ||_ t |j|_t|j|_|dS)NFr)rTrUrrCrDrYlm_headrHZtext_vocab_sizeembed_text_tokensr:Z _from_configbackbone_modelr Zdepth_decoder_config depth_decoderrZ from_configZ codec_config codec_modelrrr[r9r:rUks   z$CsmForConditionalGeneration.__init__cCr"rr>rrer9r9r:rwr$z0CsmForConditionalGeneration.get_input_embeddingscCr%rrArr9r9r:rzr&z0CsmForConditionalGeneration.set_input_embeddingscCrrr<rer9r9r:get_output_embeddings}rz1CsmForConditionalGeneration.get_output_embeddingscCrrrB)rNZnew_embeddingsr9r9r:set_output_embeddingsrz1CsmForConditionalGeneration.set_output_embeddingscCs(|jjr||jjj|jjjdSdSr)rAZtie_codebooks_embeddingsZ_tie_or_clone_weightsr>rr8r?r<rer9r9r: _tie_weightss z(CsmForConditionalGeneration._tie_weightscs|ddrtj|i|\}}n tj|i|}dtfddt|jD}t|jjddi||D] }t |j|q?d|vrR||fS|S)NZoutput_loading_infoFdepth_decoder_cs(i|]\}}|r|d|qSr) startswith)rattrrprefix prefix_lenr9r: s  z?CsmForConditionalGeneration.from_pretrained..Z_from_model_config) rqrTfrom_pretrainedlenvarsgeneration_configitemsr?rdelattr)clsargsrr<Z loading_infodepth_decoder_attrsrHr[rIr:rMs   z+CsmForConditionalGeneration.from_pretrainedcsVd}|jj}|dd|D] \}}t|j|||qtj|i|dS)NrFZtransformers_version)r?rPZ to_diff_dictr2rQsetattrrTsave_pretrained)rNrTrrJrUrHrr[r9r:rWs   z+CsmForConditionalGeneration.save_pretrainedNr input_valuesinput_values_cutoffsr+rcs||}|durtj|d}||dk}||dk}tj||jd t |d}|| dk}g}t ||D]?\} } | | dk} t | jddD]+} | | } | | d} | d| | f}|j| d}|jdd}||dqPq=tdd |Dtfd d |D}|j|}|jj}||k}|j|}||||<tjdd|jjf|jtjd |jj}|j|d}||jjk}| |!d||<|dur| d dd|jj}||||<|d kj"dd}d||d|dddf<|}||dS)a Merges the input_ids and input_values to produce a single inputs_embeds tensor: 1 - Infers the codec model on the input_values to retreive codebook token. 2 - Embeds codebook tokens and places them at the correct positions in the inputs_embeds tensor. 3 - If labels are provided, expands them to match codebook dimensions and position the target codebook tokens in the inputs_embeds tensor. Args: input_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`): The input ids to embed. input_values (`torch.Tensor` of shape `(batch_size, channels, audio_sequence_length)`): The audio input values to embed. input_values_cutoffs (`torch.Tensor` of shape `(batch_size, max_num_audio)`): The cutoffs of the audio input values relative to its batch index, padded with -1 when no audio. Nr rrrr]r .css|]}|jdVqdS)rN)rdrelr9r9r: szQCsmForConditionalGeneration._merge_input_ids_with_input_values..c s,g|]}tj|ddd|jdfqS)r)rCrpadrdr[Zmax_audio_framesr9r:rs,zRCsmForConditionalGeneration._merge_input_ids_with_input_values..)rtr^iTas_tuple)rr+)#r=rCrr^diffr6rmaxrtr{rNrziprLrdr@encodeZ audio_codesr}appendrZget_audio_codes_maskrAaudio_token_idr>rrWrKlongZcodebook_eos_token_idZsqueezeZaudio_eos_token_idrepeatr9nonzero)rNrrXrYr+rZ audio_lengthsZinput_values_maskZaudio_tokens_listZbatch_input_valuesZbatch_input_values_cutoffsrPZ start_idxZend_idxZ audio_batchZ codec_outputsZ codebook_idsZbatched_audio_token_idsZaudio_codes_maskrhZaudio_token_maskZ audio_embedsZaudio_eos_frame_idsZaudio_eos_embedsZaudio_eos_token_maskZlabels_expandedZ depth_decoder_ignore_frames_idxsr9r_r:"_merge_input_ids_with_input_valuessV            z>CsmForConditionalGeneration._merge_input_ids_with_input_valuesr)rrrc  stjd |||||d|}|dur>|jdkr>|ddur>|j||d|d|dd}||d|ddd|S) N)rr)rrrrrrXrYr+)rrXrYr+)rr+rr9)rTr1ndimrqrlr) rNrr)rrrrr3 merged_inputsr[r9r:r1s(  z9CsmForConditionalGeneration.prepare_inputs_for_generationrrrrrr,rc Ks| dur| n|jj} | dur| n|jj} |dur/|jdkr/|||||}|d}|d}d}|jd|||||| | | | d |}|d}t| trPt| dn| }| |dd|ddf}d}d}d}d}|dur|dddddf}|j d|||jj d|}|ddddddfd kj d d }||d d|jj df}tjj|d dd}|jdd}||d|ddddf}||}|j||| | | d|d}|j}||}t|||||j|j|j|dur|jnd|dur|jnd|dur|jnd|dur |jd Sdd S)a input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length, num_codebooks) or (batch_size, sequence_length)`): 1. (batch_size, sequence_length): corresponds to the input sequence prepared with the processor from the text prompt. Such input requires `input_values` to be provided so that audio can be encoded in codebook tokens and then merged with the text tokens. 2. (batch_size, sequence_length, num_codebooks): codebook tokens generated during the autoregressive decoding. Such input is not meant to be used by end users. Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details. [What are input IDs?](../glossary#input-ids) input_values_cutoffs (`torch.Tensor` of shape `(batch_size, max_num_audio)`, *optional*): Specify the end positions of audio segments within each batch entry, relative to the concatenated audio input. If a batch entry has fewer segments than the maximum, it is padded with -1. For example, in a batch of 2 sequences where the first contains 2 audio segments of length l1, and the second contains 1 audio segment of length l2, the input_values_cutoffs would be: [[l1, 2 * l1], [l2, -1]]. labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*): Labels for computing the masked language modeling loss. Indices should be in `[config.audio_token_id, -100, -101]`. Requires targeted `input_values` to be provided as audio tokens will be infered from it using the `codec_model`. - `config.audio_token_id` indicates an audio frames (considering sequence length elements as frames) - `-100` will be ignored in the loss computation - `-101` indicates the audio frame will be used only for the backbone model (using the first codebook token as labels) Such labels can be prepared using `output_labels=True` when calling [`CsmProcessor`]. logits_to_keep (`int` or `torch.Tensor`, *optional*): Kept for compatibility. Does not support another value than: 1. `0`, which is equivalent to keeping all logits, used in the training regime 2. `1`, which is equivalent to keeping only the last logit, used in the generation regime Example: ```python >>> import torch >>> from transformers import CsmForConditionalGeneration, AutoProcessor >>> from datasets import load_dataset, Audio >>> model_id = "eustlb/csm-1b" >>> torch_device = "cuda" if torch.cuda.is_available() else "cpu" >>> processor = AutoProcessor.from_pretrained(model_id) >>> ds = load_dataset("hf-internal-testing/dailytalk-dummy", split="train") >>> # ensure the audio is 24kHz >>> ds = ds.cast_column("audio", Audio(sampling_rate=24000)) >>> conversation = [] >>> # prepare a conversation with text and corresponding audio >>> for text, audio, speaker_id in zip(ds[:4]["text"], ds[:4]["audio"], ds[:4]["speaker_id"]): ... conversation.append( ... { ... "role": f"{speaker_id}", ... "content": [{"type": "text", "text": text}, {"type": "audio", "path": audio["array"]}], ... } ... ) >>> inputs = processor.apply_chat_template( ... conversation, ... tokenize=True, ... return_dict=True, ... output_labels=True, ... ).to(torch_device) >>> model = CsmForConditionalGeneration.from_pretrained(model_id, device_map=torch_device) >>> output = model(**inputs) >>> output.loss.backward() ```Nrrr+) rrrr)rrrrrr)r(r+rr rbr]rx.rZ)rTr`)rrrrrZ return_dictr+) r'r1r,r(r)r*r+r-r.r/r0r9)rArrrmrlr>rBrr.r<r/rallrKrCrr^rkr?r'r&r)r*r+r()rNrrXrrYrr)rr+rrrrr,rrnZbackbone_outputsZbackbone_hidden_statesr0Zbackbone_logitsr'r1r,Zdepth_decoder_outputsZbackbone_labelsZ train_maskZdepth_decoder_input_idsZ train_idxsZbackbone_last_hidden_statesZdepth_decoder_labelsr9r9r:rbsV   (  z#CsmForConditionalGeneration.forwardrrr^rc Krrrrr9r9r:r rzQCsmForConditionalGeneration._prepare_4d_causal_attention_mask_with_cache_positionr4) NNNNNNNNNNNNr)$r2r3r4r5rUrrrCrDrE classmethodrMrWrr6rrlrr r7r1rrrrrrrrrr&rbrr^r rgr9r9r[r:r;`s     S       )r;)r;r:rr r;r)r>)Q dataclassesrtypingrrrrrr6Ztorch.nnrCZ activationsr Z cache_utilsr r Z generationr Z integrationsr Zmodeling_attn_mask_utilsrZmodeling_flash_attention_utilsrZmodeling_layersrZmodeling_outputsrrZmodeling_rope_utilsrrZmodeling_utilsrrZprocessing_utilsrutilsrrrrrrautorZconfiguration_csmr!r"Zgeneration_csmr#Z!torch.nn.attention.flex_attentionr$Zintegrations.flex_attentionr%Z get_loggerr2rr&r;ModulerMrhrrrrrrrzrrr=rrJrr r6r:r;__all__r9r9r9r:s               :!"  M5