Aluode/PerceptionLabPortable
0
1# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2# This file was automatically generated from src/transformers/models/t5gemma/modular_t5gemma.py.3# Do NOT edit this file manually as any edits will be overwritten by the generation of4# the file from the modular. If any change should be done, please apply the change to the5# modular_t5gemma.py file directly. One of our CI enforces this.6# 🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2025 Google Inc. HuggingFace Inc. team. All rights reserved.9#10#11# Licensed under the Apache License, Version 2.0 (the "License");12# you may not use this file except in compliance with the License.13# You may obtain a copy of the License at14#15# http://www.apache.org/licenses/LICENSE-2.016#17# Unless required by applicable law or agreed to in writing, software18# distributed under the License is distributed on an "AS IS" BASIS,19# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.20# See the License for the specific language governing permissions and21# limitations under the License.22from typing import Callable, Optional, Union23 24import torch25import torch.nn as nn26 27from ...activations import ACT2FN28from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache29from ...generation import GenerationMixin30from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask31from ...modeling_flash_attention_utils import FlashAttentionKwargs32from ...modeling_layers import GradientCheckpointingLayer33from ...modeling_outputs import (34 BaseModelOutput,35 BaseModelOutputWithPastAndCrossAttentions,36 Seq2SeqLMOutput,37 Seq2SeqModelOutput,38 SequenceClassifierOutput,39 TokenClassifierOutput,40)41from ...modeling_rope_utils import ROPE_INIT_FUNCTIONS, dynamic_rope_update42from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel43from ...processing_utils import Unpack44from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging45from ...utils.deprecation import deprecate_kwarg46from ...utils.generic import OutputRecorder, check_model_inputs47from .configuration_t5gemma import T5GemmaConfig, T5GemmaModuleConfig48 49 50logger = logging.get_logger(__name__)51 52 53class T5GemmaRMSNorm(nn.Module):54 def __init__(self, dim: int, eps: float = 1e-6):55 super().__init__()56 self.eps = eps57 self.weight = nn.Parameter(torch.zeros(dim))58 59 def _norm(self, x):60 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)61 62 def forward(self, x):63 output = self._norm(x.float())64 # Llama does x.to(float16) * w whilst T5Gemma is (x * w).to(float16)65 # See https://github.com/huggingface/transformers/pull/2940266 output = output * (1.0 + self.weight.float())67 return output.type_as(x)68 69 def extra_repr(self):70 return f"{tuple(self.weight.shape)}, eps={self.eps}"71 72 73class T5GemmaMLP(nn.Module):74 def __init__(self, config):75 super().__init__()76 self.config = config77 self.hidden_size = config.hidden_size78 self.intermediate_size = config.intermediate_size79 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)80 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)81 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)82 self.act_fn = ACT2FN[config.hidden_activation]83 self.dropout = nn.Dropout(config.dropout_rate)84 85 def forward(self, x):86 hidden_states = self.act_fn(self.gate_proj(x)) * self.up_proj(x)87 hidden_states = self.dropout(hidden_states)88 down_proj = self.down_proj(hidden_states)89 return down_proj90 91 92class T5GemmaRotaryEmbedding(nn.Module):93 inv_freq: torch.Tensor # fix linting for `register_buffer`94 95 def __init__(self, config, device=None):96 super().__init__()97 # BC: "rope_type" was originally "type"98 if hasattr(config, "rope_scaling") and isinstance(config.rope_scaling, dict):99 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))100 else:101 self.rope_type = "default"102 self.max_seq_len_cached = config.max_position_embeddings103 self.original_max_seq_len = config.max_position_embeddings104 105 self.config = config106 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]107 108 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)109 self.register_buffer("inv_freq", inv_freq, persistent=False)110 self.original_inv_freq = self.inv_freq111 112 @torch.no_grad()113 @dynamic_rope_update # power user: used with advanced RoPE types (e.g. dynamic rope)114 def forward(self, x, position_ids):115 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1).to(x.device)116 position_ids_expanded = position_ids[:, None, :].float()117 118 device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu"119 with torch.autocast(device_type=device_type, enabled=False): # Force float32120 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)121 emb = torch.cat((freqs, freqs), dim=-1)122 cos = emb.cos() * self.attention_scaling123 sin = emb.sin() * self.attention_scaling124 125 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)126 127 128def rotate_half(x):129 """Rotates half the hidden dims of the input."""130 x1 = x[..., : x.shape[-1] // 2]131 x2 = x[..., x.shape[-1] // 2 :]132 return torch.cat((-x2, x1), dim=-1)133 134 135def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):136 """Applies Rotary Position Embedding to the query and key tensors.137 138 Args:139 q (`torch.Tensor`): The query tensor.140 k (`torch.Tensor`): The key tensor.141 cos (`torch.Tensor`): The cosine part of the rotary embedding.142 sin (`torch.Tensor`): The sine part of the rotary embedding.143 position_ids (`torch.Tensor`, *optional*):144 Deprecated and unused.145 unsqueeze_dim (`int`, *optional*, defaults to 1):146 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and147 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note148 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and149 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes150 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have151 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.152 Returns:153 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.154 """155 cos = cos.unsqueeze(unsqueeze_dim)156 sin = sin.unsqueeze(unsqueeze_dim)157 q_embed = (q * cos) + (rotate_half(q) * sin)158 k_embed = (k * cos) + (rotate_half(k) * sin)159 return q_embed, k_embed160 161 162def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:163 """164 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,165 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)166 """167 batch, num_key_value_heads, slen, head_dim = hidden_states.shape168 if n_rep == 1:169 return hidden_states170 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)171 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)172 173 174def eager_attention_forward(175 module: nn.Module,176 query: torch.Tensor,177 key: torch.Tensor,178 value: torch.Tensor,179 attention_mask: Optional[torch.Tensor],180 dropout: float = 0.0,181 scaling: Optional[float] = None,182 softcap: Optional[float] = None,183 **kwargs,184) -> tuple[torch.Tensor, torch.Tensor]:185 if scaling is None:186 scaling = module.head_dim**-0.5187 188 key_states = repeat_kv(key, module.num_key_value_groups)189 value_states = repeat_kv(value, module.num_key_value_groups)190 191 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling192 193 if softcap is not None:194 attn_weights = attn_weights / softcap195 attn_weights = torch.tanh(attn_weights)196 attn_weights = attn_weights * softcap197 if attention_mask is not None: # no matter the length, we just slice it198 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]199 attn_weights = attn_weights + causal_mask200 201 # upcast attention to fp32202 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)203 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)204 attn_output = torch.matmul(attn_weights, value_states)205 attn_output = attn_output.transpose(1, 2).contiguous()206 return attn_output, attn_weights207 208 209class T5GemmaSelfAttention(nn.Module):210 """Multi-headed attention from 'Attention Is All You Need' paper"""211 212 def __init__(self, config: T5GemmaModuleConfig, layer_idx: int):213 super().__init__()214 self.config = config215 self.layer_idx = layer_idx216 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)217 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads218 self.scaling = config.query_pre_attn_scalar**-0.5219 self.attention_dropout = self.config.attention_dropout220 # Required by flash attention: encoder selfattention is non-causal221 self.is_causal = config.is_decoder222 223 self.q_proj = nn.Linear(224 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias225 )226 self.k_proj = nn.Linear(227 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias228 )229 self.v_proj = nn.Linear(230 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias231 )232 self.o_proj = nn.Linear(233 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias234 )235 self.attn_logit_softcapping = self.config.attn_logit_softcapping236 self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None237 238 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")239 def forward(240 self,241 hidden_states: torch.Tensor,242 position_embeddings: tuple[torch.Tensor, torch.Tensor],243 attention_mask: Optional[torch.Tensor],244 past_key_values: Optional[Cache] = None,245 cache_position: Optional[torch.LongTensor] = None,246 **kwargs: Unpack[FlashAttentionKwargs],247 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:248 input_shape = hidden_states.shape[:-1]249 hidden_shape = (*input_shape, -1, self.head_dim)250 251 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)252 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)253 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)254 255 cos, sin = position_embeddings256 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)257 258 if past_key_values is not None:259 # sin and cos are specific to RoPE models; cache_position needed for the static cache260 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}261 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)262 263 attention_interface: Callable = eager_attention_forward264 if self.config._attn_implementation != "eager":265 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]266 267 attn_output, attn_weights = attention_interface(268 self,269 query_states,270 key_states,271 value_states,272 attention_mask,273 dropout=self.attention_dropout if self.training else 0.0,274 scaling=self.scaling,275 sliding_window=self.sliding_window,276 softcap=self.attn_logit_softcapping,277 **kwargs,278 )279 280 attn_output = attn_output.reshape(*input_shape, -1).contiguous()281 attn_output = self.o_proj(attn_output)282 return attn_output, attn_weights283 284 285class T5GemmaCrossAttention(nn.Module):286 """Multi-headed attention from 'Attention Is All You Need' paper"""287 288 def __init__(self, config: T5GemmaModuleConfig, layer_idx: int):289 super().__init__()290 self.config = config291 self.layer_idx = layer_idx292 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)293 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads294 self.scaling = config.query_pre_attn_scalar**-0.5295 self.attention_dropout = self.config.attention_dropout296 self.is_causal = False297 298 self.q_proj = nn.Linear(299 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias300 )301 302 self.k_proj = nn.Linear(303 config.cross_attention_hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias304 )305 self.v_proj = nn.Linear(306 config.cross_attention_hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias307 )308 self.o_proj = nn.Linear(309 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias310 )311 self.attn_logit_softcapping = self.config.attn_logit_softcapping312 313 if config.cross_attention_hidden_size is None:314 raise ValueError("Cross-attention needs cross_attention_hidden_size to be specified.")315 316 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")317 def forward(318 self,319 hidden_states: torch.Tensor,320 attention_mask: Optional[torch.Tensor],321 encoder_hidden_states: Optional[torch.Tensor],322 past_key_values: Optional[Cache] = None,323 **kwargs: Unpack[FlashAttentionKwargs],324 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:325 if encoder_hidden_states is None:326 raise ValueError("Encoder hidden state is required for cross attention.")327 328 input_shape = hidden_states.shape[:-1]329 hidden_shape = (*input_shape, -1, self.head_dim)330 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)331 332 if past_key_values is not None:333 is_updated = past_key_values.is_updated.get(self.layer_idx)334 curr_past_key_value = past_key_values.cross_attention_cache335 336 if past_key_values is None or not is_updated:337 encoder_input_shape = encoder_hidden_states.shape[:-1]338 encoder_hidden_shape = (*encoder_input_shape, -1, self.head_dim)339 key_states = self.k_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)340 value_states = self.v_proj(encoder_hidden_states).view(encoder_hidden_shape).transpose(1, 2)341 342 if past_key_values is not None:343 key_states, value_states = curr_past_key_value.update(key_states, value_states, self.layer_idx)344 past_key_values.is_updated[self.layer_idx] = True345 else:346 key_states = curr_past_key_value.layers[self.layer_idx].keys347 value_states = curr_past_key_value.layers[self.layer_idx].values348 349 attention_interface: Callable = eager_attention_forward350 if self.config._attn_implementation != "eager":351 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]352 353 attn_output, attn_weights = attention_interface(354 self,355 query_states,356 key_states,357 value_states,358 attention_mask,359 dropout=self.attention_dropout if self.training else 0.0,360 scaling=self.scaling,361 sliding_window=None,362 softcap=self.attn_logit_softcapping,363 **kwargs,364 )365 366 attn_output = attn_output.reshape(*input_shape, -1).contiguous()367 attn_output = self.o_proj(attn_output)368 return attn_output, attn_weights369 370 371class T5GemmaEncoderLayer(GradientCheckpointingLayer):372 """Encoder sub-layer."""373 374 def __init__(self, config, layer_idx: int):375 super().__init__()376 self.hidden_size = config.hidden_size377 self.config = config378 self.layer_idx = layer_idx379 self.attention_type = config.layer_types[layer_idx]380 381 self.self_attn = T5GemmaSelfAttention(382 config=config,383 layer_idx=layer_idx,384 )385 self.pre_self_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)386 self.post_self_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)387 388 self.mlp = T5GemmaMLP(config)389 self.pre_feedforward_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)390 self.post_feedforward_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)391 392 self.dropout = nn.Dropout(config.dropout_rate)393 394 def forward(395 self,396 hidden_states: torch.Tensor,397 position_embeddings: tuple[torch.Tensor, torch.Tensor],398 attention_mask: Optional[torch.Tensor] = None,399 position_ids: Optional[torch.LongTensor] = None,400 **kwargs,401 ) -> tuple[torch.FloatTensor,]:402 residual = hidden_states403 hidden_states = self.pre_self_attn_layernorm(hidden_states)404 hidden_states, _ = self.self_attn(405 hidden_states=hidden_states,406 position_embeddings=position_embeddings,407 attention_mask=attention_mask,408 position_ids=position_ids,409 past_key_values=None,410 **kwargs,411 )412 hidden_states = self.post_self_attn_layernorm(hidden_states)413 hidden_states = residual + self.dropout(hidden_states)414 415 residual = hidden_states416 hidden_states = self.pre_feedforward_layernorm(hidden_states)417 hidden_states = self.mlp(hidden_states)418 hidden_states = self.post_feedforward_layernorm(hidden_states)419 hidden_states = residual + self.dropout(hidden_states)420 return hidden_states421 422 423class T5GemmaDecoderLayer(T5GemmaEncoderLayer):424 """Decoder sub-layer: an extra cross-attention layer."""425 426 def __init__(self, config, layer_idx: int):427 super().__init__(config, layer_idx)428 self.cross_attn = T5GemmaCrossAttention(config=config, layer_idx=layer_idx)429 self.pre_cross_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)430 self.post_cross_attn_layernorm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)431 432 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")433 def forward(434 self,435 hidden_states: torch.Tensor,436 position_embeddings: tuple[torch.Tensor, torch.Tensor],437 attention_mask: Optional[torch.Tensor] = None,438 position_ids: Optional[torch.LongTensor] = None,439 past_key_values: Optional[EncoderDecoderCache] = None,440 use_cache: Optional[bool] = False,441 cache_position: Optional[torch.LongTensor] = None,442 encoder_hidden_states: Optional[torch.Tensor] = None,443 encoder_attention_mask: Optional[torch.Tensor] = None,444 **kwargs,445 ) -> torch.FloatTensor:446 residual = hidden_states447 hidden_states = self.pre_self_attn_layernorm(hidden_states)448 hidden_states, _ = self.self_attn(449 hidden_states=hidden_states,450 position_embeddings=position_embeddings,451 attention_mask=attention_mask,452 position_ids=position_ids,453 past_key_values=past_key_values.self_attention_cache if past_key_values is not None else None,454 use_cache=use_cache,455 cache_position=cache_position,456 **kwargs,457 )458 hidden_states = self.post_self_attn_layernorm(hidden_states)459 hidden_states = residual + self.dropout(hidden_states)460 461 residual = hidden_states462 hidden_states = self.pre_cross_attn_layernorm(hidden_states)463 hidden_states, _ = self.cross_attn(464 hidden_states=hidden_states,465 encoder_hidden_states=encoder_hidden_states,466 attention_mask=encoder_attention_mask,467 past_key_values=past_key_values,468 use_cache=use_cache,469 **kwargs,470 )471 hidden_states = self.post_cross_attn_layernorm(hidden_states)472 hidden_states = residual + self.dropout(hidden_states)473 474 residual = hidden_states475 hidden_states = self.pre_feedforward_layernorm(hidden_states)476 hidden_states = self.mlp(hidden_states)477 hidden_states = self.post_feedforward_layernorm(hidden_states)478 hidden_states = residual + self.dropout(hidden_states)479 return hidden_states480 481 482class T5GemmaClassificationHead(nn.Module):483 """Head for sentence-level classification tasks."""484 485 def __init__(self, hidden_size: int, num_labels: int, classifier_dropout_rate: float = 0.0):486 super().__init__()487 self.dropout = nn.Dropout(p=classifier_dropout_rate)488 self.out_proj = nn.Linear(hidden_size, num_labels)489 490 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:491 hidden_states = self.dropout(hidden_states)492 hidden_states = self.out_proj(hidden_states)493 return hidden_states494 495 496class T5GemmaLMHead(nn.Module):497 """Head for language modeling (generation) tasks."""498 499 def __init__(self, hidden_size: int, vocab_size: int, bias: bool = False):500 super().__init__()501 self.out_proj = nn.Linear(hidden_size, vocab_size, bias=bias)502 503 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:504 logits = self.out_proj(hidden_states)505 return logits506 507 508class T5GemmaAttention(nn.Module):509 """Multi-headed attention from 'Attention Is All You Need' paper"""510 511 def __init__(self, config: T5GemmaConfig, layer_idx: int):512 super().__init__()513 self.config = config514 self.layer_idx = layer_idx515 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)516 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads517 self.scaling = config.query_pre_attn_scalar**-0.5518 self.attention_dropout = self.config.attention_dropout519 self.is_causal = True520 521 self.q_proj = nn.Linear(522 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias523 )524 self.k_proj = nn.Linear(525 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias526 )527 self.v_proj = nn.Linear(528 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias529 )530 self.o_proj = nn.Linear(531 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias532 )533 self.attn_logit_softcapping = self.config.attn_logit_softcapping534 self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None535 536 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")537 def forward(538 self,539 hidden_states: torch.Tensor,540 position_embeddings: tuple[torch.Tensor, torch.Tensor],541 attention_mask: Optional[torch.Tensor],542 past_key_values: Optional[Cache] = None,543 cache_position: Optional[torch.LongTensor] = None,544 **kwargs: Unpack[FlashAttentionKwargs],545 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:546 input_shape = hidden_states.shape[:-1]547 hidden_shape = (*input_shape, -1, self.head_dim)548 549 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)550 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)551 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)552 553 cos, sin = position_embeddings554 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)555 556 if past_key_values is not None:557 # sin and cos are specific to RoPE models; cache_position needed for the static cache558 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}559 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)560 561 attention_interface: Callable = eager_attention_forward562 if self.config._attn_implementation != "eager":563 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]564 565 attn_output, attn_weights = attention_interface(566 self,567 query_states,568 key_states,569 value_states,570 attention_mask,571 dropout=self.attention_dropout if self.training else 0.0,572 scaling=self.scaling,573 sliding_window=self.sliding_window,574 softcap=self.attn_logit_softcapping,575 **kwargs,576 )577 578 attn_output = attn_output.reshape(*input_shape, -1).contiguous()579 attn_output = self.o_proj(attn_output)580 return attn_output, attn_weights581 582 583@auto_docstring584class T5GemmaPreTrainedModel(PreTrainedModel):585 config: T5GemmaConfig586 base_model_prefix = "model"587 supports_gradient_checkpointing = True588 _no_split_modules = ["T5GemmaEncoderLayer", "T5GemmaDecoderLayer"]589 _skip_keys_device_placement = ["past_key_values"]590 _supports_flash_attn = True591 _supports_sdpa = True592 _supports_flex_attn = True593 594 _can_compile_fullgraph = True595 _supports_attention_backend = True596 _can_record_outputs = {597 "hidden_states": T5GemmaDecoderLayer,598 "attentions": T5GemmaAttention,599 }600 601 def _init_weights(self, module):602 # TODO: support initialization for encoders and decoders separately(?)603 super()._init_weights(module)604 std = self.config.initializer_range605 if isinstance(module, T5GemmaClassificationHead):606 scale = module.out_proj.weight.shape[0] ** -0.5607 module.out_proj.weight.data.normal_(mean=0.0, std=std * scale)608 if hasattr(module.out_proj, "bias") and module.out_proj.bias is not None:609 module.out_proj.bias.data.zero_()610 elif isinstance(module, T5GemmaLMHead):611 if not self.config.tie_word_embeddings:612 scale = module.out_proj.weight.shape[0] ** -0.5613 module.out_proj.weight.data.normal_(mean=0.0, std=std * scale)614 # We initialize with 0s to be 1 centered as the RMSNorm here does (1 + weight)615 elif "RMSNorm" in module.__class__.__name__:616 module.weight.data.zero_()617 618 def _shift_right(self, input_ids):619 """620 Shifts input_ids to the right, prepends the decoder_start_token_id, and handles621 pad_token_id replacement for labels that were -100.622 This is a common preparation step for decoder inputs in sequence-to-sequence models.623 """624 decoder_start_token_id = self.config.decoder.bos_token_id625 pad_token_id = self.config.decoder.pad_token_id626 627 if decoder_start_token_id is None:628 raise ValueError("self.model.config.decoder.bos_token_id has to be defined. ")629 630 # shift inputs to the right631 shifted_input_ids = input_ids.new_zeros(input_ids.shape)632 shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()633 shifted_input_ids[..., 0] = decoder_start_token_id634 635 if pad_token_id is None:636 raise ValueError("self.model.config.decoder.pad_token_id has to be defined.")637 638 # Is this T5 specific?639 # replace possible -100 values in labels by `pad_token_id`640 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)641 642 return shifted_input_ids643 644 645def bidirectional_mask_function(attention_mask: Optional[torch.Tensor]) -> Callable:646 """647 This creates bidirectional attention mask.648 """649 650 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:651 if attention_mask is None:652 return torch.ones((), dtype=torch.bool)653 return attention_mask[batch_idx, kv_idx].to(torch.bool)654 655 return inner_mask656 657 658def sliding_window_bidirectional_mask_function(sliding_window: int) -> Callable:659 """660 This creates bidirectional attention mask with sliding window.661 """662 663 def inner_mask(batch_idx: int, head_idx: int, q_idx: int, kv_idx: int) -> bool:664 return (q_idx - sliding_window < kv_idx) & (kv_idx < q_idx + sliding_window)665 666 return inner_mask667 668 669def make_default_2d_attention_mask(670 token_ids: Optional[torch.LongTensor],671 hidden_states: torch.Tensor,672 pad_token_id: Optional[int],673) -> torch.Tensor:674 """Construct the default attention mask."""675 if token_ids is not None:676 if pad_token_id is None:677 raise ValueError("`pad_token_id` is required for padding information.")678 attention_mask = (token_ids != pad_token_id).to(hidden_states.device, torch.long)679 else:680 attention_mask = torch.ones(681 (hidden_states.shape[0], hidden_states.shape[1]), device=hidden_states.device, dtype=torch.long682 )683 return attention_mask684 685 686class T5GemmaEncoder(T5GemmaPreTrainedModel):687 _can_record_outputs = {688 "attentions": T5GemmaSelfAttention,689 "hidden_states": T5GemmaEncoderLayer,690 }691 692 def __init__(self, config):693 super().__init__(config)694 self.padding_idx = config.pad_token_id695 self.vocab_size = config.vocab_size696 697 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)698 self.norm = T5GemmaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)699 self.rotary_emb = T5GemmaRotaryEmbedding(config=config)700 self.gradient_checkpointing = False701 702 self.layers = nn.ModuleList(703 [T5GemmaEncoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]704 )705 self.dropout = nn.Dropout(config.dropout_rate)706 707 # Initialize weights and apply final processing708 self.post_init()709 710 @check_model_inputs()711 def forward(712 self,713 input_ids: Optional[torch.LongTensor] = None,714 attention_mask: Optional[torch.Tensor] = None,715 position_ids: Optional[torch.LongTensor] = None,716 inputs_embeds: Optional[torch.FloatTensor] = None,717 **kwargs: Unpack[TransformersKwargs],718 ) -> BaseModelOutput:719 if (input_ids is None) ^ (inputs_embeds is not None):720 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")721 722 # As we want to pass `past_key_values=None` explicitly everywhere, we need to pop them from kwargs if present723 kwargs.pop("past_key_values", None)724 725 if inputs_embeds is None:726 inputs_embeds = self.embed_tokens(input_ids)727 728 cache_position = torch.arange(0, inputs_embeds.shape[1], device=inputs_embeds.device)729 730 if position_ids is None:731 position_ids = cache_position.unsqueeze(0)732 733 if attention_mask is None:734 attention_mask = make_default_2d_attention_mask(input_ids, inputs_embeds, self.config.pad_token_id)735 736 if not isinstance(self_attn_mask_mapping := attention_mask, dict):737 mask_kwargs = {738 "config": self.config,739 "input_embeds": inputs_embeds,740 "attention_mask": attention_mask,741 "cache_position": cache_position,742 "past_key_values": None,743 "position_ids": position_ids,744 }745 self_attn_mask_mapping = {746 "full_attention": create_causal_mask(747 **mask_kwargs,748 or_mask_function=bidirectional_mask_function(attention_mask),749 ),750 "sliding_attention": create_sliding_window_causal_mask(751 **mask_kwargs,752 or_mask_function=sliding_window_bidirectional_mask_function(self.config.sliding_window),753 and_mask_function=bidirectional_mask_function(attention_mask),754 ),755 }756 757 hidden_states = inputs_embeds758 position_embeddings = self.rotary_emb(hidden_states, position_ids)759 760 normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)761 hidden_states = hidden_states * normalizer762 hidden_states = self.dropout(hidden_states)763 764 for layer_module in self.layers[: self.config.num_hidden_layers]:765 hidden_states = layer_module(766 hidden_states,767 position_embeddings,768 self_attn_mask_mapping[layer_module.attention_type],769 position_ids,770 **kwargs,771 )772 hidden_states = self.norm(hidden_states)773 hidden_states = self.dropout(hidden_states)774 return BaseModelOutput(775 last_hidden_state=hidden_states,776 )777 778 779class T5GemmaDecoder(T5GemmaEncoder):780 _can_record_outputs = {781 "attentions": OutputRecorder(T5GemmaSelfAttention, index=1),782 "cross_attentions": OutputRecorder(T5GemmaCrossAttention, index=1),783 "hidden_states": T5GemmaDecoderLayer,784 }785 786 def __init__(self, config):787 super().__init__(config)788 self.layers = nn.ModuleList(789 [T5GemmaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]790 )791 792 self.post_init()793 794 @check_model_inputs()795 def forward(796 self,797 input_ids: Optional[torch.LongTensor] = None,798 attention_mask: Optional[torch.Tensor] = None,799 position_ids: Optional[torch.LongTensor] = None,800 past_key_values: Optional[EncoderDecoderCache] = None,801 inputs_embeds: Optional[torch.FloatTensor] = None,802 use_cache: Optional[bool] = None,803 cache_position: Optional[torch.LongTensor] = None,804 encoder_hidden_states: Optional[torch.Tensor] = None,805 encoder_attention_mask: Optional[torch.Tensor] = None,806 **kwargs: Unpack[TransformersKwargs],807 ) -> BaseModelOutputWithPastAndCrossAttentions:808 if (input_ids is None) ^ (inputs_embeds is not None):809 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")810 if encoder_hidden_states is None:811 raise ValueError("`encoder_hidden_states` must be given in decoder")812 813 if inputs_embeds is None:814 inputs_embeds = self.embed_tokens(input_ids)815 816 if not self.training and use_cache and past_key_values is None:817 past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))818 if cache_position is None:819 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0820 cache_position = torch.arange(821 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device822 )823 824 if position_ids is None:825 position_ids = cache_position.unsqueeze(0)826 827 if attention_mask is None and past_key_values is None:828 attention_mask = make_default_2d_attention_mask(input_ids, inputs_embeds, self.config.pad_token_id)829 830 if not isinstance(self_attn_mask_mapping := attention_mask, dict):831 mask_kwargs = {832 "config": self.config,833 "input_embeds": inputs_embeds,834 "attention_mask": attention_mask,835 "cache_position": cache_position,836 "past_key_values": past_key_values.self_attention_cache if past_key_values is not None else None,837 "position_ids": position_ids,838 }839 self_attn_mask_mapping = {840 "full_attention": create_causal_mask(**mask_kwargs),841 "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),842 }843 844 if not isinstance(cross_attn_mask_mapping := encoder_attention_mask, dict):845 mask_kwargs = {846 "config": self.config,847 "input_embeds": encoder_hidden_states,848 "attention_mask": encoder_attention_mask,849 "cache_position": cache_position,850 "past_key_values": None,851 "position_ids": None,852 }853 cross_attn_mask_mapping = {854 "full_attention": create_causal_mask(855 **mask_kwargs,856 or_mask_function=bidirectional_mask_function(encoder_attention_mask),857 ),858 }859 860 hidden_states = inputs_embeds861 position_embeddings = self.rotary_emb(hidden_states, position_ids)862 863 normalizer = torch.tensor(self.config.hidden_size**0.5, dtype=hidden_states.dtype)864 hidden_states = hidden_states * normalizer865 hidden_states = self.dropout(hidden_states)866 867 for layer_module in self.layers[: self.config.num_hidden_layers]:868 hidden_states = layer_module(869 hidden_states,870 position_embeddings,871 self_attn_mask_mapping[layer_module.attention_type],872 position_ids,873 past_key_values,874 use_cache,875 cache_position,876 encoder_hidden_states,877 cross_attn_mask_mapping["full_attention"],878 **kwargs,879 )880 hidden_states = self.norm(hidden_states)881 hidden_states = self.dropout(hidden_states)882 return BaseModelOutputWithPastAndCrossAttentions(883 last_hidden_state=hidden_states,884 past_key_values=past_key_values,885 )886 887 888@auto_docstring889class T5GemmaModel(T5GemmaPreTrainedModel):890 def __init__(self, config: T5GemmaConfig):891 super().__init__(config)892 893 if not config.is_encoder_decoder:894 raise ValueError("T5GemmaModel only support encoder-decoder modeling. Use `T5GemmaEncoderModel` instead.")895 896 self.encoder = T5GemmaEncoder(config.encoder)897 self.decoder = T5GemmaDecoder(config.decoder)898 899 self.post_init()900 901 def get_encoder(self):902 return self.encoder903 904 def get_input_embeddings(self):905 return self.encoder.get_input_embeddings()906 907 def set_input_embeddings(self, new_embeddings):908 return self.encoder.set_input_embeddings(new_embeddings)909 910 @can_return_tuple911 @auto_docstring912 def forward(913 self,914 input_ids: Optional[torch.LongTensor] = None,915 attention_mask: Optional[torch.FloatTensor] = None,916 position_ids: Optional[torch.LongTensor] = None,917 decoder_input_ids: Optional[torch.LongTensor] = None,918 decoder_attention_mask: Optional[torch.BoolTensor] = None,919 decoder_position_ids: Optional[torch.LongTensor] = None,920 encoder_outputs: Optional[BaseModelOutput] = None,921 past_key_values: Optional[EncoderDecoderCache] = None,922 inputs_embeds: Optional[torch.Tensor] = None,923 decoder_inputs_embeds: Optional[torch.Tensor] = None,924 use_cache: Optional[bool] = None,925 cache_position: Optional[torch.LongTensor] = None,926 **kwargs: Unpack[TransformersKwargs],927 ) -> Seq2SeqModelOutput:928 r"""929 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, decoder_sequence_length)`, *optional*):930 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the range `[0,931 config.decoder.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)932 """933 if encoder_outputs is None:934 encoder_outputs = self.encoder(935 input_ids=input_ids,936 attention_mask=attention_mask,937 position_ids=position_ids,938 inputs_embeds=inputs_embeds,939 **kwargs,940 )941 942 encoder_hidden_states = encoder_outputs.last_hidden_state943 944 decoder_outputs = self.decoder(945 input_ids=decoder_input_ids,946 attention_mask=decoder_attention_mask,947 position_ids=decoder_position_ids,948 inputs_embeds=decoder_inputs_embeds,949 past_key_values=past_key_values,950 encoder_hidden_states=encoder_hidden_states,951 encoder_attention_mask=attention_mask,952 use_cache=use_cache,953 cache_position=cache_position,954 **kwargs,955 )956 957 return Seq2SeqModelOutput(958 last_hidden_state=decoder_outputs.last_hidden_state,959 past_key_values=decoder_outputs.past_key_values,960 decoder_hidden_states=decoder_outputs.hidden_states961 if kwargs.get("output_hidden_states", False)962 else (decoder_outputs.last_hidden_state,),963 decoder_attentions=decoder_outputs.attentions,964 cross_attentions=decoder_outputs.cross_attentions,965 encoder_last_hidden_state=encoder_outputs.last_hidden_state,966 encoder_hidden_states=encoder_outputs.hidden_states,967 encoder_attentions=encoder_outputs.attentions,968 )969 970 971@auto_docstring972class T5GemmaEncoderModel(T5GemmaPreTrainedModel):973 def __init__(self, config: T5GemmaConfig):974 super().__init__(config)975 976 if config.is_encoder_decoder:977 raise ValueError("T5GemmaEncoderModel only supports encoder-only model. Use `T5GemmaModel` instead.")978 979 self.encoder = T5GemmaEncoder(config.encoder)980 self.post_init()981 982 def get_input_embeddings(self):983 return self.encoder.get_input_embeddings()984 985 def set_input_embeddings(self, new_embeddings):986 return self.encoder.set_input_embeddings(new_embeddings)987 988 @can_return_tuple989 @auto_docstring990 def forward(991 self,992 input_ids: Optional[torch.LongTensor] = None,993 attention_mask: Optional[torch.FloatTensor] = None,994 position_ids: Optional[torch.LongTensor] = None,995 inputs_embeds: Optional[torch.Tensor] = None,996 **kwargs: Unpack[TransformersKwargs],997 ) -> BaseModelOutput:998 encoder_outputs = self.encoder(999 input_ids=input_ids,1000 attention_mask=attention_mask,1001 position_ids=position_ids,1002 inputs_embeds=inputs_embeds,1003 **kwargs,1004 )1005 return encoder_outputs1006 1007 1008class T5GemmaForConditionalGeneration(T5GemmaPreTrainedModel, GenerationMixin):1009 _tied_weights_keys = ["model.decoder.embed_tokens.weight", "lm_head.out_proj.weight"]1010 _tp_plan = {"lm_head.out_proj": "colwise_rep"}1011 _pp_plan = {"lm_head.out_proj": (["hidden_states"], ["logits"])}1012 1013 def __init__(self, config: T5GemmaConfig):1014 config.is_encoder_decoder = True1015 super().__init__(config)1016 1017 self.model = T5GemmaModel(config)1018 self.vocab_size = config.decoder.vocab_size1019 self.lm_head = T5GemmaLMHead(config.decoder.hidden_size, self.vocab_size)1020 self.loss_type = "ForMaskedLM"1021 1022 self.post_init()1023 1024 def set_output_embeddings(self, new_embeddings):1025 self.lm_head.out_proj = new_embeddings1026 1027 def get_output_embeddings(self):1028 return self.lm_head.out_proj1029 1030 def _tie_weights(self):1031 # Decoder input and output embeddings are tied.1032 if self.config.tie_word_embeddings:1033 self._tie_or_clone_weights(self.lm_head.out_proj, self.get_decoder().get_input_embeddings())1034 1035 def get_encoder(self):1036 return self.model.encoder1037 1038 def get_decoder(self):1039 return self.model.decoder1040 1041 @can_return_tuple1042 @auto_docstring1043 def forward(1044 self,1045 input_ids: Optional[torch.LongTensor] = None,1046 attention_mask: Optional[torch.FloatTensor] = None,1047 position_ids: Optional[torch.LongTensor] = None,1048 decoder_input_ids: Optional[torch.LongTensor] = None,1049 decoder_attention_mask: Optional[torch.BoolTensor] = None,1050 decoder_position_ids: Optional[torch.LongTensor] = None,1051 encoder_outputs: Optional[BaseModelOutput] = None,1052 past_key_values: Optional[EncoderDecoderCache] = None,1053 inputs_embeds: Optional[torch.FloatTensor] = None,1054 decoder_inputs_embeds: Optional[torch.FloatTensor] = None,1055 labels: Optional[torch.LongTensor] = None,1056 use_cache: Optional[bool] = None,1057 cache_position: Optional[torch.LongTensor] = None,1058 logits_to_keep: Union[int, torch.Tensor] = 0,1059 **kwargs: Unpack[TransformersKwargs],1060 ) -> Union[tuple[torch.FloatTensor], Seq2SeqLMOutput]:1061 r"""1062 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, decoder_sequence_length)`, *optional*):1063 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the range `[0,1064 config.decoder.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)1065 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1066 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1067 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1068 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1069 """1070 1071 if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:1072 # get decoder inputs from shifting lm labels to the right1073 decoder_input_ids = self._shift_right(labels)1074 1075 decoder_outputs: Seq2SeqModelOutput = self.model(1076 input_ids=input_ids,1077 attention_mask=attention_mask,1078 position_ids=position_ids,1079 decoder_input_ids=decoder_input_ids,1080 decoder_attention_mask=decoder_attention_mask,1081 decoder_position_ids=decoder_position_ids,1082 encoder_outputs=encoder_outputs,1083 past_key_values=past_key_values,1084 inputs_embeds=inputs_embeds,1085 decoder_inputs_embeds=decoder_inputs_embeds,1086 use_cache=use_cache,1087 cache_position=cache_position,1088 **kwargs,1089 )1090 1091 hidden_states = decoder_outputs.last_hidden_state1092 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss1093 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep1094 logits = self.lm_head(hidden_states[:, slice_indices, :])1095 decoder_config = self.get_decoder().config1096 if decoder_config.final_logit_softcapping is not None:1097 logits = logits / decoder_config.final_logit_softcapping1098 logits = torch.tanh(logits)1099 logits = logits * decoder_config.final_logit_softcapping1100 1101 loss = None1102 if labels is not None:1103 # Input has right-shifted so we directly perform masked lm loss1104 loss = self.loss_function(logits, labels, self.vocab_size, **kwargs)1105 1106 return Seq2SeqLMOutput(1107 loss=loss,1108 logits=logits,1109 past_key_values=decoder_outputs.past_key_values,1110 decoder_hidden_states=decoder_outputs.decoder_hidden_states,1111 decoder_attentions=decoder_outputs.decoder_attentions,1112 cross_attentions=decoder_outputs.cross_attentions,1113 encoder_last_hidden_state=decoder_outputs.encoder_last_hidden_state,1114 encoder_hidden_states=decoder_outputs.encoder_hidden_states,1115 encoder_attentions=decoder_outputs.encoder_attentions,1116 )1117 1118 def prepare_decoder_input_ids_from_labels(self, labels: torch.Tensor):1119 return self._shift_right(labels)1120 1121 1122@auto_docstring1123class T5GemmaForSequenceClassification(T5GemmaPreTrainedModel):1124 def __init__(self, config: T5GemmaConfig, is_encoder_decoder: Optional[bool] = None):1125 r"""1126 is_encoder_decoder (`Optional`, *optional*):1127 Whether use encoder_decoder for sequence classification. When set to False, only encoder is used.1128 """1129 if is_encoder_decoder is not None:1130 config.is_encoder_decoder = is_encoder_decoder1131 super().__init__(config)1132 self.num_labels = config.num_labels1133 1134 if config.is_encoder_decoder:1135 self.model = T5GemmaModel(config)1136 else:1137 self.model = T5GemmaEncoderModel(config)1138 1139 hidden_size = config.encoder.hidden_size1140 if config.is_encoder_decoder:1141 hidden_size = config.decoder.hidden_size1142 1143 classifier_dropout = getattr(config, "classifier_dropout_rate", 0.1)1144 self.score = T5GemmaClassificationHead(hidden_size, self.num_labels, classifier_dropout)1145 self.post_init()1146 1147 def get_input_embeddings(self):1148 return self.model.get_input_embeddings()1149 1150 def set_input_embeddings(self, value):1151 self.model.set_input_embeddings(value)1152 1153 @can_return_tuple1154 @auto_docstring1155 def forward(1156 self,1157 input_ids: Optional[torch.LongTensor] = None,1158 attention_mask: Optional[torch.Tensor] = None,1159 position_ids: Optional[torch.LongTensor] = None,1160 decoder_input_ids: Optional[torch.LongTensor] = None,1161 decoder_attention_mask: Optional[torch.Tensor] = None,1162 decoder_position_ids: Optional[torch.LongTensor] = None,1163 encoder_outputs: Optional[BaseModelOutput] = None,1164 inputs_embeds: Optional[torch.FloatTensor] = None,1165 decoder_inputs_embeds: Optional[torch.FloatTensor] = None,1166 labels: Optional[torch.LongTensor] = None,1167 **kwargs: Unpack[TransformersKwargs],1168 ) -> SequenceClassifierOutput:1169 r"""1170 decoder_position_ids (`torch.LongTensor` of shape `(batch_size, decoder_sequence_length)`, *optional*):1171 Indices of positions of each decoder input sequence tokens in the position embeddings. Selected in the range `[0,1172 config.decoder.n_positions - 1]`. [What are position IDs?](../glossary#position-ids)1173 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1174 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1175 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1176 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1177 """1178 if self.config.is_encoder_decoder and (input_ids is None and inputs_embeds is not None):1179 raise NotImplementedError(1180 f"Passing input embeddings is currently not supported for {self.__class__.__name__} in encoder-decoder mode."1181 )1182 1183 # Following T5, we automatically creates decoder_input_ids from input_ids if no decoder_input_ids are provided1184 if self.config.is_encoder_decoder and (decoder_input_ids is None and decoder_inputs_embeds is None):1185 if input_ids is None:1186 raise ValueError(1187 "If no `decoder_input_ids` or `decoder_inputs_embeds` are "1188 "passed, `input_ids` cannot be `None`. Please pass either "1189 "`input_ids` or `decoder_input_ids` or `decoder_inputs_embeds`."1190 )1191 decoder_input_ids = self._shift_right(input_ids)1192 1193 if self.config.is_encoder_decoder:1194 outputs: Seq2SeqModelOutput = self.model(1195 input_ids,1196 attention_mask=attention_mask,1197 position_ids=position_ids,1198 decoder_input_ids=decoder_input_ids,1199 decoder_attention_mask=decoder_attention_mask,1200 decoder_position_ids=decoder_position_ids,