Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2023 The Pop2Piano Authors and The HuggingFace Inc. team. All rights reserved.3#4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15"""PyTorch Pop2Piano model."""16 17import copy18import math19from typing import Optional, Union20 21import torch22from torch import nn23from torch.nn import CrossEntropyLoss24 25from transformers.generation import GenerationConfig26 27from ...activations import ACT2FN28from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache29from ...generation import GenerationMixin30from ...modeling_attn_mask_utils import AttentionMaskConverter31from ...modeling_layers import GradientCheckpointingLayer32from ...modeling_outputs import BaseModelOutput, BaseModelOutputWithPastAndCrossAttentions, Seq2SeqLMOutput33from ...modeling_utils import PreTrainedModel34from ...pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer35from ...utils import auto_docstring, is_torch_flex_attn_available, is_torch_fx_proxy, is_torchdynamo_compiling, logging36from ...utils.deprecation import deprecate_kwarg37from .configuration_pop2piano import Pop2PianoConfig38 39 40if is_torch_flex_attn_available():41 from torch.nn.attention.flex_attention import BlockMask42 43 from ...integrations.flex_attention import make_flex_block_causal_mask44 45 46logger = logging.get_logger(__name__)47 48_load_pop2piano_layer_norm = True49 50try:51 from apex.normalization import FusedRMSNorm52 53 _load_pop2piano_layer_norm = False54 55 logger.info("Discovered apex.normalization.FusedRMSNorm - will use it instead of Pop2PianoLayerNorm")56except ImportError:57 # using the normal Pop2PianoLayerNorm58 pass59except Exception:60 logger.warning("Discovered apex but it failed to load, falling back to Pop2PianoLayerNorm")61 pass62 63 64# Copied from transformers.models.t5.modeling_t5.T5LayerNorm with T5->Pop2Piano65class Pop2PianoLayerNorm(nn.Module):66 def __init__(self, hidden_size, eps=1e-6):67 """68 Construct a layernorm module in the Pop2Piano style. No bias and no subtraction of mean.69 """70 super().__init__()71 self.weight = nn.Parameter(torch.ones(hidden_size))72 self.variance_epsilon = eps73 74 def forward(self, hidden_states):75 # Pop2Piano uses a layer_norm which only scales and doesn't shift, which is also known as Root Mean76 # Square Layer Normalization https://huggingface.co/papers/1910.07467 thus variance is calculated77 # w/o mean and there is no bias. Additionally we want to make sure that the accumulation for78 # half-precision inputs is done in fp3279 80 variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)81 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)82 83 # convert into half-precision if necessary84 if self.weight.dtype in [torch.float16, torch.bfloat16]:85 hidden_states = hidden_states.to(self.weight.dtype)86 87 return self.weight * hidden_states88 89 90if not _load_pop2piano_layer_norm:91 Pop2PianoLayerNorm = FusedRMSNorm92 93 94# Copied from transformers.models.t5.modeling_t5.T5DenseActDense with T5->Pop2Piano,t5->pop2piano95class Pop2PianoDenseActDense(nn.Module):96 def __init__(self, config: Pop2PianoConfig):97 super().__init__()98 self.wi = nn.Linear(config.d_model, config.d_ff, bias=False)99 self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)100 self.dropout = nn.Dropout(config.dropout_rate)101 self.act = ACT2FN[config.dense_act_fn]102 103 def forward(self, hidden_states):104 hidden_states = self.wi(hidden_states)105 hidden_states = self.act(hidden_states)106 hidden_states = self.dropout(hidden_states)107 if (108 isinstance(self.wo.weight, torch.Tensor)109 and hidden_states.dtype != self.wo.weight.dtype110 and self.wo.weight.dtype != torch.int8111 ):112 hidden_states = hidden_states.to(self.wo.weight.dtype)113 hidden_states = self.wo(hidden_states)114 return hidden_states115 116 117# Copied from transformers.models.t5.modeling_t5.T5DenseGatedActDense with T5->Pop2Piano118class Pop2PianoDenseGatedActDense(nn.Module):119 def __init__(self, config: Pop2PianoConfig):120 super().__init__()121 self.wi_0 = nn.Linear(config.d_model, config.d_ff, bias=False)122 self.wi_1 = nn.Linear(config.d_model, config.d_ff, bias=False)123 self.wo = nn.Linear(config.d_ff, config.d_model, bias=False)124 self.dropout = nn.Dropout(config.dropout_rate)125 self.act = ACT2FN[config.dense_act_fn]126 127 def forward(self, hidden_states):128 hidden_gelu = self.act(self.wi_0(hidden_states))129 hidden_linear = self.wi_1(hidden_states)130 hidden_states = hidden_gelu * hidden_linear131 hidden_states = self.dropout(hidden_states)132 133 # To make 8bit quantization work for google/flan-t5-xxl, self.wo is kept in float32.134 # See https://github.com/huggingface/transformers/issues/20287135 # we also make sure the weights are not in `int8` in case users will force `_keep_in_fp32_modules` to be `None``136 if (137 isinstance(self.wo.weight, torch.Tensor)138 and hidden_states.dtype != self.wo.weight.dtype139 and self.wo.weight.dtype != torch.int8140 ):141 hidden_states = hidden_states.to(self.wo.weight.dtype)142 143 hidden_states = self.wo(hidden_states)144 return hidden_states145 146 147# Copied from transformers.models.t5.modeling_t5.T5LayerFF with T5->Pop2Piano148class Pop2PianoLayerFF(nn.Module):149 def __init__(self, config: Pop2PianoConfig):150 super().__init__()151 if config.is_gated_act:152 self.DenseReluDense = Pop2PianoDenseGatedActDense(config)153 else:154 self.DenseReluDense = Pop2PianoDenseActDense(config)155 156 self.layer_norm = Pop2PianoLayerNorm(config.d_model, eps=config.layer_norm_epsilon)157 self.dropout = nn.Dropout(config.dropout_rate)158 159 def forward(self, hidden_states):160 forwarded_states = self.layer_norm(hidden_states)161 forwarded_states = self.DenseReluDense(forwarded_states)162 hidden_states = hidden_states + self.dropout(forwarded_states)163 return hidden_states164 165 166# Copied from transformers.models.t5.modeling_t5.T5Attention with T5->Pop2Piano,t5->pop2piano167class Pop2PianoAttention(nn.Module):168 def __init__(169 self,170 config: Pop2PianoConfig,171 has_relative_attention_bias=False,172 layer_idx: Optional[int] = None,173 ):174 super().__init__()175 self.is_decoder = config.is_decoder176 self.has_relative_attention_bias = has_relative_attention_bias177 self.relative_attention_num_buckets = config.relative_attention_num_buckets178 self.relative_attention_max_distance = config.relative_attention_max_distance179 self.d_model = config.d_model180 self.key_value_proj_dim = config.d_kv181 self.n_heads = config.num_heads182 self.dropout = config.dropout_rate183 self.inner_dim = self.n_heads * self.key_value_proj_dim184 self.layer_idx = layer_idx185 if layer_idx is None and self.is_decoder:186 logger.warning_once(187 f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "188 "will to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "189 "when creating this class."190 )191 192 # Mesh TensorFlow initialization to avoid scaling before softmax193 self.q = nn.Linear(self.d_model, self.inner_dim, bias=False)194 self.k = nn.Linear(self.d_model, self.inner_dim, bias=False)195 self.v = nn.Linear(self.d_model, self.inner_dim, bias=False)196 self.o = nn.Linear(self.inner_dim, self.d_model, bias=False)197 198 if self.has_relative_attention_bias:199 self.relative_attention_bias = nn.Embedding(self.relative_attention_num_buckets, self.n_heads)200 self.pruned_heads = set()201 self.gradient_checkpointing = False202 203 def prune_heads(self, heads):204 if len(heads) == 0:205 return206 heads, index = find_pruneable_heads_and_indices(207 heads, self.n_heads, self.key_value_proj_dim, self.pruned_heads208 )209 # Prune linear layers210 self.q = prune_linear_layer(self.q, index)211 self.k = prune_linear_layer(self.k, index)212 self.v = prune_linear_layer(self.v, index)213 self.o = prune_linear_layer(self.o, index, dim=1)214 # Update hyper params215 self.n_heads = self.n_heads - len(heads)216 self.inner_dim = self.key_value_proj_dim * self.n_heads217 self.pruned_heads = self.pruned_heads.union(heads)218 219 @staticmethod220 def _relative_position_bucket(relative_position, bidirectional=True, num_buckets=32, max_distance=128):221 """222 Adapted from Mesh Tensorflow:223 https://github.com/tensorflow/mesh/blob/0cb87fe07da627bf0b7e60475d59f95ed6b5be3d/mesh_tensorflow/transformer/transformer_layers.py#L593224 225 Translate relative position to a bucket number for relative attention. The relative position is defined as226 memory_position - query_position, i.e. the distance in tokens from the attending position to the attended-to227 position. If bidirectional=False, then positive relative positions are invalid. We use smaller buckets for228 small absolute relative_position and larger buckets for larger absolute relative_positions. All relative229 positions >=max_distance map to the same bucket. All relative positions <=-max_distance map to the same bucket.230 This should allow for more graceful generalization to longer sequences than the model has been trained on231 232 Args:233 relative_position: an int32 Tensor234 bidirectional: a boolean - whether the attention is bidirectional235 num_buckets: an integer236 max_distance: an integer237 238 Returns:239 a Tensor with the same shape as relative_position, containing int32 values in the range [0, num_buckets)240 """241 relative_buckets = 0242 if bidirectional:243 num_buckets //= 2244 relative_buckets += (relative_position > 0).to(torch.long) * num_buckets245 relative_position = torch.abs(relative_position)246 else:247 relative_position = -torch.min(relative_position, torch.zeros_like(relative_position))248 # now relative_position is in the range [0, inf)249 250 # half of the buckets are for exact increments in positions251 max_exact = num_buckets // 2252 is_small = relative_position < max_exact253 254 # The other half of the buckets are for logarithmically bigger bins in positions up to max_distance255 relative_position_if_large = max_exact + (256 torch.log(relative_position.float() / max_exact)257 / math.log(max_distance / max_exact)258 * (num_buckets - max_exact)259 ).to(torch.long)260 relative_position_if_large = torch.min(261 relative_position_if_large, torch.full_like(relative_position_if_large, num_buckets - 1)262 )263 264 relative_buckets += torch.where(is_small, relative_position, relative_position_if_large)265 return relative_buckets266 267 def compute_bias(self, query_length, key_length, device=None, cache_position=None):268 """Compute binned relative position bias"""269 if device is None:270 device = self.relative_attention_bias.weight.device271 if cache_position is None:272 context_position = torch.arange(query_length, dtype=torch.long, device=device)[:, None]273 else:274 context_position = cache_position[:, None].to(device)275 memory_position = torch.arange(key_length, dtype=torch.long, device=device)[None, :]276 relative_position = memory_position - context_position # shape (query_length, key_length)277 relative_position_bucket = self._relative_position_bucket(278 relative_position, # shape (query_length, key_length)279 bidirectional=(not self.is_decoder),280 num_buckets=self.relative_attention_num_buckets,281 max_distance=self.relative_attention_max_distance,282 )283 values = self.relative_attention_bias(relative_position_bucket) # shape (query_length, key_length, num_heads)284 values = values.permute([2, 0, 1]).unsqueeze(0) # shape (1, num_heads, query_length, key_length)285 return values286 287 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")288 def forward(289 self,290 hidden_states,291 mask=None,292 key_value_states=None,293 position_bias=None,294 past_key_values=None,295 layer_head_mask=None,296 query_length=None,297 use_cache=False,298 output_attentions=False,299 cache_position=None,300 ):301 """302 Self-attention (if key_value_states is None) or attention over source sentence (provided by key_value_states).303 """304 # Input is (batch_size, seq_length, dim)305 # Mask is (batch_size, 1, 1, key_length) (non-causal encoder) or (batch_size, 1, seq_length, key_length) (causal decoder)306 batch_size, seq_length = hidden_states.shape[:2]307 308 # if key_value_states are provided this layer is used as a cross-attention layer for the decoder309 is_cross_attention = key_value_states is not None310 311 query_states = self.q(hidden_states)312 query_states = query_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)313 314 # Check is encoder-decoder model is being used. Otherwise we'll get `DynamicCache`315 is_updated = False316 if isinstance(past_key_values, EncoderDecoderCache):317 is_updated = past_key_values.is_updated.get(self.layer_idx)318 if is_cross_attention:319 # after the first generated id, we can subsequently re-use all key/value_states from cache320 curr_past_key_value = past_key_values.cross_attention_cache321 else:322 curr_past_key_value = past_key_values.self_attention_cache323 else:324 curr_past_key_value = past_key_values325 326 current_states = key_value_states if is_cross_attention else hidden_states327 if is_cross_attention and past_key_values is not None and is_updated:328 # reuse k,v, cross_attentions329 key_states = curr_past_key_value.layers[self.layer_idx].keys330 value_states = curr_past_key_value.layers[self.layer_idx].values331 else:332 key_states = self.k(current_states)333 value_states = self.v(current_states)334 key_states = key_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)335 value_states = value_states.view(batch_size, -1, self.n_heads, self.key_value_proj_dim).transpose(1, 2)336 337 if past_key_values is not None:338 # save all key/value_states to cache to be re-used for fast auto-regressive generation339 cache_position = cache_position if not is_cross_attention else None340 key_states, value_states = curr_past_key_value.update(341 key_states, value_states, self.layer_idx, {"cache_position": cache_position}342 )343 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls344 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):345 past_key_values.is_updated[self.layer_idx] = True346 347 # compute scores, equivalent of torch.einsum("bnqd,bnkd->bnqk", query_states, key_states), compatible with onnx op>9348 scores = torch.matmul(query_states, key_states.transpose(3, 2))349 350 if position_bias is None:351 key_length = key_states.shape[-2]352 # cache position is 0-indexed so we add 1 to get the real length of queries (aka with past)353 real_seq_length = query_length if query_length is not None else cache_position[-1] + 1354 if not self.has_relative_attention_bias:355 position_bias = torch.zeros(356 (1, self.n_heads, seq_length, key_length), device=scores.device, dtype=scores.dtype357 )358 if self.gradient_checkpointing and self.training:359 position_bias.requires_grad = True360 else:361 position_bias = self.compute_bias(362 real_seq_length, key_length, device=scores.device, cache_position=cache_position363 )364 position_bias = position_bias[:, :, -seq_length:, :]365 366 if mask is not None:367 causal_mask = mask[:, :, :, : key_states.shape[-2]]368 position_bias = position_bias + causal_mask369 370 if self.pruned_heads:371 mask = torch.ones(position_bias.shape[1])372 mask[list(self.pruned_heads)] = 0373 position_bias_masked = position_bias[:, mask.bool()]374 else:375 position_bias_masked = position_bias376 377 scores += position_bias_masked378 379 # (batch_size, n_heads, seq_length, key_length)380 attn_weights = nn.functional.softmax(scores.float(), dim=-1).type_as(scores)381 attn_weights = nn.functional.dropout(attn_weights, p=self.dropout, training=self.training)382 383 # Mask heads if we want to384 if layer_head_mask is not None:385 attn_weights = attn_weights * layer_head_mask386 387 attn_output = torch.matmul(attn_weights, value_states)388 389 attn_output = attn_output.transpose(1, 2).contiguous()390 attn_output = attn_output.view(batch_size, -1, self.inner_dim)391 attn_output = self.o(attn_output)392 393 outputs = (attn_output, position_bias)394 395 if output_attentions:396 outputs = outputs + (attn_weights,)397 return outputs398 399 400# Copied from transformers.models.t5.modeling_t5.T5LayerSelfAttention with T5->Pop2Piano,t5->pop2piano401class Pop2PianoLayerSelfAttention(nn.Module):402 def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None):403 super().__init__()404 self.SelfAttention = Pop2PianoAttention(405 config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx406 )407 self.layer_norm = Pop2PianoLayerNorm(config.d_model, eps=config.layer_norm_epsilon)408 self.dropout = nn.Dropout(config.dropout_rate)409 410 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")411 def forward(412 self,413 hidden_states,414 attention_mask=None,415 position_bias=None,416 layer_head_mask=None,417 past_key_values=None,418 use_cache=False,419 output_attentions=False,420 cache_position=None,421 ):422 normed_hidden_states = self.layer_norm(hidden_states)423 attention_output = self.SelfAttention(424 normed_hidden_states,425 mask=attention_mask,426 position_bias=position_bias,427 layer_head_mask=layer_head_mask,428 past_key_values=past_key_values,429 use_cache=use_cache,430 output_attentions=output_attentions,431 cache_position=cache_position,432 )433 hidden_states = hidden_states + self.dropout(attention_output[0])434 outputs = (hidden_states,) + attention_output[1:] # add attentions if we output them435 return outputs436 437 438# Copied from transformers.models.t5.modeling_t5.T5LayerCrossAttention with T5->Pop2Piano,t5->pop2piano439class Pop2PianoLayerCrossAttention(nn.Module):440 def __init__(self, config, layer_idx: Optional[int] = None):441 super().__init__()442 self.EncDecAttention = Pop2PianoAttention(config, has_relative_attention_bias=False, layer_idx=layer_idx)443 self.layer_norm = Pop2PianoLayerNorm(config.d_model, eps=config.layer_norm_epsilon)444 self.dropout = nn.Dropout(config.dropout_rate)445 446 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")447 def forward(448 self,449 hidden_states,450 key_value_states,451 attention_mask=None,452 position_bias=None,453 layer_head_mask=None,454 past_key_values=None,455 use_cache=False,456 query_length=None,457 output_attentions=False,458 cache_position=None,459 ):460 normed_hidden_states = self.layer_norm(hidden_states)461 attention_output = self.EncDecAttention(462 normed_hidden_states,463 mask=attention_mask,464 key_value_states=key_value_states,465 position_bias=position_bias,466 layer_head_mask=layer_head_mask,467 past_key_values=past_key_values,468 use_cache=use_cache,469 query_length=query_length,470 output_attentions=output_attentions,471 cache_position=cache_position,472 )473 layer_output = hidden_states + self.dropout(attention_output[0])474 outputs = (layer_output,) + attention_output[1:] # add attentions if we output them475 return outputs476 477 478# Copied from transformers.models.t5.modeling_t5.T5Block with T5->Pop2Piano,t5->pop2piano479class Pop2PianoBlock(GradientCheckpointingLayer):480 def __init__(self, config, has_relative_attention_bias=False, layer_idx: Optional[int] = None):481 super().__init__()482 self.is_decoder = config.is_decoder483 self.layer = nn.ModuleList()484 self.layer.append(485 Pop2PianoLayerSelfAttention(486 config, has_relative_attention_bias=has_relative_attention_bias, layer_idx=layer_idx487 )488 )489 if self.is_decoder:490 self.layer.append(Pop2PianoLayerCrossAttention(config, layer_idx=layer_idx))491 492 self.layer.append(Pop2PianoLayerFF(config))493 494 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")495 def forward(496 self,497 hidden_states,498 attention_mask=None,499 position_bias=None,500 encoder_hidden_states=None,501 encoder_attention_mask=None,502 encoder_decoder_position_bias=None,503 layer_head_mask=None,504 cross_attn_layer_head_mask=None,505 past_key_values=None,506 use_cache=False,507 output_attentions=False,508 return_dict=True,509 cache_position=None,510 ):511 self_attention_outputs = self.layer[0](512 hidden_states,513 attention_mask=attention_mask,514 position_bias=position_bias,515 layer_head_mask=layer_head_mask,516 past_key_values=past_key_values,517 use_cache=use_cache,518 output_attentions=output_attentions,519 cache_position=cache_position,520 )521 hidden_states = self_attention_outputs[0]522 attention_outputs = self_attention_outputs[1:] # Keep self-attention outputs and relative position weights523 524 # clamp inf values to enable fp16 training525 if hidden_states.dtype == torch.float16:526 clamp_value = torch.where(527 torch.isinf(hidden_states).any(),528 torch.finfo(hidden_states.dtype).max - 1000,529 torch.finfo(hidden_states.dtype).max,530 )531 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)532 533 do_cross_attention = self.is_decoder and encoder_hidden_states is not None534 if do_cross_attention:535 cross_attention_outputs = self.layer[1](536 hidden_states,537 key_value_states=encoder_hidden_states,538 attention_mask=encoder_attention_mask,539 position_bias=encoder_decoder_position_bias,540 layer_head_mask=cross_attn_layer_head_mask,541 past_key_values=past_key_values,542 query_length=cache_position[-1] + 1,543 use_cache=use_cache,544 output_attentions=output_attentions,545 )546 hidden_states = cross_attention_outputs[0]547 548 # clamp inf values to enable fp16 training549 if hidden_states.dtype == torch.float16:550 clamp_value = torch.where(551 torch.isinf(hidden_states).any(),552 torch.finfo(hidden_states.dtype).max - 1000,553 torch.finfo(hidden_states.dtype).max,554 )555 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)556 557 # Keep cross-attention outputs and relative position weights558 attention_outputs = attention_outputs + cross_attention_outputs[1:]559 560 # Apply Feed Forward layer561 hidden_states = self.layer[-1](hidden_states)562 563 # clamp inf values to enable fp16 training564 if hidden_states.dtype == torch.float16:565 clamp_value = torch.where(566 torch.isinf(hidden_states).any(),567 torch.finfo(hidden_states.dtype).max - 1000,568 torch.finfo(hidden_states.dtype).max,569 )570 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)571 572 outputs = (hidden_states,)573 574 return (575 outputs + attention_outputs576 ) # hidden-states, (self-attention position bias), (self-attention weights), (cross-attention position bias), (cross-attention weights)577 578 579@auto_docstring580class Pop2PianoPreTrainedModel(PreTrainedModel):581 config: Pop2PianoConfig582 base_model_prefix = "transformer"583 is_parallelizable = False584 supports_gradient_checkpointing = True585 586 _can_compile_fullgraph = False587 _no_split_modules = ["Pop2PianoBlock"]588 _keep_in_fp32_modules = ["wo"]589 590 def _init_weights(self, module):591 """Initialize the weights"""592 factor = self.config.initializer_factor # Used for testing weights initialization593 if isinstance(module, Pop2PianoLayerNorm):594 module.weight.data.fill_(factor * 1.0)595 elif isinstance(module, Pop2PianoConcatEmbeddingToMel):596 module.embedding.weight.data.normal_(mean=0.0, std=factor * 1.0)597 elif isinstance(module, Pop2PianoForConditionalGeneration):598 # Mesh TensorFlow embeddings initialization599 # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L1624600 module.shared.weight.data.normal_(mean=0.0, std=factor * 1.0)601 if hasattr(module, "lm_head") and not self.config.tie_word_embeddings:602 module.lm_head.weight.data.normal_(mean=0.0, std=factor * 1.0)603 elif isinstance(module, Pop2PianoDenseActDense):604 # Mesh TensorFlow FF initialization605 # See https://github.com/tensorflow/mesh/blob/master/mesh_tensorflow/transformer/transformer_layers.py#L56606 # and https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L89607 module.wi.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))608 if hasattr(module.wi, "bias") and module.wi.bias is not None:609 module.wi.bias.data.zero_()610 module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))611 if hasattr(module.wo, "bias") and module.wo.bias is not None:612 module.wo.bias.data.zero_()613 elif isinstance(module, Pop2PianoDenseGatedActDense):614 module.wi_0.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))615 if hasattr(module.wi_0, "bias") and module.wi_0.bias is not None:616 module.wi_0.bias.data.zero_()617 module.wi_1.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_model) ** -0.5))618 if hasattr(module.wi_1, "bias") and module.wi_1.bias is not None:619 module.wi_1.bias.data.zero_()620 module.wo.weight.data.normal_(mean=0.0, std=factor * ((self.config.d_ff) ** -0.5))621 if hasattr(module.wo, "bias") and module.wo.bias is not None:622 module.wo.bias.data.zero_()623 elif isinstance(module, Pop2PianoAttention):624 # Mesh TensorFlow attention initialization to avoid scaling before softmax625 # See https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/transformer/attention.py#L136626 d_model = self.config.d_model627 key_value_proj_dim = self.config.d_kv628 n_heads = self.config.num_heads629 module.q.weight.data.normal_(mean=0.0, std=factor * ((d_model * key_value_proj_dim) ** -0.5))630 module.k.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))631 module.v.weight.data.normal_(mean=0.0, std=factor * (d_model**-0.5))632 module.o.weight.data.normal_(mean=0.0, std=factor * ((n_heads * key_value_proj_dim) ** -0.5))633 if module.has_relative_attention_bias:634 module.relative_attention_bias.weight.data.normal_(mean=0.0, std=factor * ((d_model) ** -0.5))635 636 def _shift_right(self, input_ids):637 decoder_start_token_id = self.config.decoder_start_token_id638 pad_token_id = self.config.pad_token_id639 640 if decoder_start_token_id is None:641 raise ValueError(642 "self.model.config.decoder_start_token_id has to be defined. In Pop2Piano it is usually set to the pad_token_id."643 )644 645 # shift inputs to the right646 if is_torch_fx_proxy(input_ids):647 # Item assignment is not supported natively for proxies.648 shifted_input_ids = torch.full(input_ids.shape[:-1] + (1,), decoder_start_token_id)649 shifted_input_ids = torch.cat([shifted_input_ids, input_ids[..., :-1]], dim=-1)650 else:651 shifted_input_ids = input_ids.new_zeros(input_ids.shape)652 shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()653 shifted_input_ids[..., 0] = decoder_start_token_id654 655 if pad_token_id is None:656 raise ValueError("self.model.config.pad_token_id has to be defined.")657 # replace possible -100 values in labels by `pad_token_id`658 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)659 660 return shifted_input_ids661 662 663class Pop2PianoStack(Pop2PianoPreTrainedModel):664 # Copied from transformers.models.t5.modeling_t5.T5Stack.__init__ with T5->Pop2Piano,t5->pop2piano665 def __init__(self, config, embed_tokens=None):666 super().__init__(config)667 668 self.embed_tokens = embed_tokens669 self.is_decoder = config.is_decoder670 671 self.block = nn.ModuleList(672 [673 Pop2PianoBlock(config, has_relative_attention_bias=bool(i == 0), layer_idx=i)674 for i in range(config.num_layers)675 ]676 )677 self.final_layer_norm = Pop2PianoLayerNorm(config.d_model, eps=config.layer_norm_epsilon)678 self.dropout = nn.Dropout(config.dropout_rate)679 680 # Initialize weights and apply final processing681 self.post_init()682 # Model parallel683 self.model_parallel = False684 self.device_map = None685 self.gradient_checkpointing = False686 687 # Copied from transformers.models.t5.modeling_t5.T5Stack.set_input_embeddings688 def set_input_embeddings(self, new_embeddings):689 self.embed_tokens = new_embeddings690 691 def forward(692 self,693 input_ids=None,694 attention_mask=None,695 encoder_hidden_states=None,696 encoder_attention_mask=None,697 inputs_embeds=None,698 head_mask=None,699 cross_attn_head_mask=None,700 past_key_values=None,701 use_cache=None,702 output_attentions=None,703 output_hidden_states=None,704 return_dict=None,705 cache_position=None,706 ):707 use_cache = use_cache if use_cache is not None else self.config.use_cache708 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions709 output_hidden_states = (710 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states711 )712 return_dict = return_dict if return_dict is not None else self.config.use_return_dict713 714 if input_ids is not None and inputs_embeds is not None:715 err_msg_prefix = "decoder_" if self.is_decoder else ""716 raise ValueError(717 f"You cannot specify both {err_msg_prefix}input_ids and {err_msg_prefix}inputs_embeds at the same time"718 )719 elif input_ids is not None:720 input_shape = input_ids.size()721 input_ids = input_ids.view(-1, input_shape[-1])722 elif inputs_embeds is not None:723 input_shape = inputs_embeds.size()[:-1]724 else:725 err_msg_prefix = "decoder_" if self.is_decoder else ""726 raise ValueError(f"You have to specify either {err_msg_prefix}input_ids or {err_msg_prefix}inputs_embeds")727 728 if self.gradient_checkpointing and self.training:729 if use_cache:730 logger.warning_once(731 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."732 )733 use_cache = False734 735 if inputs_embeds is None:736 if self.embed_tokens is None:737 raise ValueError("You have to initialize the model with valid token embeddings")738 inputs_embeds = self.embed_tokens(input_ids)739 740 batch_size, seq_length = input_shape741 742 if use_cache is True:743 if not self.is_decoder:744 raise ValueError(f"`use_cache` can only be set to `True` if {self} is used as a decoder")745 746 if self.is_decoder:747 if use_cache and past_key_values is None:748 if self.config.is_encoder_decoder:749 past_key_values = EncoderDecoderCache(750 DynamicCache(config=self.config), DynamicCache(config=self.config)751 )752 else:753 past_key_values = DynamicCache(config=self.config)754 elif not self.is_decoder:755 # do not pass cache object down the line for encoder stack756 # it messes indexing later in decoder-stack because cache object is modified in-place757 past_key_values = None758 759 past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0760 if cache_position is None:761 cache_position = torch.arange(762 past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device763 )764 765 if attention_mask is None and not is_torchdynamo_compiling():766 # required mask seq length can be calculated via length of past cache767 mask_seq_length = past_key_values_length + seq_length768 attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)769 770 if self.config.is_decoder:771 causal_mask = self._update_causal_mask(772 attention_mask,773 inputs_embeds,774 cache_position,775 past_key_values.self_attention_cache776 if isinstance(past_key_values, EncoderDecoderCache)777 else past_key_values,778 output_attentions,779 )780 else:781 causal_mask = attention_mask[:, None, None, :]782 causal_mask = causal_mask.to(dtype=inputs_embeds.dtype)783 causal_mask = (1.0 - causal_mask) * torch.finfo(inputs_embeds.dtype).min784 785 # If a 2D or 3D attention mask is provided for the cross-attention786 # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]787 if self.is_decoder and encoder_hidden_states is not None:788 encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()789 encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)790 if encoder_attention_mask is None:791 encoder_attention_mask = torch.ones(encoder_hidden_shape, device=inputs_embeds.device)792 encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)793 else:794 encoder_extended_attention_mask = None795 796 # Prepare head mask if needed797 head_mask = self.get_head_mask(head_mask, self.config.num_layers)798 cross_attn_head_mask = self.get_head_mask(cross_attn_head_mask, self.config.num_layers)799 all_hidden_states = () if output_hidden_states else None800 all_attentions = () if output_attentions else None801 all_cross_attentions = () if (output_attentions and self.is_decoder) else None802 position_bias = None803 encoder_decoder_position_bias = None804 805 hidden_states = self.dropout(inputs_embeds)806 807 for i, layer_module in enumerate(self.block):808 layer_head_mask = head_mask[i]809 cross_attn_layer_head_mask = cross_attn_head_mask[i]810 if output_hidden_states:811 all_hidden_states = all_hidden_states + (hidden_states,)812 813 layer_outputs = layer_module(814 hidden_states,815 causal_mask,816 position_bias,817 encoder_hidden_states,818 encoder_extended_attention_mask,819 encoder_decoder_position_bias, # as a positional argument for gradient checkpointing820 layer_head_mask=layer_head_mask,821 cross_attn_layer_head_mask=cross_attn_layer_head_mask,822 past_key_values=past_key_values,823 use_cache=use_cache,824 output_attentions=output_attentions,825 cache_position=cache_position,826 )827 828 hidden_states = layer_outputs[0]829 830 # We share the position biases between the layers - the first layer store them831 # layer_outputs = hidden-states, key-value-states (self-attention position bias), (self-attention weights),832 # (cross-attention position bias), (cross-attention weights)833 position_bias = layer_outputs[1]834 if self.is_decoder and encoder_hidden_states is not None:835 encoder_decoder_position_bias = layer_outputs[3 if output_attentions else 2]836 837 if output_attentions:838 all_attentions = all_attentions + (layer_outputs[2],)839 if self.is_decoder:840 all_cross_attentions = all_cross_attentions + (layer_outputs[4],)841 842 hidden_states = self.final_layer_norm(hidden_states)843 hidden_states = self.dropout(hidden_states)844 845 # Add last layer846 if output_hidden_states:847 all_hidden_states = all_hidden_states + (hidden_states,)848 849 if not return_dict:850 return tuple(851 v852 for v in [853 hidden_states,854 past_key_values,855 all_hidden_states,856 all_attentions,857 all_cross_attentions,858 ]859 if v is not None860 )861 return BaseModelOutputWithPastAndCrossAttentions(862 last_hidden_state=hidden_states,863 past_key_values=past_key_values,864 hidden_states=all_hidden_states,865 attentions=all_attentions,866 cross_attentions=all_cross_attentions,867 )868 869 # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._update_causal_mask870 def _update_causal_mask(871 self,872 attention_mask: Union[torch.Tensor, "BlockMask"],873 input_tensor: torch.Tensor,874 cache_position: torch.Tensor,875 past_key_values: Cache,876 output_attentions: bool = False,877 ):878 if self.config._attn_implementation == "flash_attention_2":879 if attention_mask is not None and (attention_mask == 0.0).any():880 return attention_mask881 return None882 if self.config._attn_implementation == "flex_attention":883 if isinstance(attention_mask, torch.Tensor):884 attention_mask = make_flex_block_causal_mask(attention_mask)885 return attention_mask886 887 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in888 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail889 # to infer the attention mask.890 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0891 using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False892 893 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward894 if self.config._attn_implementation == "sdpa" and not using_compilable_cache and not output_attentions:895 if AttentionMaskConverter._ignore_causal_mask_sdpa(896 attention_mask,897 inputs_embeds=input_tensor,898 past_key_values_length=past_seen_tokens,899 is_training=self.training,900 ):901 return None902 903 dtype = input_tensor.dtype904 sequence_length = input_tensor.shape[1]905 if using_compilable_cache:906 target_length = past_key_values.get_max_cache_shape()907 else:908 target_length = (909 attention_mask.shape[-1]910 if isinstance(attention_mask, torch.Tensor)911 else past_seen_tokens + sequence_length + 1912 )913 914 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).915 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(916 attention_mask,917 sequence_length=sequence_length,918 target_length=target_length,919 dtype=dtype,920 cache_position=cache_position,921 batch_size=input_tensor.shape[0],922 )923 924 if (925 self.config._attn_implementation == "sdpa"926 and attention_mask is not None927 and attention_mask.device.type in ["cuda", "xpu", "npu"]928 and not output_attentions929 ):930 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when931 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.932 # Details: https://github.com/pytorch/pytorch/issues/110213933 min_dtype = torch.finfo(dtype).min934 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)935 936 return causal_mask937 938 @staticmethod939 # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position940 def _prepare_4d_causal_attention_mask_with_cache_position(941 attention_mask: torch.Tensor,942 sequence_length: int,943 target_length: int,944 dtype: torch.dtype,945 cache_position: torch.Tensor,946 batch_size: int,947 **kwargs,948 ):949 """950 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape951 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.952 953 Args:954 attention_mask (`torch.Tensor`):955 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape956 `(batch_size, 1, query_length, key_value_length)`.957 sequence_length (`int`):958 The sequence length being processed.959 target_length (`int`):960 The target length: when generating with static cache, the mask should be as long as the static cache,961 to account for the 0 padding, the part of the cache that is not filled yet.962 dtype (`torch.dtype`):963 The dtype to use for the 4D attention mask.964 cache_position (`torch.Tensor`):965 Indices depicting the position of the input sequence tokens in the sequence.966 batch_size (`torch.Tensor`):967 Batch size.968 """969 if attention_mask is not None and attention_mask.dim() == 4:970 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.971 causal_mask = attention_mask972 else:973 min_dtype = torch.finfo(dtype).min974 causal_mask = torch.full(975 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device976 )977 if sequence_length != 1:978 causal_mask = torch.triu(causal_mask, diagonal=1)979 causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)980 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)981 if attention_mask is not None:982 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit983 mask_length = attention_mask.shape[-1]984 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(985 causal_mask.device986 )987 padding_mask = padding_mask == 0988 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(989 padding_mask, min_dtype990 )991 992 return causal_mask993 994 995class Pop2PianoConcatEmbeddingToMel(nn.Module):996 """Embedding Matrix for `composer` tokens."""997 998 def __init__(self, config):999 super().__init__()1000 self.embedding = nn.Embedding(num_embeddings=config.composer_vocab_size, embedding_dim=config.d_model)1001 1002 def forward(self, feature, index_value, embedding_offset):1003 index_shifted = index_value - embedding_offset1004 composer_embedding = self.embedding(index_shifted).unsqueeze(1)1005 inputs_embeds = torch.cat([composer_embedding, feature], dim=1)1006 return inputs_embeds1007 1008 1009@auto_docstring(1010 custom_intro="""1011 Pop2Piano Model with a `language modeling` head on top.1012 """1013)1014class Pop2PianoForConditionalGeneration(Pop2PianoPreTrainedModel, GenerationMixin):1015 _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight", "lm_head.weight"]1016 1017 def __init__(self, config: Pop2PianoConfig):1018 super().__init__(config)1019 self.config = config1020 self.model_dim = config.d_model1021 1022 self.shared = nn.Embedding(config.vocab_size, config.d_model)1023 1024 self.mel_conditioner = Pop2PianoConcatEmbeddingToMel(config)1025 1026 encoder_config = copy.deepcopy(config)1027 encoder_config.is_decoder = False1028 encoder_config.use_cache = False1029 encoder_config.tie_encoder_decoder = False1030 1031 self.encoder = Pop2PianoStack(encoder_config, self.shared)1032 1033 decoder_config = copy.deepcopy(config)1034 decoder_config.is_decoder = True1035 decoder_config.tie_encoder_decoder = False1036 decoder_config.num_layers = config.num_decoder_layers1037 self.decoder = Pop2PianoStack(decoder_config, self.shared)1038 1039 self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)1040 1041 # Initialize weights and apply final processing1042 self.post_init()1043 1044 def get_input_embeddings(self):1045 return self.shared1046 1047 def set_input_embeddings(self, new_embeddings):1048 self.shared = new_embeddings1049 self.encoder.set_input_embeddings(new_embeddings)1050 self.decoder.set_input_embeddings(new_embeddings)1051 1052 def get_encoder(self):1053 return self.encoder1054 1055 def get_mel_conditioner_outputs(1056 self,1057 input_features: torch.FloatTensor,1058 composer: str,1059 generation_config: GenerationConfig,1060 attention_mask: Optional[torch.FloatTensor] = None,1061 ):1062 """1063 This method is used to concatenate mel conditioner tokens at the front of the input_features in order to1064 control the type of MIDI token generated by the model.1065 1066 Args:1067 input_features (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`):1068 input features extracted from the feature extractor.1069 composer (`str`):1070 composer token which determines the type of MIDI tokens to be generated.1071 generation_config (`~generation.GenerationConfig`):1072 The generation is used to get the composer-feature_token pair.1073 attention_mask (``, *optional*):1074 For batched generation `input_features` are padded to have the same shape across all examples.1075 `attention_mask` helps to determine which areas were padded and which were not.1076 - 1 for tokens that are **not padded**,1077 - 0 for tokens that are **padded**.1078 """1079 composer_to_feature_token = generation_config.composer_to_feature_token1080 if composer not in composer_to_feature_token:1081 raise ValueError(1082 f"Please choose a composer from {list(composer_to_feature_token.keys())}. Composer received - {composer}"1083 )1084 composer_value = composer_to_feature_token[composer]1085 composer_value = torch.tensor(composer_value, device=self.device)1086 composer_value = composer_value.repeat(input_features.shape[0])1087 1088 embedding_offset = min(composer_to_feature_token.values())1089 1090 input_features = self.mel_conditioner(1091 feature=input_features,1092 index_value=composer_value,1093 embedding_offset=embedding_offset,1094 )1095 if attention_mask is not None:1096 input_features[~attention_mask[:, 0].bool()] = 0.01097 1098 # since self.mel_conditioner adds a new array at the front of inputs_embeds we need to do the same for attention_mask to keep the shapes same1099 attention_mask = torch.concatenate([attention_mask[:, 0].view(-1, 1), attention_mask], axis=1)1100 return input_features, attention_mask1101 1102 return input_features, None1103 1104 @auto_docstring1105 def forward(1106 self,1107 input_ids: Optional[torch.LongTensor] = None,1108 attention_mask: Optional[torch.FloatTensor] = None,1109 decoder_input_ids: Optional[torch.LongTensor] = None,1110 decoder_attention_mask: Optional[torch.BoolTensor] = None,1111 head_mask: Optional[torch.FloatTensor] = None,1112 decoder_head_mask: Optional[torch.FloatTensor] = None,1113 cross_attn_head_mask: Optional[torch.Tensor] = None,1114 encoder_outputs: Optional[tuple[tuple[torch.Tensor]]] = None,1115 past_key_values: Optional[Cache] = None,1116 inputs_embeds: Optional[torch.FloatTensor] = None,1117 input_features: Optional[torch.FloatTensor] = None,1118 decoder_inputs_embeds: Optional[torch.FloatTensor] = None,1119 labels: Optional[torch.LongTensor] = None,1120 use_cache: Optional[bool] = None,1121 output_attentions: Optional[bool] = None,1122 output_hidden_states: Optional[bool] = None,1123 return_dict: Optional[bool] = None,1124 cache_position: Optional[torch.LongTensor] = None,1125 ) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]:1126 r"""1127 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):1128 Indices of input sequence tokens in the vocabulary. Pop2Piano is a model with relative position embeddings1129 so you should be able to pad the inputs on both the right and the left. Indices can be obtained using1130 [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for detail.1131 [What are input IDs?](../glossary#input-ids) To know more on how to prepare `input_ids` for pretraining1132 take a look a [Pop2Piano Training](./Pop2Piano#training).1133 decoder_input_ids (`torch.LongTensor` of shape `(batch_size, target_sequence_length)`, *optional*):1134 Indices of decoder input sequence tokens in the vocabulary. Indices can be obtained using1135 [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and [`PreTrainedTokenizer.__call__`] for details.1136 [What are decoder input IDs?](../glossary#decoder-input-ids) Pop2Piano uses the `pad_token_id` as the1137 starting token for `decoder_input_ids` generation. If `past_key_values` is used, optionally only the last1138 `decoder_input_ids` have to be input (see `past_key_values`). To know more on how to prepare1139 decoder_attention_mask (`torch.BoolTensor` of shape `(batch_size, target_sequence_length)`, *optional*):1140 Default behavior: generate a tensor that ignores pad tokens in `decoder_input_ids`. Causal mask will also1141 be used by default.1142 decoder_head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):1143 Mask to nullify selected heads of the self-attention modules in the decoder. Mask values selected in `[0,1144 1]`:1145 - 1 indicates the head is **not masked**,1146 - 0 indicates the head is **masked**.1147 cross_attn_head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):1148 Mask to nullify selected heads of the cross-attention modules in the decoder. Mask values selected in1149 `[0, 1]`:1150 - 1 indicates the head is **not masked**,1151 - 0 indicates the head is **masked**.1152 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1153 Labels for computing the sequence classification/regression loss. Indices should be in `[-100, 0, ...,1154 config.vocab_size - 1]`. All labels set to `-100` are ignored (masked), the loss is only computed for1155 labels in `[0, ..., config.vocab_size]`1156 """1157 use_cache = use_cache if use_cache is not None else self.config.use_cache1158 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1159 1160 if inputs_embeds is not None and input_features is not None:1161 raise ValueError("Both `inputs_embeds` and `input_features` received! Please provide only one of them")1162 elif input_features is not None and inputs_embeds is None:1163 inputs_embeds = input_features1164 1165 # Encode if needed (training, first prediction pass)1166 if encoder_outputs is None:1167 # Convert encoder inputs in embeddings if needed1168 encoder_outputs = self.encoder(1169 input_ids=input_ids,1170 attention_mask=attention_mask,1171 inputs_embeds=inputs_embeds,1172 head_mask=head_mask,1173 output_attentions=output_attentions,1174 output_hidden_states=output_hidden_states,1175 return_dict=return_dict,1176 )1177 elif return_dict and not isinstance(encoder_outputs, BaseModelOutput):1178 encoder_outputs = BaseModelOutput(1179 last_hidden_state=encoder_outputs[0],1180 hidden_states=encoder_outputs[1] if len(encoder_outputs) > 1 else None,1181 attentions=encoder_outputs[2] if len(encoder_outputs) > 2 else None,1182 )1183 1184 hidden_states = encoder_outputs[0]1185 1186 if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:1187 # get decoder inputs from shifting lm labels to the right1188 decoder_input_ids = self._shift_right(labels)1189 1190 # Decode1191 decoder_outputs = self.decoder(1192 input_ids=decoder_input_ids,1193 attention_mask=decoder_attention_mask,1194 inputs_embeds=decoder_inputs_embeds,1195 past_key_values=past_key_values,1196 encoder_hidden_states=hidden_states,1197 encoder_attention_mask=attention_mask,1198 head_mask=decoder_head_mask,1199 cross_attn_head_mask=cross_attn_head_mask,1200 use_cache=use_cache,