# coding=utf-8 # Copyright 2025 Sesame and The HuggingFace Inc. team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. from dataclasses import dataclass from typing import List, Optional, Tuple, Union import torch import torch.nn as nn from ...cache_utils import Cache, DynamicCache from ...generation import GenerationMixin from ...modeling_flash_attention_utils import FlashAttentionKwargs from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast from ...modeling_utils import PreTrainedModel from ...processing_utils import Unpack from ...utils import ModelOutput, auto_docstring, can_return_tuple, logging from ..auto import AutoModel from ..llama.modeling_llama import ( KwargsForCausalLM, LlamaAttention, LlamaDecoderLayer, LlamaForCausalLM, LlamaMLP, LlamaModel, LlamaRMSNorm, LlamaRotaryEmbedding, ) from .configuration_csm import CsmConfig, CsmDepthDecoderConfig from .generation_csm import CsmGenerationMixin logger = logging.get_logger(__name__) @dataclass class CsmOutputWithPast(ModelOutput): """ 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. """ loss: Optional[torch.FloatTensor] = None logits: torch.FloatTensor = None past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None attentions: Optional[Tuple[torch.FloatTensor, ...]] = None depth_decoder_loss: Optional[torch.FloatTensor] = None depth_decoder_logits: torch.FloatTensor = None depth_decoder_past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None depth_decoder_hidden_states: Optional[Tuple[torch.FloatTensor, ...]] = None depth_decoder_attentions: Optional[Tuple[torch.FloatTensor, ...]] = None backbone_loss: Optional[torch.FloatTensor] = None @auto_docstring( custom_intro=""" The bare Csm Model outputting raw hidden-states without any specific head on top. """ ) @auto_docstring class CsmPreTrainedModel(PreTrainedModel): config_class = CsmConfig base_model_prefix = "model" supports_gradient_checkpointing = True _no_split_modules = ["CsmDecoderLayer"] _skip_keys_device_placement = ["past_key_values"] _supports_flash_attn_2 = True _supports_sdpa = True # does not because of Mimi codec model # _supports_flex_attn = True _supports_cache_class = True _supports_quantized_cache = True _supports_static_cache = True _supports_attention_backend = True def _init_weights(self, module): std = self.config.initializer_range if isinstance(module, nn.Linear): module.weight.data.normal_(mean=0.0, std=std) if module.bias is not None: module.bias.data.zero_() elif isinstance(module, nn.Embedding): module.weight.data.normal_(mean=0.0, std=std) if module.padding_idx is not None: module.weight.data[module.padding_idx].zero_() elif isinstance(module, CsmCodebooksHead): num_codebooks = module.num_codebooks for i in range(num_codebooks - 1): module.weight.data[i].normal_(mean=0.0, std=std) elif isinstance(module, CsmRMSNorm): module.weight.data.fill_(1.0) # manually specify names for correct naming when converting from modualr class CsmRMSNorm(LlamaRMSNorm): pass class CsmRotaryEmbedding(LlamaRotaryEmbedding): pass class CsmMLP(LlamaMLP): pass class CsmAttention(LlamaAttention): pass class CsmDecoderLayer(LlamaDecoderLayer): pass @auto_docstring class CsmDepthDecoderModel(LlamaModel): config_class = CsmDepthDecoderConfig def __init__(self, config): super().__init__(config) self.embed_tokens = nn.Embedding((config.num_codebooks * config.vocab_size), config.backbone_hidden_size) self.inputs_embeds_projector = nn.Linear(config.backbone_hidden_size, config.hidden_size, bias=False) @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor = None, backbone_last_hidden_state: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Cache] = None, inputs_embeds: Optional[torch.FloatTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, **flash_attn_kwargs: Unpack[FlashAttentionKwargs], ) -> Union[Tuple, BaseModelOutputWithPast]: r""" 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. """ if position_ids is not None and not torch.compiler.is_compiling(): logger.warning_once( "Custom `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." ) position_ids = None output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) use_cache = use_cache if use_cache is not None else self.config.use_cache if (input_ids is None) ^ (inputs_embeds is not None): raise ValueError("You must specify exactly one of input_ids or inputs_embeds.") if self.gradient_checkpointing and self.training and use_cache: logger.warning_once( "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`." ) use_cache = False if use_cache and past_key_values is None: past_key_values = DynamicCache() if cache_position is None: past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 inputs_seq_length = inputs_embeds.shape[1] if inputs_embeds is not None else input_ids.shape[1] device = inputs_embeds.device if inputs_embeds is not None else input_ids.device cache_position = torch.arange(past_seen_tokens, past_seen_tokens + inputs_seq_length, device=device) if inputs_embeds is None: codebook_idxs = torch.clamp(cache_position - 1, min=0) offset = codebook_idxs * self.vocab_size inputs_embeds = self.embed_tokens(input_ids + offset) input_ids_are_first_codebook = cache_position[0] == 0 if backbone_last_hidden_state is not None: inputs_embeds[:, 0] = backbone_last_hidden_state else: if not torch.compiler.is_compiling() and input_ids_are_first_codebook: logger.warning( "When the first codebook token is provided, `backbone_last_hidden_state` should also be provided for correct inference." ) inputs_embeds = self.inputs_embeds_projector(inputs_embeds) causal_mask = self._update_causal_mask( attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions ) hidden_states = inputs_embeds # create position embeddings to be shared across the decoder layers position_ids = cache_position.unsqueeze(0) position_embeddings = self.rotary_emb(hidden_states, position_ids) # decoder layers all_hidden_states = () if output_hidden_states else None all_self_attns = () if output_attentions else None for decoder_layer in self.layers[: self.config.num_hidden_layers]: if output_hidden_states: all_hidden_states += (hidden_states,) layer_outputs = decoder_layer( hidden_states, attention_mask=causal_mask, position_ids=position_ids, past_key_value=past_key_values, output_attentions=output_attentions, use_cache=use_cache, cache_position=cache_position, position_embeddings=position_embeddings, **flash_attn_kwargs, ) hidden_states = layer_outputs[0] if output_attentions: all_self_attns += (layer_outputs[1],) hidden_states = self.norm(hidden_states) # add hidden states from the last decoder layer if output_hidden_states: all_hidden_states += (hidden_states,) return BaseModelOutputWithPast( last_hidden_state=hidden_states, past_key_values=past_key_values if use_cache else None, hidden_states=all_hidden_states, attentions=all_self_attns, ) class CsmCodebooksHead(nn.Module): def __init__(self, hidden_size, num_codebooks, vocab_size): super().__init__() self.num_codebooks = num_codebooks self.weight = nn.Parameter(torch.empty(self.num_codebooks - 1, hidden_size, vocab_size)) def forward(self, hidden_states, cache_position=None): if cache_position is None: seq_length = hidden_states.shape[1] codebook_weight = self.weight[torch.arange(seq_length)] else: codebook_idxs = cache_position - 1 codebook_weight = self.weight[codebook_idxs] hidden_states = [ nn.functional.linear(hidden_states[:, codebook_idx, :], codebook_weight[codebook_idx].T) for codebook_idx in range(codebook_weight.shape[0]) ] hidden_states = torch.stack(hidden_states, dim=1) return hidden_states @auto_docstring( custom_intro=""" 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.) """ ) class CsmDepthDecoderForCausalLM(LlamaForCausalLM, GenerationMixin): _tied_weights_keys = None _tp_plan = None _pp_plan = None def __init__(self, config): super().__init__(config) del self.lm_head self.codebooks_head = CsmCodebooksHead(config.hidden_size, config.num_codebooks, config.vocab_size) self.model = CsmDepthDecoderModel(config) def get_output_embeddings(self): raise AttributeError("Not needed for Csm") def set_output_embeddings(self, new_embeddings): raise AttributeError("Not needed for Csm") def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[Cache] = None, attention_mask: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ): model_inputs = super().prepare_inputs_for_generation( input_ids, past_key_values, attention_mask, inputs_embeds, cache_position, **kwargs ) is_first_generation_step = model_inputs["cache_position"][0] == 0 if not is_first_generation_step: model_inputs.pop("backbone_last_hidden_state") # csm depth decoder does not use position_ids model_inputs.pop("position_ids") return model_inputs @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor = None, backbone_last_hidden_state: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CausalLMOutputWithPast]: r""" 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]`. """ output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn) outputs = self.model( input_ids=input_ids, backbone_last_hidden_state=backbone_last_hidden_state, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, **kwargs, ) hidden_states = outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss if isinstance(logits_to_keep, int): if logits_to_keep == 0: # skip idx 0 logits since it's for the concatenated backbone last hidden state slice_indices = slice(1, None) else: slice_indices = slice(-logits_to_keep, None) else: slice_indices = logits_to_keep logits = self.codebooks_head( hidden_states[:, slice_indices, :], cache_position[slice_indices] if cache_position is not None else None ) logits = logits.contiguous() loss = None if labels is not None: shift_labels = labels[..., 1:].contiguous() loss = self.loss_function( logits=logits, labels=None, vocab_size=self.config.vocab_size, shift_labels=shift_labels, **kwargs ) return CausalLMOutputWithPast( loss=loss, logits=logits, past_key_values=outputs.past_key_values, hidden_states=outputs.hidden_states, attentions=outputs.attentions, ) class CsmBackboneModelEmbeddings(nn.Module): def __init__(self, config): super().__init__() self.embed_audio_tokens = nn.Embedding((config.num_codebooks * config.vocab_size), config.hidden_size) self.register_buffer( "audio_tokens_offsets", torch.arange(config.num_codebooks) * config.vocab_size, persistent=False ) def forward(self, input_ids): input_embeds = self.embed_audio_tokens(input_ids + self.audio_tokens_offsets) input_embeds = input_embeds.sum(dim=2) return input_embeds @auto_docstring class CsmBackboneModel(LlamaModel): def __init__(self, config): super().__init__(config) self.embed_tokens = CsmBackboneModelEmbeddings(config) @can_return_tuple @auto_docstring def forward(self, **super_kwargs): r""" 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) """ return super().forward(**super_kwargs) @auto_docstring( custom_intro=""" 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. """ ) class CsmForConditionalGeneration(CsmPreTrainedModel, CsmGenerationMixin): _tied_weights_keys = [ "backbone_model.embed_tokens.embed_audio_tokens.weight", "depth_decoder.model.embed_tokens.weight", ] def __init__(self, config): super().__init__(config) self.vocab_size = config.vocab_size self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False) self.embed_text_tokens = nn.Embedding(config.text_vocab_size, config.hidden_size) self.backbone_model = CsmBackboneModel._from_config(config) self.depth_decoder = CsmDepthDecoderForCausalLM._from_config(config.depth_decoder_config) self.codec_model = AutoModel.from_config(config.codec_config) # Initialize weights and apply final processing self.post_init() def get_input_embeddings(self): return self.backbone_model.embed_tokens def set_input_embeddings(self, value): self.backbone_model.embed_tokens = value def get_output_embeddings(self): return self.lm_head def set_output_embeddings(self, new_embeddings): self.lm_head = new_embeddings def _tie_weights(self): if self.config.tie_codebooks_embeddings: self._tie_or_clone_weights( self.backbone_model.embed_tokens.embed_audio_tokens, self.depth_decoder.model.embed_tokens, ) @classmethod def from_pretrained(cls, *args, **kwargs): if kwargs.get("output_loading_info", False): model, loading_info = super().from_pretrained(*args, **kwargs) else: model = super().from_pretrained(*args, **kwargs) # copy depth decoder generation conf attr to the depth decoder generation config prefix = "depth_decoder_" prefix_len = len(prefix) depth_decoder_attrs = { attr[prefix_len:]: value for attr, value in vars(model.generation_config).items() if attr.startswith(prefix) } vars(model.depth_decoder.generation_config).update({"_from_model_config": False, **depth_decoder_attrs}) # remove the depth decoder generation conf attr from the model generation config for attr in depth_decoder_attrs: delattr(model.generation_config, prefix + attr) if "output_loading_info" in kwargs: return model, loading_info else: return model def save_pretrained(self, *args, **kwargs): # copy the depth decoder generation config attributes to the model generation config prefix = "depth_decoder_" depth_decoder_attrs = self.depth_decoder.generation_config.to_diff_dict() depth_decoder_attrs.pop("transformers_version", None) for attr, value in depth_decoder_attrs.items(): setattr(self.generation_config, prefix + attr, value) super().save_pretrained(*args, **kwargs) def _merge_input_ids_with_input_values( self, input_ids: Optional[torch.Tensor] = None, input_values: Optional[torch.Tensor] = None, input_values_cutoffs: Optional[torch.Tensor] = None, labels: Optional[torch.Tensor] = None, ) -> Optional[torch.Tensor]: """ 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. """ inputs_embeds = self.embed_text_tokens(input_ids) if input_values is not None: # infer input_values_mask input_values_cutoffs = nn.functional.pad(input_values_cutoffs, (1, 0)) audio_lengths = input_values_cutoffs[input_values_cutoffs >= 0].diff() audio_lengths = audio_lengths[audio_lengths > 0] input_values_mask = torch.arange(input_values_cutoffs.max(), device=input_values.device).expand( len(audio_lengths), -1 ) input_values_mask = input_values_mask < audio_lengths.unsqueeze(1) # ======================================= # TODO: @eustlb, this should be batched !!! # but requires making sure batched inference of the codec model works as intended audio_tokens_list = [] for batch_input_values, batch_input_values_cutoffs in zip(input_values, input_values_cutoffs): batch_input_values_cutoffs = batch_input_values_cutoffs[batch_input_values_cutoffs >= 0] for i in range(batch_input_values_cutoffs.shape[0] - 1): start_idx = batch_input_values_cutoffs[i] end_idx = batch_input_values_cutoffs[i + 1] audio_batch = batch_input_values[..., start_idx:end_idx] codec_outputs = self.codec_model.encode(audio_batch.unsqueeze(0)) codebook_ids = codec_outputs.audio_codes.transpose(1, -1) audio_tokens_list.append(codebook_ids[0]) max_audio_frames = max(el.shape[0] for el in audio_tokens_list) batched_audio_token_ids = torch.stack( [nn.functional.pad(el, (0, 0, 0, max_audio_frames - el.shape[0])) for el in audio_tokens_list] ) audio_codes_mask = self.codec_model.get_audio_codes_mask(input_values_mask) # ======================================= audio_token_id = self.config.audio_token_id audio_token_mask = input_ids == audio_token_id audio_embeds = self.backbone_model.embed_tokens(batched_audio_token_ids) inputs_embeds[audio_token_mask] = audio_embeds[audio_codes_mask] # same for the audio eos token audio_eos_frame_ids = ( torch.ones((1, 1, self.config.num_codebooks), device=input_ids.device, dtype=torch.long) * self.config.codebook_eos_token_id ) audio_eos_embeds = self.backbone_model.embed_tokens(audio_eos_frame_ids).squeeze(1) audio_eos_token_mask = input_ids == self.config.audio_eos_token_id inputs_embeds[audio_eos_token_mask] = audio_eos_embeds.repeat(audio_eos_token_mask.sum(), 1) # if the labels are provided, we need to expand the labels to (batch_size, seq_length, num_codebooks) if labels is not None: labels_expanded = labels.unsqueeze(-1).repeat(1, 1, self.config.num_codebooks) labels_expanded[audio_token_mask] = batched_audio_token_ids[audio_codes_mask] # mask depth decoder depth_decoder_ignore_frames_idxs = (labels == -101).nonzero(as_tuple=True) labels_expanded[depth_decoder_ignore_frames_idxs[0], depth_decoder_ignore_frames_idxs[1], 1:] = -100 labels = labels_expanded return {"inputs_embeds": inputs_embeds, "labels": labels} def prepare_inputs_for_generation( self, input_ids: torch.LongTensor, past_key_values: Optional[Cache] = None, attention_mask: Optional[torch.LongTensor] = None, inputs_embeds: Optional[torch.FloatTensor] = None, cache_position: Optional[torch.LongTensor] = None, **kwargs, ): model_inputs = super().prepare_inputs_for_generation( input_ids=input_ids, past_key_values=past_key_values, attention_mask=attention_mask, inputs_embeds=inputs_embeds, cache_position=cache_position, **kwargs, ) if input_ids is not None and input_ids.ndim == 2 and model_inputs.get("inputs_embeds") is None: merged_inputs = self._merge_input_ids_with_input_values( input_ids=input_ids, input_values=kwargs.get("input_values"), input_values_cutoffs=kwargs.get("input_values_cutoffs"), labels=kwargs.get("labels"), ) model_inputs.update( {"inputs_embeds": merged_inputs["inputs_embeds"], "labels": merged_inputs["labels"], "input_ids": None} ) return model_inputs @can_return_tuple @auto_docstring def forward( self, input_ids: torch.LongTensor = None, input_values: Optional[torch.Tensor] = None, attention_mask: Optional[torch.Tensor] = None, input_values_cutoffs: Optional[torch.Tensor] = None, position_ids: Optional[torch.LongTensor] = None, past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, cache_position: Optional[torch.LongTensor] = None, logits_to_keep: Union[int, torch.Tensor] = 0, **kwargs: Unpack[KwargsForCausalLM], ) -> Union[Tuple, CsmOutputWithPast]: r""" 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() ```""" output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions output_hidden_states = ( output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states ) if input_ids is not None and input_ids.ndim == 2: merged_inputs = self._merge_input_ids_with_input_values( input_ids, input_values, input_values_cutoffs, labels ) inputs_embeds = merged_inputs["inputs_embeds"] labels = merged_inputs["labels"] input_ids = None backbone_outputs = self.backbone_model( input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, past_key_values=past_key_values, inputs_embeds=inputs_embeds, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, cache_position=cache_position, **kwargs, ) backbone_hidden_states = backbone_outputs[0] # Only compute necessary logits, and do not upcast them to float if we are not computing the loss slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep backbone_logits = self.lm_head(backbone_hidden_states[:, slice_indices, :]) loss = None backbone_loss = None depth_decoder_loss = None depth_decoder_outputs = None if labels is not None: # select first codebook as labels for the backbone model backbone_labels = labels[:, :, 0] backbone_loss = self.loss_function( logits=backbone_logits, labels=backbone_labels, vocab_size=self.config.vocab_size, **kwargs ) # for the depth decoder, we need to select the frames to train on # those are frames where the label is not uniformly `ignore_index` along the codebook dimension train_mask = ~(labels[:, :, 1:] == -100).all(dim=-1) depth_decoder_input_ids = labels[train_mask][..., : self.config.num_codebooks - 1] # add place holder in position 0 that will be replaced by the backbone_last_hidden_state depth_decoder_input_ids = nn.functional.pad(depth_decoder_input_ids, (1, 0), value=0) train_idxs = train_mask.nonzero(as_tuple=True) backbone_last_hidden_states = backbone_hidden_states[train_idxs[0], train_idxs[1] - 1, :] depth_decoder_labels = labels[train_mask] depth_decoder_outputs = self.depth_decoder( input_ids=depth_decoder_input_ids, backbone_last_hidden_state=backbone_last_hidden_states, use_cache=use_cache, output_attentions=output_attentions, output_hidden_states=output_hidden_states, return_dict=True, labels=depth_decoder_labels, ) depth_decoder_loss = depth_decoder_outputs.loss loss = backbone_loss + depth_decoder_loss return CsmOutputWithPast( loss=loss, backbone_loss=backbone_loss, depth_decoder_loss=depth_decoder_loss, logits=backbone_logits, past_key_values=backbone_outputs.past_key_values, hidden_states=backbone_outputs.hidden_states, attentions=backbone_outputs.attentions, depth_decoder_logits=depth_decoder_outputs.logits if depth_decoder_outputs is not None else None, depth_decoder_past_key_values=depth_decoder_outputs.past_key_values if depth_decoder_outputs is not None else None, depth_decoder_hidden_states=depth_decoder_outputs.hidden_states if depth_decoder_outputs is not None else None, depth_decoder_attentions=depth_decoder_outputs.attentions if depth_decoder_outputs is not None else None, ) @staticmethod def _prepare_4d_causal_attention_mask_with_cache_position( attention_mask: torch.Tensor, sequence_length: int, target_length: int, dtype: torch.dtype, cache_position: torch.Tensor, batch_size: int, **kwargs, ): """ 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. """ if attention_mask is not None and attention_mask.dim() == 4: # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing. causal_mask = attention_mask else: min_dtype = torch.finfo(dtype).min causal_mask = torch.full( (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device ) if sequence_length != 1: causal_mask = torch.triu(causal_mask, diagonal=1) causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1) causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1) if attention_mask is not None: causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit mask_length = attention_mask.shape[-1] padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to( causal_mask.device ) padding_mask = padding_mask == 0 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill( padding_mask, min_dtype ) return causal_mask __all__ = [ "CsmPreTrainedModel", "CsmBackboneModel", "CsmDepthDecoderModel", "CsmDepthDecoderForCausalLM", "CsmForConditionalGeneration", ]