Aluode/PerceptionLabPortable
0
1# coding=utf-82# Copyright 2021 The Fairseq 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 BART model."""16 17import math18import warnings19from typing import Callable, Optional, Union20 21import torch22from torch import nn23from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss24 25from ...activations import ACT2FN26from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache27from ...generation import GenerationMixin28from ...modeling_attn_mask_utils import (29 AttentionMaskConverter,30 _prepare_4d_attention_mask,31 _prepare_4d_attention_mask_for_sdpa,32)33from ...modeling_flash_attention_utils import FlashAttentionKwargs34from ...modeling_layers import GradientCheckpointingLayer35from ...modeling_outputs import (36 BaseModelOutput,37 BaseModelOutputWithPastAndCrossAttentions,38 CausalLMOutputWithCrossAttentions,39 Seq2SeqLMOutput,40 Seq2SeqModelOutput,41 Seq2SeqQuestionAnsweringModelOutput,42 Seq2SeqSequenceClassifierOutput,43)44from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel45from ...processing_utils import Unpack46from ...utils import (47 auto_docstring,48 is_torch_flex_attn_available,49 is_torchdynamo_compiling,50 logging,51)52from ...utils.deprecation import deprecate_kwarg53from .configuration_bart import BartConfig54 55 56if is_torch_flex_attn_available():57 from ...integrations.flex_attention import BlockMask, make_flex_block_causal_mask58 59 60logger = logging.get_logger(__name__)61 62 63def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int, decoder_start_token_id: int):64 """65 Shift input ids one token to the right.66 """67 shifted_input_ids = input_ids.new_zeros(input_ids.shape)68 shifted_input_ids[:, 1:] = input_ids[:, :-1].clone()69 shifted_input_ids[:, 0] = decoder_start_token_id70 71 if pad_token_id is None:72 raise ValueError("self.model.config.pad_token_id has to be defined.")73 # replace possible -100 values in labels by `pad_token_id`74 shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)75 76 return shifted_input_ids77 78 79class BartLearnedPositionalEmbedding(nn.Embedding):80 """81 This module learns positional embeddings up to a fixed maximum size.82 """83 84 def __init__(self, num_embeddings: int, embedding_dim: int):85 # Bart is set up so that if padding_idx is specified then offset the embedding ids by 286 # and adjust num_embeddings appropriately. Other models don't have this hack87 self.offset = 288 super().__init__(num_embeddings + self.offset, embedding_dim)89 90 def forward(91 self, input_ids: torch.Tensor, past_key_values_length: int = 0, position_ids: Optional[torch.Tensor] = None92 ):93 """`input_ids' shape is expected to be [bsz x seqlen]."""94 95 if position_ids is None:96 bsz, seq_len = input_ids.shape[:2]97 position_ids = torch.arange(98 past_key_values_length, past_key_values_length + seq_len, dtype=torch.long, device=self.weight.device99 ).expand(bsz, -1)100 else:101 position_ids = position_ids.unsqueeze(0)102 103 return super().forward(position_ids + self.offset)104 105 106class BartScaledWordEmbedding(nn.Embedding):107 """108 This module overrides nn.Embeddings' forward by multiplying with embeddings scale.109 """110 111 def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0):112 super().__init__(num_embeddings, embedding_dim, padding_idx)113 self.embed_scale = embed_scale114 115 def forward(self, input_ids: torch.Tensor):116 return super().forward(input_ids) * self.embed_scale117 118 119def eager_attention_forward(120 module: nn.Module,121 query: torch.Tensor,122 key: torch.Tensor,123 value: torch.Tensor,124 attention_mask: Optional[torch.Tensor],125 scaling: Optional[float] = None,126 dropout: float = 0.0,127 head_mask: Optional[torch.Tensor] = None,128 **kwargs,129):130 if scaling is None:131 scaling = query.size(-1) ** -0.5132 133 attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling134 if attention_mask is not None:135 attn_weights = attn_weights + attention_mask136 137 attn_weights = nn.functional.softmax(attn_weights, dim=-1)138 139 if head_mask is not None:140 attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)141 142 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)143 attn_output = torch.matmul(attn_weights, value)144 attn_output = attn_output.transpose(1, 2).contiguous()145 146 return attn_output, attn_weights147 148 149class BartAttention(nn.Module):150 """Multi-headed attention from 'Attention Is All You Need' paper"""151 152 def __init__(153 self,154 embed_dim: int,155 num_heads: int,156 dropout: float = 0.0,157 is_decoder: bool = False,158 bias: bool = True,159 is_causal: bool = False,160 config: Optional[BartConfig] = None,161 layer_idx: Optional[int] = None,162 ):163 super().__init__()164 self.embed_dim = embed_dim165 self.num_heads = num_heads166 self.dropout = dropout167 self.head_dim = embed_dim // num_heads168 self.config = config169 170 if (self.head_dim * num_heads) != self.embed_dim:171 raise ValueError(172 f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"173 f" and `num_heads`: {num_heads})."174 )175 self.scaling = self.head_dim**-0.5176 self.is_decoder = is_decoder177 self.is_causal = is_causal178 self.layer_idx = layer_idx179 if layer_idx is None and self.is_decoder:180 logger.warning_once(181 f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "182 "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "183 "when creating this class."184 )185 186 self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)187 self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)188 self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)189 self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)190 191 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")192 def forward(193 self,194 hidden_states: torch.Tensor,195 key_value_states: Optional[torch.Tensor] = None,196 past_key_values: Optional[Cache] = None,197 attention_mask: Optional[torch.Tensor] = None,198 layer_head_mask: Optional[torch.Tensor] = None,199 output_attentions: bool = False,200 cache_position: Optional[torch.Tensor] = None,201 # TODO: we need a refactor so that the different attention modules can get their specific kwargs202 # ATM, we have mixed things encoder, decoder, and encoder-decoder attn203 **kwargs: Unpack[FlashAttentionKwargs],204 ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:205 """Input shape: Batch x Time x Channel"""206 207 # if key_value_states are provided this layer is used as a cross-attention layer208 # for the decoder209 is_cross_attention = key_value_states is not None210 211 # determine input shapes212 bsz, tgt_len = hidden_states.shape[:-1]213 src_len = key_value_states.shape[1] if is_cross_attention else tgt_len214 215 q_input_shape = (bsz, tgt_len, -1, self.head_dim)216 kv_input_shape = (bsz, src_len, -1, self.head_dim)217 218 # get query proj219 query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)220 221 is_updated = False222 if past_key_values is not None:223 if isinstance(past_key_values, EncoderDecoderCache):224 is_updated = past_key_values.is_updated.get(self.layer_idx)225 if is_cross_attention:226 # after the first generated id, we can subsequently re-use all key/value_states from cache227 curr_past_key_value = past_key_values.cross_attention_cache228 else:229 curr_past_key_value = past_key_values.self_attention_cache230 else:231 curr_past_key_value = past_key_values232 233 current_states = key_value_states if is_cross_attention else hidden_states234 if is_cross_attention and past_key_values is not None and is_updated:235 # reuse k,v, cross_attentions236 key_states = curr_past_key_value.layers[self.layer_idx].keys237 value_states = curr_past_key_value.layers[self.layer_idx].values238 else:239 key_states = self.k_proj(current_states)240 value_states = self.v_proj(current_states)241 key_states = key_states.view(*kv_input_shape).transpose(1, 2)242 value_states = value_states.view(*kv_input_shape).transpose(1, 2)243 244 if past_key_values is not None:245 # save all key/value_states to cache to be re-used for fast auto-regressive generation246 cache_position = cache_position if not is_cross_attention else None247 key_states, value_states = curr_past_key_value.update(248 key_states, value_states, self.layer_idx, {"cache_position": cache_position}249 )250 # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls251 if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):252 past_key_values.is_updated[self.layer_idx] = True253 254 attention_interface: Callable = eager_attention_forward255 if self.config._attn_implementation != "eager":256 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]257 258 attn_output, attn_weights = attention_interface(259 self,260 query_states,261 key_states,262 value_states,263 attention_mask,264 dropout=0.0 if not self.training else self.dropout,265 scaling=self.scaling,266 output_attentions=output_attentions,267 head_mask=layer_head_mask,268 **kwargs,269 )270 271 attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()272 attn_output = self.out_proj(attn_output)273 274 return attn_output, attn_weights275 276 277class BartEncoderLayer(GradientCheckpointingLayer):278 def __init__(self, config: BartConfig, layer_idx: Optional[int] = None):279 super().__init__()280 self.embed_dim = config.d_model281 282 self.self_attn = BartAttention(283 embed_dim=self.embed_dim,284 num_heads=config.encoder_attention_heads,285 dropout=config.attention_dropout,286 config=config,287 layer_idx=layer_idx,288 )289 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)290 self.dropout = config.dropout291 self.activation_fn = ACT2FN[config.activation_function]292 self.activation_dropout = config.activation_dropout293 self.fc1 = nn.Linear(self.embed_dim, config.encoder_ffn_dim)294 self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)295 self.final_layer_norm = nn.LayerNorm(self.embed_dim)296 297 def forward(298 self,299 hidden_states: torch.FloatTensor,300 attention_mask: torch.FloatTensor,301 layer_head_mask: torch.FloatTensor,302 output_attentions: Optional[bool] = False,303 ) -> tuple[torch.FloatTensor, Optional[torch.FloatTensor]]:304 """305 Args:306 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`307 attention_mask (`torch.FloatTensor`): attention mask of size308 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.309 layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size310 `(encoder_attention_heads,)`.311 output_attentions (`bool`, *optional*):312 Whether or not to return the attentions tensors of all attention layers. See `attentions` under313 returned tensors for more detail.314 """315 residual = hidden_states316 hidden_states, attn_weights = self.self_attn(317 hidden_states=hidden_states,318 attention_mask=attention_mask,319 layer_head_mask=layer_head_mask,320 output_attentions=output_attentions,321 )322 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)323 hidden_states = residual + hidden_states324 hidden_states = self.self_attn_layer_norm(hidden_states)325 326 residual = hidden_states327 hidden_states = self.activation_fn(self.fc1(hidden_states))328 hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)329 hidden_states = self.fc2(hidden_states)330 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)331 hidden_states = residual + hidden_states332 hidden_states = self.final_layer_norm(hidden_states)333 334 if hidden_states.dtype == torch.float16 and (335 torch.isinf(hidden_states).any() or torch.isnan(hidden_states).any()336 ):337 clamp_value = torch.finfo(hidden_states.dtype).max - 1000338 hidden_states = torch.clamp(hidden_states, min=-clamp_value, max=clamp_value)339 340 outputs = (hidden_states,)341 342 if output_attentions:343 outputs += (attn_weights,)344 345 return outputs346 347 348class BartDecoderLayer(GradientCheckpointingLayer):349 def __init__(self, config: BartConfig, layer_idx: Optional[int] = None):350 super().__init__()351 self.embed_dim = config.d_model352 353 self.self_attn = BartAttention(354 embed_dim=self.embed_dim,355 num_heads=config.decoder_attention_heads,356 dropout=config.attention_dropout,357 is_decoder=True,358 is_causal=True,359 config=config,360 layer_idx=layer_idx,361 )362 self.dropout = config.dropout363 self.activation_fn = ACT2FN[config.activation_function]364 self.activation_dropout = config.activation_dropout365 366 self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)367 self.encoder_attn = BartAttention(368 self.embed_dim,369 config.decoder_attention_heads,370 dropout=config.attention_dropout,371 is_decoder=True,372 config=config,373 layer_idx=layer_idx,374 )375 self.encoder_attn_layer_norm = nn.LayerNorm(self.embed_dim)376 self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)377 self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)378 self.final_layer_norm = nn.LayerNorm(self.embed_dim)379 380 @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")381 def forward(382 self,383 hidden_states: torch.Tensor,384 attention_mask: Optional[torch.Tensor] = None,385 encoder_hidden_states: Optional[torch.Tensor] = None,386 encoder_attention_mask: Optional[torch.Tensor] = None,387 layer_head_mask: Optional[torch.Tensor] = None,388 cross_attn_layer_head_mask: Optional[torch.Tensor] = None,389 past_key_values: Optional[Cache] = None,390 output_attentions: Optional[bool] = False,391 use_cache: Optional[bool] = True,392 cache_position: Optional[torch.Tensor] = None,393 ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:394 """395 Args:396 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`397 attention_mask (`torch.FloatTensor`): attention mask of size398 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.399 encoder_hidden_states (`torch.FloatTensor`):400 cross attention input to the layer of shape `(batch, seq_len, embed_dim)`401 encoder_attention_mask (`torch.FloatTensor`): encoder attention mask of size402 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.403 layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size404 `(encoder_attention_heads,)`.405 cross_attn_layer_head_mask (`torch.FloatTensor`): mask for cross-attention heads in a given layer of406 size `(decoder_attention_heads,)`.407 past_key_values (`Cache`): cached past key and value projection states408 output_attentions (`bool`, *optional*):409 Whether or not to return the attentions tensors of all attention layers. See `attentions` under410 returned tensors for more detail.411 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):412 Indices depicting the position of the input sequence tokens in the sequence. It is used to update the413 cache in the correct position and to infer the complete sequence length.414 """415 residual = hidden_states416 417 # Self Attention418 hidden_states, self_attn_weights = self.self_attn(419 hidden_states=hidden_states,420 past_key_values=past_key_values,421 attention_mask=attention_mask,422 layer_head_mask=layer_head_mask,423 output_attentions=output_attentions,424 cache_position=cache_position,425 )426 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)427 hidden_states = residual + hidden_states428 hidden_states = self.self_attn_layer_norm(hidden_states)429 430 # Cross-Attention Block431 cross_attn_weights = None432 if encoder_hidden_states is not None:433 residual = hidden_states434 435 hidden_states, cross_attn_weights = self.encoder_attn(436 hidden_states=hidden_states,437 key_value_states=encoder_hidden_states,438 attention_mask=encoder_attention_mask,439 layer_head_mask=cross_attn_layer_head_mask,440 past_key_values=past_key_values,441 output_attentions=output_attentions,442 cache_position=cache_position,443 )444 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)445 hidden_states = residual + hidden_states446 hidden_states = self.encoder_attn_layer_norm(hidden_states)447 448 # Fully Connected449 residual = hidden_states450 hidden_states = self.activation_fn(self.fc1(hidden_states))451 hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)452 hidden_states = self.fc2(hidden_states)453 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)454 hidden_states = residual + hidden_states455 hidden_states = self.final_layer_norm(hidden_states)456 457 outputs = (hidden_states,)458 459 if output_attentions:460 outputs += (self_attn_weights, cross_attn_weights)461 462 return outputs463 464 465class BartClassificationHead(nn.Module):466 """Head for sentence-level classification tasks."""467 468 def __init__(469 self,470 input_dim: int,471 inner_dim: int,472 num_classes: int,473 pooler_dropout: float,474 ):475 super().__init__()476 self.dense = nn.Linear(input_dim, inner_dim)477 self.dropout = nn.Dropout(p=pooler_dropout)478 self.out_proj = nn.Linear(inner_dim, num_classes)479 480 def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:481 hidden_states = self.dropout(hidden_states)482 hidden_states = self.dense(hidden_states)483 hidden_states = torch.tanh(hidden_states)484 hidden_states = self.dropout(hidden_states)485 hidden_states = self.out_proj(hidden_states)486 return hidden_states487 488 489@auto_docstring490class BartPreTrainedModel(PreTrainedModel):491 config: BartConfig492 base_model_prefix = "model"493 supports_gradient_checkpointing = True494 _keys_to_ignore_on_load_unexpected = ["encoder.version", "decoder.version"]495 _no_split_modules = [r"BartEncoderLayer", r"BartDecoderLayer"]496 _skip_keys_device_placement = "past_key_values"497 _supports_flash_attn = True498 _supports_sdpa = True499 _supports_flex_attn = True500 501 _can_compile_fullgraph = True502 503 def _init_weights(self, module):504 std = self.config.init_std505 if isinstance(module, nn.Linear):506 module.weight.data.normal_(mean=0.0, std=std)507 if module.bias is not None:508 module.bias.data.zero_()509 elif isinstance(module, nn.Embedding):510 module.weight.data.normal_(mean=0.0, std=std)511 if module.padding_idx is not None:512 module.weight.data[module.padding_idx].zero_()513 elif isinstance(module, nn.LayerNorm):514 module.weight.data.fill_(1.0)515 module.bias.data.zero_()516 517 @property518 def dummy_inputs(self):519 pad_token = self.config.pad_token_id520 input_ids = torch.tensor([[0, 6, 10, 4, 2], [0, 8, 12, 2, pad_token]], device=self.device)521 dummy_inputs = {522 "attention_mask": input_ids.ne(pad_token),523 "input_ids": input_ids,524 }525 return dummy_inputs526 527 def _update_full_mask(528 self,529 attention_mask: Union[torch.Tensor, None],530 inputs_embeds: torch.Tensor,531 ):532 if attention_mask is not None:533 if self.config._attn_implementation == "flash_attention_2":534 attention_mask = attention_mask if 0 in attention_mask else None535 elif self.config._attn_implementation == "sdpa":536 # output_attentions=True & head_mask can not be supported when using SDPA, fall back to537 # the manual implementation that requires a 4D causal mask in all cases.538 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]539 attention_mask = _prepare_4d_attention_mask_for_sdpa(attention_mask, inputs_embeds.dtype)540 elif self.config._attn_implementation == "flex_attention":541 if isinstance(attention_mask, torch.Tensor):542 attention_mask = make_flex_block_causal_mask(attention_mask, is_causal=False)543 else:544 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]545 attention_mask = _prepare_4d_attention_mask(attention_mask, inputs_embeds.dtype)546 547 return attention_mask548 549 def _update_causal_mask(550 self,551 attention_mask: Optional[Union[torch.Tensor, "BlockMask"]],552 input_tensor: torch.Tensor,553 cache_position: torch.Tensor,554 past_key_values: Cache,555 ):556 if self.config._attn_implementation == "flex_attention":557 if isinstance(attention_mask, torch.Tensor):558 attention_mask = make_flex_block_causal_mask(attention_mask)559 # Other attention flavors support in-built causal (when `mask is None`)560 # while we need to create our specific block mask regardless561 elif attention_mask is None:562 attention_mask = make_flex_block_causal_mask(563 torch.ones(564 size=(input_tensor.shape[0], input_tensor.shape[1]),565 device=attention_mask.device,566 )567 )568 return attention_mask569 570 if self.config._attn_implementation == "flash_attention_2":571 if attention_mask is not None and (attention_mask == 0.0).any():572 return attention_mask573 return None574 575 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in576 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail577 # to infer the attention mask.578 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0579 using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False580 581 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward582 if self.config._attn_implementation == "sdpa" and not using_compilable_cache:583 if AttentionMaskConverter._ignore_causal_mask_sdpa(584 attention_mask,585 inputs_embeds=input_tensor,586 past_key_values_length=past_seen_tokens,587 is_training=self.training,588 ):589 return None590 591 dtype = input_tensor.dtype592 sequence_length = input_tensor.shape[1]593 if using_compilable_cache:594 target_length = past_key_values.get_max_cache_shape()595 else:596 target_length = (597 attention_mask.shape[-1]598 if isinstance(attention_mask, torch.Tensor)599 else past_seen_tokens + sequence_length + 1600 )601 602 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).603 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(604 attention_mask,605 sequence_length=sequence_length,606 target_length=target_length,607 dtype=dtype,608 cache_position=cache_position,609 batch_size=input_tensor.shape[0],610 )611 612 if (613 self.config._attn_implementation == "sdpa"614 and attention_mask is not None615 and attention_mask.device.type in ["cuda", "xpu", "npu"]616 ):617 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when618 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.619 # Details: https://github.com/pytorch/pytorch/issues/110213620 min_dtype = torch.finfo(dtype).min621 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)622 623 return causal_mask624 625 @staticmethod626 # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position627 def _prepare_4d_causal_attention_mask_with_cache_position(628 attention_mask: torch.Tensor,629 sequence_length: int,630 target_length: int,631 dtype: torch.dtype,632 cache_position: torch.Tensor,633 batch_size: int,634 **kwargs,635 ):636 """637 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape638 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.639 640 Args:641 attention_mask (`torch.Tensor`):642 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape643 `(batch_size, 1, query_length, key_value_length)`.644 sequence_length (`int`):645 The sequence length being processed.646 target_length (`int`):647 The target length: when generating with static cache, the mask should be as long as the static cache,648 to account for the 0 padding, the part of the cache that is not filled yet.649 dtype (`torch.dtype`):650 The dtype to use for the 4D attention mask.651 cache_position (`torch.Tensor`):652 Indices depicting the position of the input sequence tokens in the sequence.653 batch_size (`torch.Tensor`):654 Batch size.655 """656 if attention_mask is not None and attention_mask.dim() == 4:657 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.658 causal_mask = attention_mask659 else:660 min_dtype = torch.finfo(dtype).min661 causal_mask = torch.full(662 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device663 )664 if sequence_length != 1:665 causal_mask = torch.triu(causal_mask, diagonal=1)666 causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)667 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)668 if attention_mask is not None:669 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit670 mask_length = attention_mask.shape[-1]671 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(672 causal_mask.device673 )674 padding_mask = padding_mask == 0675 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(676 padding_mask, min_dtype677 )678 679 return causal_mask680 681 def _update_cross_attn_mask(682 self,683 encoder_hidden_states: Union[torch.Tensor, None],684 encoder_attention_mask: Union[torch.Tensor, None],685 input_shape: torch.Size,686 inputs_embeds: torch.Tensor,687 ):688 # expand encoder attention mask689 if encoder_hidden_states is not None and encoder_attention_mask is not None:690 if self.config._attn_implementation == "flash_attention_2":691 encoder_attention_mask = encoder_attention_mask if 0 in encoder_attention_mask else None692 elif self.config._attn_implementation == "sdpa":693 # output_attentions=True & cross_attn_head_mask can not be supported when using SDPA, and we fall back on694 # the manual implementation that requires a 4D causal mask in all cases.695 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]696 encoder_attention_mask = _prepare_4d_attention_mask_for_sdpa(697 encoder_attention_mask,698 inputs_embeds.dtype,699 tgt_len=input_shape[-1],700 )701 elif self.config._attn_implementation == "flex_attention":702 if isinstance(encoder_attention_mask, torch.Tensor):703 encoder_attention_mask = make_flex_block_causal_mask(704 encoder_attention_mask,705 query_length=input_shape[-1],706 is_causal=False,707 )708 else:709 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]710 encoder_attention_mask = _prepare_4d_attention_mask(711 encoder_attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]712 )713 714 return encoder_attention_mask715 716 717class PretrainedBartModel(BartPreTrainedModel):718 def __init_subclass__(self):719 warnings.warn(720 "The class `PretrainedBartModel` has been depreciated, please use `BartPreTrainedModel` instead.",721 FutureWarning,722 )723 724 725class BartPretrainedModel(BartPreTrainedModel):726 def __init_subclass__(self):727 warnings.warn(728 "The class `PretrainedBartModel` has been depreciated, please use `BartPreTrainedModel` instead.",729 FutureWarning,730 )731 732 733class BartEncoder(BartPreTrainedModel):734 """735 Transformer encoder consisting of *config.encoder_layers* self attention layers. Each layer is a736 [`BartEncoderLayer`].737 738 Args:739 config: BartConfig740 embed_tokens (nn.Embedding): output embedding741 """742 743 def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):744 super().__init__(config)745 746 self.dropout = config.dropout747 self.layerdrop = config.encoder_layerdrop748 749 embed_dim = config.d_model750 self.padding_idx = config.pad_token_id751 self.max_source_positions = config.max_position_embeddings752 embed_scale = math.sqrt(embed_dim) if config.scale_embedding else 1.0753 754 self.embed_tokens = BartScaledWordEmbedding(755 config.vocab_size, embed_dim, self.padding_idx, embed_scale=embed_scale756 )757 758 if embed_tokens is not None:759 self.embed_tokens.weight = embed_tokens.weight760 761 self.embed_positions = BartLearnedPositionalEmbedding(762 config.max_position_embeddings,763 embed_dim,764 )765 self.layers = nn.ModuleList([BartEncoderLayer(config, layer_idx=i) for i in range(config.encoder_layers)])766 self.layernorm_embedding = nn.LayerNorm(embed_dim)767 768 self.gradient_checkpointing = False769 # Initialize weights and apply final processing770 self.post_init()771 772 def forward(773 self,774 input_ids: Optional[torch.LongTensor] = None,775 attention_mask: Optional[torch.Tensor] = None,776 head_mask: Optional[torch.Tensor] = None,777 inputs_embeds: Optional[torch.FloatTensor] = None,778 output_attentions: Optional[bool] = None,779 output_hidden_states: Optional[bool] = None,780 return_dict: Optional[bool] = None,781 ) -> Union[tuple, BaseModelOutput]:782 r"""783 Args:784 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):785 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you786 provide it.787 788 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and789 [`PreTrainedTokenizer.__call__`] for details.790 791 [What are input IDs?](../glossary#input-ids)792 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):793 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:794 795 - 1 for tokens that are **not masked**,796 - 0 for tokens that are **masked**.797 798 [What are attention masks?](../glossary#attention-mask)799 head_mask (`torch.Tensor` of shape `(encoder_layers, encoder_attention_heads)`, *optional*):800 Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:801 802 - 1 indicates the head is **not masked**,803 - 0 indicates the head is **masked**.804 805 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):806 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.807 This is useful if you want more control over how to convert `input_ids` indices into associated vectors808 than the model's internal embedding lookup matrix.809 output_attentions (`bool`, *optional*):810 Whether or not to return the attentions tensors of all attention layers. See `attentions` under811 returned tensors for more detail.812 output_hidden_states (`bool`, *optional*):813 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors814 for more detail.815 return_dict (`bool`, *optional*):816 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.817 """818 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions819 output_hidden_states = (820 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states821 )822 return_dict = return_dict if return_dict is not None else self.config.use_return_dict823 824 # retrieve input_ids and inputs_embeds825 if input_ids is not None and inputs_embeds is not None:826 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")827 elif input_ids is not None:828 input = input_ids829 input_ids = input_ids.view(-1, input_ids.shape[-1])830 elif inputs_embeds is not None:831 input = inputs_embeds[:, :, -1]832 else:833 raise ValueError("You have to specify either input_ids or inputs_embeds")834 835 if inputs_embeds is None:836 inputs_embeds = self.embed_tokens(input_ids)837 838 embed_pos = self.embed_positions(input)839 embed_pos = embed_pos.to(inputs_embeds.device)840 841 hidden_states = inputs_embeds + embed_pos842 hidden_states = self.layernorm_embedding(hidden_states)843 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)844 845 attention_mask = self._update_full_mask(846 attention_mask,847 inputs_embeds,848 )849 850 encoder_states = () if output_hidden_states else None851 all_attentions = () if output_attentions else None852 853 # check if head_mask has a correct number of layers specified if desired854 if head_mask is not None:855 if head_mask.size()[0] != (len(self.layers)):856 raise ValueError(857 f"The head_mask should be specified for {len(self.layers)} layers, but it is for"858 f" {head_mask.size()[0]}."859 )860 861 for idx, encoder_layer in enumerate(self.layers):862 if output_hidden_states:863 encoder_states = encoder_states + (hidden_states,)864 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)865 to_drop = False866 if self.training:867 dropout_probability = torch.rand([])868 if dropout_probability < self.layerdrop: # skip the layer869 to_drop = True870 871 if to_drop:872 layer_outputs = (None, None)873 else:874 layer_outputs = encoder_layer(875 hidden_states,876 attention_mask,877 layer_head_mask=(head_mask[idx] if head_mask is not None else None),878 output_attentions=output_attentions,879 )880 881 hidden_states = layer_outputs[0]882 883 if output_attentions:884 all_attentions = all_attentions + (layer_outputs[1],)885 886 if output_hidden_states:887 encoder_states = encoder_states + (hidden_states,)888 889 if not return_dict:890 return tuple(v for v in [hidden_states, encoder_states, all_attentions] if v is not None)891 return BaseModelOutput(892 last_hidden_state=hidden_states, hidden_states=encoder_states, attentions=all_attentions893 )894 895 896class BartDecoder(BartPreTrainedModel):897 """898 Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`BartDecoderLayer`]899 900 Args:901 config: BartConfig902 embed_tokens (nn.Embedding): output embedding903 """904 905 def __init__(self, config: BartConfig, embed_tokens: Optional[nn.Embedding] = None):906 super().__init__(config)907 self.dropout = config.dropout908 self.layerdrop = config.decoder_layerdrop909 self.padding_idx = config.pad_token_id910 self.max_target_positions = config.max_position_embeddings911 embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0912 913 self.embed_tokens = BartScaledWordEmbedding(914 config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale915 )916 917 if embed_tokens is not None:918 self.embed_tokens.weight = embed_tokens.weight919 920 self.embed_positions = BartLearnedPositionalEmbedding(921 config.max_position_embeddings,922 config.d_model,923 )924 self.layers = nn.ModuleList([BartDecoderLayer(config, layer_idx=i) for i in range(config.decoder_layers)])925 926 self.layernorm_embedding = nn.LayerNorm(config.d_model)927 928 self.gradient_checkpointing = False929 # Initialize weights and apply final processing930 self.post_init()931 932 def forward(933 self,934 input_ids: Optional[torch.LongTensor] = None,935 attention_mask: Optional[torch.Tensor] = None,936 encoder_hidden_states: Optional[torch.FloatTensor] = None,937 encoder_attention_mask: Optional[torch.LongTensor] = None,938 head_mask: Optional[torch.Tensor] = None,939 cross_attn_head_mask: Optional[torch.Tensor] = None,940 past_key_values: Optional[Cache] = None,941 inputs_embeds: Optional[torch.FloatTensor] = None,942 use_cache: Optional[bool] = None,943 output_attentions: Optional[bool] = None,944 output_hidden_states: Optional[bool] = None,945 return_dict: Optional[bool] = None,946 cache_position: Optional[torch.LongTensor] = None,947 ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:948 r"""949 Args:950 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):951 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you952 provide it.953 954 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and955 [`PreTrainedTokenizer.__call__`] for details.956 957 [What are input IDs?](../glossary#input-ids)958 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):959 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:960 961 - 1 for tokens that are **not masked**,962 - 0 for tokens that are **masked**.963 964 [What are attention masks?](../glossary#attention-mask)965 encoder_hidden_states (`torch.FloatTensor` of shape `(batch_size, encoder_sequence_length, hidden_size)`, *optional*):966 Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention967 of the decoder.968 encoder_attention_mask (`torch.LongTensor` of shape `(batch_size, encoder_sequence_length)`, *optional*):969 Mask to avoid performing cross-attention on padding tokens indices of encoder input_ids. Mask values970 selected in `[0, 1]`:971 972 - 1 for tokens that are **not masked**,973 - 0 for tokens that are **masked**.974 975 [What are attention masks?](../glossary#attention-mask)976 head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):977 Mask to nullify selected heads of the attention modules. Mask values selected in `[0, 1]`:978 979 - 1 indicates the head is **not masked**,980 - 0 indicates the head is **masked**.981 982 cross_attn_head_mask (`torch.Tensor` of shape `(decoder_layers, decoder_attention_heads)`, *optional*):983 Mask to nullify selected heads of the cross-attention modules in the decoder to avoid performing984 cross-attention on hidden heads. Mask values selected in `[0, 1]`:985 986 - 1 indicates the head is **not masked**,987 - 0 indicates the head is **masked**.988 989 past_key_values (`Cache`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):990 It is a [`~cache_utils.Cache`] instance. For more details, see our [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache).991 992 Contains pre-computed hidden-states (key and values in the self-attention blocks and in the993 cross-attention blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.994 995 If `past_key_values` are used, the user can optionally input only the last `decoder_input_ids` (those996 that don't have their past key value states given to this model) of shape `(batch_size, 1)` instead of997 all `decoder_input_ids` of shape `(batch_size, sequence_length)`.998 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):999 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation.1000 This is useful if you want more control over how to convert `input_ids` indices into associated vectors1001 than the model's internal embedding lookup matrix.1002 output_attentions (`bool`, *optional*):1003 Whether or not to return the attentions tensors of all attention layers. See `attentions` under1004 returned tensors for more detail.1005 output_hidden_states (`bool`, *optional*):1006 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors1007 for more detail.1008 return_dict (`bool`, *optional*):1009 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.1010 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):1011 Indices depicting the position of the input sequence tokens in the sequence. It is used to update the1012 cache in the correct position and to infer the complete sequence length.1013 """1014 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1015 output_hidden_states = (1016 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1017 )1018 use_cache = use_cache if use_cache is not None else self.config.use_cache1019 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1020 1021 if self.gradient_checkpointing and self.training:1022 if use_cache:1023 logger.warning_once(1024 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."1025 )1026 use_cache = False1027 1028 # retrieve input_ids and inputs_embeds1029 if (input_ids is None) ^ (inputs_embeds is not None):1030 raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")1031 elif input_ids is not None:1032 input = input_ids1033 input_shape = input.shape1034 input_ids = input_ids.view(-1, input_shape[-1])1035 elif inputs_embeds is not None:1036 input_shape = inputs_embeds.size()[:-1]1037 input = inputs_embeds[:, :, -1]1038 else:1039 raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")1040 1041 if inputs_embeds is None:1042 inputs_embeds = self.embed_tokens(input)1043 1044 # initialize `past_key_values`1045 if use_cache and past_key_values is None:1046 past_key_values = (1047 EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))1048 if encoder_hidden_states is not None1049 else DynamicCache(config=self.config)1050 )1051 if use_cache and isinstance(past_key_values, tuple):1052 logger.warning_once(1053 "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. "1054 "You should pass an instance of `EncoderDecoderCache` instead, e.g. "1055 "`past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`."1056 )1057 past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)1058 1059 batch_size, seq_length = inputs_embeds.size()[:-1]1060 past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 01061 if cache_position is None:1062 cache_position = torch.arange(1063 past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device1064 )1065 1066 if attention_mask is None and not is_torchdynamo_compiling():1067 # required mask seq length can be calculated via length of past cache1068 mask_seq_length = past_key_values_length + seq_length1069 attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)1070 1071 self_attn_cache = (1072 past_key_values.self_attention_cache1073 if isinstance(past_key_values, EncoderDecoderCache)1074 else past_key_values1075 )1076 1077 attention_mask = self._update_causal_mask(1078 attention_mask,1079 inputs_embeds,1080 cache_position,1081 self_attn_cache,1082 )1083 encoder_attention_mask = self._update_cross_attn_mask(1084 encoder_hidden_states,1085 encoder_attention_mask,1086 input_shape,1087 inputs_embeds,1088 )1089 1090 # embed positions1091 positions = self.embed_positions(input, past_key_values_length, position_ids=cache_position)1092 positions = positions.to(inputs_embeds.device)1093 1094 hidden_states = inputs_embeds + positions1095 hidden_states = self.layernorm_embedding(hidden_states)1096 1097 hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)1098 1099 # decoder layers1100 all_hidden_states = () if output_hidden_states else None1101 all_self_attns = () if output_attentions else None1102 all_cross_attentions = () if (output_attentions and encoder_hidden_states is not None) else None1103 1104 # check if head_mask/cross_attn_head_mask has a correct number of layers specified if desired1105 for attn_mask, mask_name in zip([head_mask, cross_attn_head_mask], ["head_mask", "cross_attn_head_mask"]):1106 if attn_mask is not None:1107 if attn_mask.size()[0] != (len(self.layers)):1108 raise ValueError(1109 f"The `{mask_name}` should be specified for {len(self.layers)} layers, but it is for"1110 f" {head_mask.size()[0]}."1111 )1112 1113 for idx, decoder_layer in enumerate(self.layers):1114 # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)1115 if output_hidden_states:1116 all_hidden_states += (hidden_states,)1117 if self.training:1118 dropout_probability = torch.rand([])1119 if dropout_probability < self.layerdrop:1120 continue1121 1122 layer_outputs = decoder_layer(1123 hidden_states,1124 attention_mask,1125 encoder_hidden_states, # as a positional argument for gradient checkpointing1126 encoder_attention_mask=encoder_attention_mask,1127 layer_head_mask=(head_mask[idx] if head_mask is not None else None),1128 cross_attn_layer_head_mask=(cross_attn_head_mask[idx] if cross_attn_head_mask is not None else None),1129 past_key_values=past_key_values,1130 output_attentions=output_attentions,1131 use_cache=use_cache,1132 cache_position=cache_position,1133 )1134 hidden_states = layer_outputs[0]1135 if output_attentions:1136 all_self_attns += (layer_outputs[1],)1137 1138 if encoder_hidden_states is not None:1139 all_cross_attentions += (layer_outputs[2],)1140 1141 # add hidden states from the last decoder layer1142 if output_hidden_states:1143 all_hidden_states += (hidden_states,)1144 1145 if not return_dict:1146 return tuple(1147 v1148 for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns, all_cross_attentions]1149 if v is not None1150 )1151 return BaseModelOutputWithPastAndCrossAttentions(1152 last_hidden_state=hidden_states,1153 past_key_values=past_key_values,1154 hidden_states=all_hidden_states,1155 attentions=all_self_attns,1156 cross_attentions=all_cross_attentions,1157 )1158 1159 1160@auto_docstring1161class BartModel(BartPreTrainedModel):1162 _tied_weights_keys = ["encoder.embed_tokens.weight", "decoder.embed_tokens.weight"]1163 1164 def __init__(self, config: BartConfig):1165 super().__init__(config)1166 1167 padding_idx, vocab_size = config.pad_token_id, config.vocab_size1168 embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.01169 self.shared = BartScaledWordEmbedding(vocab_size, config.d_model, padding_idx, embed_scale=embed_scale)1170 1171 self.encoder = BartEncoder(config, self.shared)1172 self.decoder = BartDecoder(config, self.shared)1173 1174 # Initialize weights and apply final processing1175 self.post_init()1176 1177 def _tie_weights(self):1178 if self.config.tie_word_embeddings:1179 # Some model checkpoints like "facebook/bart-large-cnn"'s embedding weight is in decoder.embed_tokens, need check here, see issue #362471180 if self.shared.weight.device == torch.device(1181 "meta"1182 ) and self.decoder.embed_tokens.weight.device != torch.device("meta"):1183 self._tie_or_clone_weights(self.encoder.embed_tokens, self.decoder.embed_tokens)1184 self._tie_or_clone_weights(self.shared, self.decoder.embed_tokens)1185 else:1186 self._tie_or_clone_weights(self.encoder.embed_tokens, self.shared)1187 self._tie_or_clone_weights(self.decoder.embed_tokens, self.shared)1188 1189 def get_input_embeddings(self):1190 return self.shared1191 1192 def set_input_embeddings(self, value):1193 self.shared = value1194 self.encoder.embed_tokens = self.shared1195 self.decoder.embed_tokens = self.shared1196 1197 def get_encoder(self):1198 return self.encoder1199 1200 @auto_docstring