o ei@sddlmZddlmZddlmZddlZddlmZddl m Z ddl m Z ddlmZmZdd lmZdd lmZmZmZdd lmZdd lmZdd lmZmZddlmZm Z ddl!m"Z"m#Z#ddl$m%Z%ddl&m'Z'm(Z(m)Z)m*Z*m+Z+ddl,m-Z-m.Z.ddl/m0Z0ddl1m2Z2ddl3m4Z4ddl5m6Z6m7Z7ddl8m9Z9e+:e;ZZ?Gd!d"d"ej>Z@Gd#d$d$ej>ZAd%d&ZBed'dPd(d)ZCd*ejDd+eEd,ejDfd-d.ZF /dQd0ej>d1ejDd2ejDd3ejDd4ejDdBd5eGd6eGd7e%e(fd8d9ZHeeCGd:d;d;ej>ZIGdde)Gd?d@d@e#ZKe)GdAdBdBeKZLGdCdDdDej>ZMe)dEdGdFdGdGeKeZNGdHdIdIej>ZOe)GdJdKdKeKZPe)dLdGdMdNdNeKe9ZQgdOZRdS)R)Callable) dataclass)OptionalN)initialization)ACT2FN)Cache DynamicCache)GenerationMixin)use_kernel_forward_from_hubuse_kernel_func_from_hubuse_kernelized_func)create_causal_mask)GradientCheckpointingLayer)BaseModelOutputWithPastCausalLMOutputWithPast)ROPE_INIT_FUNCTIONSdynamic_rope_update)ALL_ATTENTION_FUNCTIONSPreTrainedModel)Unpack) ModelOutputTransformersKwargsauto_docstringcan_return_tuplelogging)maybe_autocastmerge_with_config_defaults)is_torchdynamo_compiling)capture_outputs) AutoModel) CsmConfigCsmDepthDecoderConfig)CsmGenerationMixinz: Base class for the model autoregressive outputs. ) custom_introc@seZdZUdZdZejdBed<dZejdBed<dZ e dBed<dZ e ejdfdBed<dZ e ejdfdBed<dZejdBed <dZejdBed <dZe dBed <dZe ejdfdBed <dZe ejdfdBed <dZejdBed<dS)CsmOutputWithPasta 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). 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. 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 (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`): It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache). 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(torch FloatTensor__annotations__r)r*rr+tupler,r-r.r/r0r1r2r;r;b/home/ubuntu/transcripts/venv/lib/python3.10/site-packages/transformers/models/csm/modeling_csm.pyr'3s r'RMSNormcsFeZdZd deddffdd Zdejdejfdd Zd d ZZ S) CsmRMSNormư>epsreturnNcs&ttt||_||_dS)z9 CsmRMSNorm is equivalent to T5LayerNorm N)super__init__nn Parameterr7onesweightvariance_epsilon)self hidden_sizer@ __class__r;r<rCes  zCsmRMSNorm.__init__r+cCsJ|j}|tj}|djddd}|t||j}|j||S)Nr T)keepdim) dtypetor7float32powmeanrsqrtrHrG)rIr+ input_dtypevariancer;r;r<forwardms  zCsmRMSNorm.forwardcCst|jjd|jS)Nz, eps=)r:rGshaperHrIr;r;r< extra_reprtszCsmRMSNorm.extra_repr)r?) r3r4r5floatrCr7TensorrWrZ __classcell__r;r;rKr<r>csr>c s~eZdZUejed<ddeffdd Ze   ddedBde dde dBd e d e ffd d Z eed dZZS)CsmRotaryEmbeddinginv_freqNconfigcst|j|_|j|_||_|jjd|_|j}|jdkr$t |j}||j|\}|_ |j d|dd|j d| dddS)N rope_typedefaultr_F persistentoriginal_inv_freq) rBrCmax_position_embeddingsmax_seq_len_cachedoriginal_max_seq_lenr`rope_parametersracompute_default_rope_parametersrattention_scalingregister_bufferclone)rIr`device rope_init_fnr_rKr;r<rC{s   zCsmRotaryEmbedding.__init__rnz torch.deviceseq_lenrAz torch.TensorcCsZ|jd}t|ddp|j|j}d}d|tjd|dtjdj|tjd|}||fS) a Computes the inverse frequencies according to the original RoPE implementation Args: config ([`~transformers.PreTrainedConfig`]): The model configuration. device (`torch.device`): The device to use for initialization of the inverse frequencies. seq_len (`int`, *optional*): The current sequence length. Unused for this type of RoPE. Returns: Tuple of (`torch.Tensor`, `float`), containing the inverse frequencies for the RoPE embeddings and the post-processing scaling factor applied to the computed cos/sin (unused in this type of RoPE). rope_thetahead_dimNg?rr rOrnrO) rigetattrrJnum_attention_headsr7arangeint64rPr[)r`rnrpbasedimattention_factorr_r;r;r<rjs &z2CsmRotaryEmbedding.compute_default_rope_parametersc Cs|jddddf|jddd|j}|dddddf}t|jjtr6|jjdkr6|jjnd}t |dd+|| dd}t j ||fdd }| |j}||j} Wdn1slwY|j|jd | j|jd fS) NrrMr"mpscpuF) device_typeenabledr rzrs)r_r[expandrXrPrn isinstancetypestrr transposer7catcosrksinrO) rIx position_idsinv_freq_expandedposition_ids_expandedr~freqsembrrr;r;r<rWs0&zCsmRotaryEmbedding.forwardN)NNN)r3r4r5r7r\r9r#rC staticmethodrintr:r[rjno_gradrrWr]r;r;rKr<r^xs&   r^c$eZdZfddZddZZS)CsmMLPcsxt||_|j|_|j|_tj|j|j|jd|_tj|j|j|jd|_ tj|j|j|jd|_ t |j |_ dS)Nbias)rBrCr`rJintermediate_sizerDLinearmlp_bias gate_projup_proj down_projr hidden_actact_fnrIr`rKr;r<rCs zCsmMLP.__init__cCs$||||||}|Sr)rrrr)rIrrr;r;r<rWs zCsmMLP.forwardr3r4r5rCrWr]r;r;rKr<rs  rcCsH|dd|jddf}|d|jdddf}tj| |fddS)z*Rotates half the hidden dims of the input..NrMr r)rXr7r)rx1x2r;r;r< rotate_halfsrrotary_pos_embcCsD||}||}||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. 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)qkrr unsqueeze_dimq_embedk_embedr;r;r<apply_rotary_pos_embs  rr+n_reprAcCs^|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)rXrreshape)r+rbatchnum_key_value_headsslenrrr;r;r< repeat_kvs 0rmodulequerykeyvalueattention_maskscalingdropoutkwargsc Kst||j}t||j} t||dd|} |dur | |} tjj| dtjd |j } tjj | ||j d} t| | } | dd } | | fS)Nr rrM)rzrO)ptrainingr")rnum_key_value_groupsr7matmulrrD functionalsoftmaxrQrPrOrr contiguous) rrrrrrrr key_states value_states attn_weights attn_outputr;r;r<eager_attention_forwards  rcseZdZdZdedeffdd Z    ddejde ejejfdBd ejdBd e dBd ej dBd e e d e ejejffddZZS) CsmAttentionz=Multi-headed attention from 'Attention Is All You Need' paperr` 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)NrrgTr)rBrCr`rrurJrvrrrrrattention_dropout is_causalrDrattention_biasq_projk_projv_projo_projrIr`rrKr;r<rCs(   zCsmAttention.__init__Nr+position_embeddingsrr*cache_positionrrAcKs|jdd}g|d|jR}|||dd} |||dd} |||dd} |\} } t| | | | \} } |durW| | |d}|| | |j |\} } t |j j t}||| | | |f|jskdn|j|jd|\}}|jg|dR}||}||fS)NrMr"r )rrrr)rr)rXrrrviewrrrrupdaterr get_interfacer`_attn_implementationrrrrrrr)rIr+rrr*rr input_shape hidden_shape query_statesrrrr cache_kwargsattention_interfacerrr;r;r<rW*s8    zCsmAttention.forwardNNNN)r3r4r5r6r#rrCr7r\r:r LongTensorrrrWr]r;r;rKr<rs,rcseZdZdedeffdd Z      ddejdejdBd ejdBd e dBd e dBd ejdBd e ejejfdBde e dejfddZZS)CsmDecoderLayerr`rcsRt|j|_t||d|_t||_t|j|jd|_ t|j|jd|_ dS)N)r`rr@) rBrCrJr self_attnrmlpr> rms_norm_epsinput_layernormpost_attention_layernormrrKr;r<rCWs  zCsmDecoderLayer.__init__NFr+rrr* use_cacherrrrAc Ks^|} ||}|jd|||||||d|\}} | |}|} ||}||}| |}|S)N)r+rrr*rrrr;)rrrr) rIr+rrr*rrrrresidual_r;r;r<rWas&     zCsmDecoderLayer.forward)NNNFNN)r3r4r5r#rrCr7r\rrboolr:rrrWr]r;r;rKr<rVs6   rz[ The bare Csm Model outputting raw hidden-states without any specific head on top. cs`eZdZUeed<dZdZdZdgZdgZ dZ dZ dZ dZ eedZefdd ZZS) CsmPreTrainedModelr`model)audiotextTrr*)r+r,cszt|t|tr$|j}t|dD] }tj|jd|j j dqdSt|t r;t |j t|j j|j jdSdS)Nr"r)rSstd)rB _init_weightsrCsmCodebooksHead num_codebooksrangeinitnormal_rGr`initializer_rangeCsmBackboneModelEmbeddingscopy_audio_tokens_offsetsr7rw vocab_size)rIrrirKr;r<rs   $z CsmPreTrainedModel._init_weights)r3r4r5r#r9base_model_prefixinput_modalitiessupports_gradient_checkpointing_no_split_modules_skip_keys_device_placement_supports_flash_attn_supports_sdpa_can_compile_fullgraph_supports_attention_backendrr_can_record_outputsr7rrr]r;r;rKr<rs rcseZdZUeed<fddZeee        dde j dBde j dBde j dBde j dBd e dBd e j dBd edBd e j dBd eedeeBfddZZS)CsmDepthDecoderModelr`cstj|_j|_tjjj|_ t fddt j D|_ tjjd|_td|_d|_tjjjdd|_|dS)Ncg|]}t|qSr;r.0rr`r;r< z1CsmDepthDecoderModel.__init__..rr Fr)rBrC pad_token_id padding_idxrrD Embeddingrbackbone_hidden_size embed_tokens ModuleListrnum_hidden_layerslayersr>rJrnormr^ rotary_embgradient_checkpointingrinputs_embeds_projector post_initrrKr r<rCs   zCsmDepthDecoderModel.__init__N input_idsbackbone_last_hidden_staterrr* inputs_embedsrrrrAc  Ks|durtstdd}|du|duArtd|r&|dur&t|jd}|durV|dur2|nd} |dur=|jdn|jd} |durI|jn|j} t j | | | | d}|durt j |ddd} | |j }| ||}|ddk}|dur||dddf<n ts|rtd ||}t|j|||||d }|}|d}|j||d }|jd|jjD]}||f||||||d | }q||}t||r|d Sdd 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.r rr"rn)minzvWhen the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference.r`rrrr*rr)rrr*rrrlast_hidden_stater*)rlogger warning_once ValueErrorr r`get_seq_lengthrXrnr7rwclamprrwarningrrrrrrrr)rIrrrrr*rrrrpast_seen_tokensinputs_seq_lengthrn codebook_idxsoffsetinput_ids_are_first_codebook causal_maskr+r decoder_layerr;r;r<rWsr        zCsmDepthDecoderModel.forward)NNNNNNNN)r3r4r5r$r9rCrrrr7rr8r\rrrrr:rrWr]r;r;rKr<rsF     rcs&eZdZfddZdddZZS)rcs0t||_tt|jd|||_dS)Nr")rBrCrrDrEr7emptyrG)rIrJrrrKr;r<rCs  zCsmCodebooksHead.__init__Ncsf|durjd}|jt|n |d}|j|fddtjdDtjddS)Nr"c s2g|]}tjdd|ddf|jqSr)rDrlinearT)r codebook_idxcodebook_weightr+r;r<r !s$z,CsmCodebooksHead.forward..rr)rXrGr7rwrstack)rIr+r seq_lengthr*r;r3r<rWs    zCsmCodebooksHead.forwardrrr;r;rKr<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.) cseZdZdZdZdZfddZee          dde j dBde j dBde j dBde j dBd e dBd e j dBd e j dBd edBd e j dBdee j BdeedeeBfddZ    dde j d e dBde j dBd e j dBd e j dBf fdd ZZS)CsmDepthDecoderForCausalLMNcs>t|t||_|j|_t|j|j|j|_| dSr) rBrCrrrrrJrcodebooks_headrrrKr;r<rC6s   z#CsmDepthDecoderForCausalLM.__init__rrrrrr*rlabelsrrlogits_to_keeprrAc  Ks|jd|||||||| d| } | d} t| tr+| dkr$tdd}n t| d}n| }|| dd|ddf| durA| |nd}|}d}|durg|dddf}|jd|d|jj|d| }t ||| j | j | j dS) 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]`. )rrrrr*rrrrr"N.)r)r9r shift_labels)r(r)r*r+r,r;) rrrslicer8r loss_functionr`rrr*r+r,)rIrrrrr*rr9rrr:routputsr+ slice_indicesr)r(r;r;r;r<rW?sJ   & z"CsmDepthDecoderForCausalLM.forwardc  sHtj|||||fi|}|dddk}|s|d|d|S)Nrrrr)rBprepare_inputs_for_generationpop) rIrr*rrrr model_inputsis_first_generation_steprKr;r<r@s   z8CsmDepthDecoderForCausalLM.prepare_inputs_for_generation) NNNNNNNNNrr)r3r4r5_tied_weights_keys_tp_plan_pp_planrCrrr7rr8r\rrrrrr:rrWr@r]r;r;rKr<r7*sr       Er7cr)rcsDtt|j|j|j|_|jdt |j|jdddS)NrFrc) rBrCrDrrrrJembed_audio_tokensrlr7rwrrKr;r<rCs  z#CsmBackboneModelEmbeddings.__init__cCs |||j}|jdd}|S)Nr r)rGrsum)rIrrr;r;r<rWs z"CsmBackboneModelEmbeddings.forwardrr;r;rKr<rs rcseZdZfddZeee       ddejdBdej dBdejdBde dBdej dBd ejdBd e dBd e ed efd dZZS)CsmBackboneModelcsvtj|_j|_t|_tfddt j D|_ t j jd|_td|_d|_|dS)Ncrr;rrr r;r<r r z-CsmBackboneModel.__init__..rr F)rBrCr r rrrrDrrrrr>rJrrr^rrrrrKr r<rCs    zCsmBackboneModel.__init__Nrrrr*rrrrrAc Ks|du|duAr td|dur||}|r!|dur!t|jd}|dur<|dur-|nd} tj|jd|jd| }|durE| d}t |j|||||d} |} |j | |d} |j d|jj D]} | | f| | ||||d |} qb|| } t| |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_embedsr rr"rrr)rrrr*rrr )r$rr r`r%r7rwrXrnrrrrrrr)rIrrrr*rrrrr(r-r+rr.r;r;r<rWsP      zCsmBackboneModel.forward)NNNNNNN)r3r4r5rCrrrr7rr\rr8rrrrrWr]r;r;rKr<rIs>   rIz 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. cs~eZdZddiZfddZddZddZefd d Zfd d Z d"de j d Bde j d Bde j d Bde j d Bde j d Bf ddZ d"de j ded Bde j d Bde jd Bde j d Bf fdd Zee d#de j d Bde j d Bde j d Bde j d Bde j d Bded Bde jd Bde j d Bded Bde j d Bdee j BdeedeeBfd d!ZZS)$CsmForConditionalGenerationz5backbone_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)rBrCrrDrrJlm_headrtext_vocab_sizeembed_text_tokensrI _from_configbackbone_modelr7depth_decoder_config depth_decoderr! from_config codec_config codec_modelrrrKr;r<rCs   z$CsmForConditionalGeneration.__init__cCs|jjSrrOrrYr;r;r<get_input_embeddingssz0CsmForConditionalGeneration.get_input_embeddingscCs ||j_dSrrU)rIrr;r;r<set_input_embeddingss z0CsmForConditionalGeneration.set_input_embeddingscs|ddrtj|i|\}}n tj|i|}dtfddt|jD}t|jjddi||D] }t |j|q?d|vrR||fS|S)Noutput_loading_infoFdepth_decoder_cs(i|]\}}|r|d|qSr) startswith)rattrrprefix prefix_lenr;r< (s  z?CsmForConditionalGeneration.from_pretrained.._from_model_config) getrBfrom_pretrainedlenvarsgeneration_configitemsrQrdelattr)clsargsrr loading_infodepth_decoder_attrsr[rKr\r<rbs   z+CsmForConditionalGeneration.from_pretrainedcsVd}|jj}|dd|D] \}}t|j|||qtj|i|dS)NrYtransformers_version)rQre to_diff_dictrArfsetattrrBsave_pretrained)rIrirr]rkr[rrKr;r<ro9s   z+CsmForConditionalGeneration.save_pretrainedNr input_valuesinput_values_cutoffsr9rAcsF||}|durtj|d}||dk}||dk}tj||jd t |d}|| dk}t jg}t ||D]?\} } | | dk} t| jddD]+} | | } | | d} | d| | f}|j| d}|jdd}||dqUqBtdd |Dtfd d |D}|j|}Wdn1swY|jj}||k}|j|}||||<tjdd|jjf|jtjd |jj}|j|d}||jj k}|!|"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 retrieve 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"rrrrMr".css|]}|jdVqdS)rN)rXrelr;r;r< sszQCsmForConditionalGeneration._merge_input_ids_with_input_values..c s,g|]}tj|ddd|jdfqS)r)rDrpadrXrsmax_audio_framesr;r<r us,zRCsmForConditionalGeneration._merge_input_ids_with_input_values..rtiTas_tuple)rr9)$rMrDrrvdiffr7rwmaxrnrrcrrziprrXrTencode audio_codesrappendr5get_audio_codes_maskr`audio_token_idrOrrFrlongcodebook_eos_token_idsqueezeaudio_eos_token_idrepeatrHnonzero)rIrrprqr9r audio_lengthsinput_values_maskaudio_tokens_listbatch_input_valuesbatch_input_values_cutoffsr start_idxend_idx audio_batch codec_outputs codebook_idsbatched_audio_token_idsaudio_codes_maskraudio_token_mask audio_embedsaudio_eos_frame_idsaudio_eos_embedsaudio_eos_token_masklabels_expanded depth_decoder_ignore_frames_idxsr;rwr<"_merge_input_ids_with_input_valuesCs\            z>CsmForConditionalGeneration._merge_input_ids_with_input_valuesr*rrrc  stjd |||||d|}|dur>|jdkr>|ddur>|j||d|d|dd}||d|ddd|S) N)rr*rrrr rrprqr9)rrprqr9)rr9rr;)rBr@ndimrarr) rIrr*rrrrrB merged_inputsrKr;r<r@s(  z9CsmForConditionalGeneration.prepare_inputs_for_generationrrrr:rc  Ks|dur|jdkr|||||} | d}| d}d}|jd|||||| | d| }|d}t| tr:t| dn| }||dd|ddf}d}d}d}d}|dur|dddddf}|jd|||jj d| }|ddddddfd kj d d }||d d|jj df}t j j|d dd}|jdd}||d|ddddf}||}|jd||| 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 inferred 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 = "sesame/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() ```Nr rr9)rrrr*rrrr)r)r9rr"r{rMr.rr)rTry)rrr return_dictr9) r(r2r-r)r*r+r,r.r/r0r1r;)rrrOrrr<rKr=r`rallrrDrrvrrQr(r'r*r+r,r))rIrrprrqrr*rr9rrr:rrbackbone_outputsbackbone_hidden_statesr?backbone_logitsr(r2r-depth_decoder_outputsbackbone_labels train_maskdepth_decoder_input_ids train_idxsbackbone_last_hidden_statesdepth_decoder_labelsr;r;r<rWsS  (   z#CsmForConditionalGeneration.forwardr) NNNNNNNNNNr)r3r4r5rDrCrVrW classmethodrbror7r\rrrr8r@rrrrrrr:r'rWr]r;r;rKr<rJs     U     rJ)rrIrr7rJ)r")r)Scollections.abcr dataclassesrtypingrr7torch.nnrDrr activationsr cache_utilsrr generationr integrationsr r r masking_utilsrmodeling_layersrmodeling_outputsrrmodeling_rope_utilsrrmodeling_utilsrrprocessing_utilsrutilsrrrrr utils.genericrrutils.import_utilsrutils.output_capturingrautor!configuration_csmr#r$generation_csmr% get_loggerr3r"r'Moduler>r^rrrr\rrr[rrrrrrr7rrIrJ__all__r;r;r;r<s               *A  F-ki[ M