o ZhJ@s ddlmZddlmZmZmZmZddlZddlm Z ddl m Z m Z ddl mZddlmZddlmZmZdd lmZdd lmZdd lmZmZmZmZd d lmZd dlm Z m!Z!m"Z"m#Z#m$Z$m%Z%m&Z&m'Z'ddl(m)Z)m*Z*ddl+m,Z,e-e.Z/eGdddeZ0eddeGdddeZ1Gddde&Z2Gddde'Z3Gddde$Z4Gddde!Z5Gd d!d!e"Z6eGd"d#d#e%Z7Gd$d%d%e j8Z9ed&dGd'd(d(e#eZ:Gd)d*d*e j8Z;eGd+d,d,e%ZdS)1) dataclass)ListOptionalTupleUnionN)Cache DynamicCache)GenerationMixin)FlashAttentionKwargs)BaseModelOutputWithPastCausalLMOutputWithPast)PreTrainedModel)Unpack) ModelOutputauto_docstringcan_return_tuplelogging) AutoModel)KwargsForCausalLMLlamaAttentionLlamaDecoderLayerLlamaForCausalLMLlamaMLP LlamaModel LlamaRMSNormLlamaRotaryEmbedding) CsmConfigCsmDepthDecoderConfig)CsmGenerationMixinc@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+r,r-r5r5R/var/www/auris/lib/python3.10/site-packages/transformers/models/csm/modular_csm.pyr"/s ,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)Ng)meanstdrg?)configZinitializer_range isinstancennLinearweightdataZnormal_biasZzero_ EmbeddingZ padding_idxCsmCodebooksHead num_codebooksrange CsmRMSNormZfill_)selfmoduler;rEir5r5r6 _init_weightss&      z CsmPreTrainedModel._init_weightsN)r.r/r0r 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_backendrKr5r5r5r6r7js r7c@ eZdZdS)rGNr.r/r0r5r5r5r6rGrGc@rM)CsmRotaryEmbeddingNrNr5r5r5r6rPrOrPc@rM)CsmMLPNrNr5r5r5r6rQrOrQc@rM) CsmAttentionNrNr5r5r5r6rRrOrRc@rM)r9NrNr5r5r5r6r9rOr9cseZdZeZfddZee          ddej de ej de ej de ej de e d e ej d e ed e ed e ed e ej deedeeeffddZZS)CsmDepthDecoderModelcs>t|t|j|j|j|_tj|j|j dd|_ dSNF)rB) super__init__r>rCrE vocab_sizeZbackbone_hidden_size embed_tokensr? hidden_sizeinputs_embeds_projectorrHr< __class__r5r6rVs zCsmDepthDecoderModel.__init__N input_idsbackbone_last_hidden_stateattention_mask position_idsr% inputs_embeds use_cacheoutput_attentionsoutput_hidden_statescache_positionflash_attn_kwargsreturnc  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.zX`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`.Frrdevice)minzvWhen the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference.r5)r`raZpast_key_valuerdrcrfposition_embeddings)Zlast_hidden_stater%r&r')r2compilerZ is_compilingloggerZ warning_oncer<rdrerc ValueErrorZgradient_checkpointingZtrainingr Zget_seq_lengthshaperjarangeclamprWrXwarningrZZ_update_causal_mask unsqueezeZ rotary_embZlayersZnum_hidden_layersZnormr )rHr^r_r`rar%rbrcrdrerfrgZpast_seen_tokensZinputs_seq_lengthrj codebook_idxsoffsetZinput_ids_are_first_codebook causal_maskr&rlZall_hidden_statesZall_self_attnsZ decoder_layerZ layer_outputsr5r5r6forwards               zCsmDepthDecoderModel.forward) NNNNNNNNNN)r.r/r0r rLrVrrr2 LongTensorrr3Tensorrboolrr rrr rx __classcell__r5r5r\r6rSsP      rScs&eZdZfddZdddZZS)rDcs0t||_tt|jd|||_dS)Nr)rUrVrEr> Parameterr2emptyr@)rHrYrErWr\r5r6rV%s  zCsmCodebooksHead.__init__Ncsf|durjd}|jt|n |d}|j|fddtjdDtjddS)Nrc s2g|]}tjdd|ddf|jqSN)r> functionalZlinearT).0Z codebook_idxZcodebook_weightr&r5r6 2s$z,CsmCodebooksHead.forward..rdim)rpr@r2rqrFstack)rHr&rfZ seq_lengthrur5rr6rx*s    zCsmCodebooksHead.forwardrr.r/r0rVrxr|r5r5r\r6rD$s rDa$ 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!s:eZdZdZdZdZfddZddZddZ    dde j d e e d e e j d e e j d e e j f fd d Zee            dde j de e j d e e jde e j d e ee ee j fd e e j de e j de ede ede ed e e j deee jfdeedeeeffddZZS)CsmDepthDecoderForCausalLMNcs2t||`t|j|j|j|_t||_ dSr) rUrVlm_headrDrYrErWcodebooks_headrSr8r[r\r5r6rVGs z#CsmDepthDecoderForCausalLM.__init__cCtdNzNot needed for CsmAttributeErrorrHr5r5r6get_output_embeddingsMz0CsmDepthDecoderForCausalLM.get_output_embeddingscCrrrrHZnew_embeddingsr5r5r6set_output_embeddingsPrz0CsmDepthDecoderForCausalLM.set_output_embeddingsr^r%r`rbrfc  sHtj|||||fi|}|dddk}|s|d|d|S)Nrfrr_ra)rUprepare_inputs_for_generationpop) rHr^r%r`rbrfkwargs model_inputsZis_first_generation_stepr\r5r6rSs   z8CsmDepthDecoderForCausalLM.prepare_inputs_for_generationrr_ralabelsrcrdrelogits_to_keeprrhc  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) r^r_r`rar%rbrcrdrerfrr.)r$rrW shift_labels)r#r$r%r&r'r5)r<rdrer8r=intslicer contiguous loss_functionrWr r%r&r')rHr^r_r`rar%rbrrcrdrerfrrZoutputsr& slice_indicesr$r#rr5r5r6rxisT   & z"CsmDepthDecoderForCausalLM.forwardNNNN) NNNNNNNNNNNr)r.r/r0_tied_weights_keysZ_tp_planZ_pp_planrVrrr2ryrrr3rrrrzrrr{rrrrr rxr|r5r5r\r6r;s       rcs$eZdZfddZddZZS)CsmBackboneModelEmbeddingscsDtt|j|j|j|_|jdt |j|jdddS)Naudio_tokens_offsetsF) persistent) rUrVr>rCrErWrYembed_audio_tokensZregister_bufferr2rqr[r\r5r6rVs  z#CsmBackboneModelEmbeddings.__init__cCs |||j}|jdd}|S)Nrr)rrsum)rHr^Z input_embedsr5r5r6rxs z"CsmBackboneModelEmbeddings.forwardrr5r5r\r6rs rcs0eZdZfddZeefddZZS)CsmBackboneModelcst|t||_dSr)rUrVrrXr[r\r5r6rVs zCsmBackboneModel.__init__c stjdi|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) Nr5)rUrx)rHZ super_kwargsr\r5r6rxszCsmBackboneModel.forward)r.r/r0rVrrrxr|r5r5r\r6rs  rz 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|_|dSrT)rUrVrWr>r?rYrrCZtext_vocab_sizeembed_text_tokensrZ _from_configbackbone_modelrZdepth_decoder_config depth_decoderrZ from_configZ codec_config codec_modelZ post_initr[r\r5r6rVs   z$CsmForConditionalGeneration.__init__cCs|jjSrrrXrr5r5r6get_input_embeddingsrz0CsmForConditionalGeneration.get_input_embeddingscCs ||j_dSrr)rHvaluer5r5r6set_input_embeddingss z0CsmForConditionalGeneration.set_input_embeddingscCs|jSrrrr5r5r6rsz1CsmForConditionalGeneration.get_output_embeddingscCs ||_dSrrrr5r5r6rs z1CsmForConditionalGeneration.set_output_embeddingscCs(|jjr||jjj|jjjdSdSr)r<Ztie_codebooks_embeddingsZ_tie_or_clone_weightsrrXrrr8rr5r5r6 _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_lenr5r6 s  z?CsmForConditionalGeneration.from_pretrained..Z_from_model_config) getrUfrom_pretrainedlenvarsgeneration_configitemsrupdatedelattr)clsargsrr8Z loading_infodepth_decoder_attrsrr\rr6rs   z+CsmForConditionalGeneration.from_pretrainedcsVd}|jj}|dd|D] \}}t|j|||qtj|i|dS)NrZtransformers_version)rrZ to_diff_dictrrsetattrrUsave_pretrained)rHrrrrrrr\r5r6r#s   z+CsmForConditionalGeneration.save_pretrainedNr^ input_valuesinput_values_cutoffsrrhcs||}|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. Nrrrrir.css|]}|jdVqdS)rN)rprelr5r5r6 \szQCsmForConditionalGeneration._merge_input_ids_with_input_values..c s,g|]}tj|ddd|jdfqS)r)r>rpadrprZmax_audio_framesr5r6r^s,zRCsmForConditionalGeneration._merge_input_ids_with_input_values..)rjdtypeiTas_tuple)rbr)#rr>rrdiffr2rqmaxrjexpandrrtziprFrprencodeZ audio_codesZ transposeappendrZget_audio_codes_maskr<audio_token_idrrXZonesrElongZcodebook_eos_token_idZsqueezeZaudio_eos_token_idrepeatrnonzero)rHr^rrrrbZ audio_lengthsZinput_values_maskZaudio_tokens_listZbatch_input_valuesZbatch_input_values_cutoffsrJZ start_idxZend_idxZ audio_batchZ codec_outputsZ codebook_idsZbatched_audio_token_idsZaudio_codes_maskrZaudio_token_maskZ audio_embedsZaudio_eos_frame_idsZaudio_eos_embedsZaudio_eos_token_maskZlabels_expandedZ depth_decoder_ignore_frames_idxsr5rr6"_merge_input_ids_with_input_values-sV            z>CsmForConditionalGeneration._merge_input_ids_with_input_valuesr%r`rbrfc  stjd |||||d|}|dur>|jdkr>|ddur>|j||d|d|dd}||d|ddd|S) N)r^r%r`rbrfrrbrrr)r^rrr)rbrr^r5)rUrndimrrr) rHr^r%r`rbrfrr merged_inputsr\r5r6r}s(  z9CsmForConditionalGeneration.prepare_inputs_for_generationrrarcrdrerrc 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() ```Nrrbr) r^r`rar%rbrcrdrerfr)r$rrWrrrr.r)rTr)r^r_rcrdreZ return_dictr) r#r-r(r$r%r&r'r)r*r+r,r5)r<rdrerrrr=rrrrrWallrEr>rrrrr#r"r%r&r'r$)rHr^rr`rrar%rbrrcrdrerfrrrZbackbone_outputsZbackbone_hidden_statesrZbackbone_logitsr#r-r(Zdepth_decoder_outputsZbackbone_labelsZ train_maskZdepth_decoder_input_idsZ train_idxsZbackbone_last_hidden_statesZdepth_decoder_labelsr5r5r6rxsV   (  z#CsmForConditionalGeneration.forwardsequence_length target_lengthr batch_sizec KsD|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_valuerrjr)Zdiagonalrirr)rr2ZfinforkfullrjZtriurqZreshaperclonerptoZ masked_fill) r`rrrrfrrrwZ min_dtypeZ mask_lengthZ padding_maskr5r5r65_prepare_4d_causal_attention_mask_with_cache_positionFs,  $ 6  zQCsmForConditionalGeneration._prepare_4d_causal_attention_mask_with_cache_positionr) NNNNNNNNNNNNr)$r.r/r0rrVrrrrr classmethodrrrr2rzrryrr3rrrrrr{rrrrr"rx staticmethodrrr|r5r5r\r6rs     S       )r)r7rrSrr)? dataclassesrtypingrrrrr2Ztorch.nnr>Z cache_utilsrr Z generationr Zmodeling_flash_attention_utilsr Zmodeling_outputsr r Zmodeling_utilsrZprocessing_utilsrutilsrrrrautorZllama.modeling_llamarrrrrrrrZconfiguration_csmrr Zgeneration_csmr!Z get_loggerr.rnr"r7rGrPrQrRr9rSModulerDrrrr__all__r5r5r5r6s\       (  :"}v