cyrilvallez/test_remote_code_dummy_llama
024
1# coding=utf-82# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.3#4# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX5# and OPT implementations in this library. It has been modified from its6# original forms to accommodate minor architectural differences compared7# to GPT-NeoX and OPT used by the Meta AI team that trained the model.8#9# Licensed under the Apache License, Version 2.0 (the "License");10# you may not use this file except in compliance with the License.11# You may obtain a copy of the License at12#13# http://www.apache.org/licenses/LICENSE-2.014#15# Unless required by applicable law or agreed to in writing, software16# distributed under the License is distributed on an "AS IS" BASIS,17# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.18# See the License for the specific language governing permissions and19# limitations under the License.20from typing import Callable, List, Optional, Tuple, Union21 22import torch23import torch.utils.checkpoint24from torch import nn25 26from transformers.activations import ACT2FN27from transformers.cache_utils import Cache, DynamicCache, StaticCache28from transformers.generation import GenerationMixin29from transformers.modeling_attn_mask_utils import AttentionMaskConverter30from transformers.modeling_flash_attention_utils import FlashAttentionKwargs31from transformers.modeling_outputs import (32 BaseModelOutputWithPast,33 CausalLMOutputWithPast,34 QuestionAnsweringModelOutput,35 SequenceClassifierOutputWithPast,36 TokenClassifierOutput,37)38from transformers.modeling_rope_utils import ROPE_INIT_FUNCTIONS39from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel40from transformers.processing_utils import Unpack41from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS42from transformers.utils import (43 add_code_sample_docstrings,44 add_start_docstrings,45 add_start_docstrings_to_model_forward,46 logging,47 replace_return_docstrings,48)49from transformers.utils.deprecation import deprecate_kwarg50from .configuration_dummy_llama import DummyLlamaConfig51 52 53logger = logging.get_logger(__name__)54 55_CHECKPOINT_FOR_DOC = "meta-llama/DummyLlama-2-7b-hf"56_CONFIG_FOR_DOC = "DummyLlamaConfig"57 58 59class DummyLlamaRMSNorm(nn.Module):60 def __init__(self, hidden_size, eps=1e-6):61 """62 DummyLlamaRMSNorm is equivalent to T5LayerNorm63 """64 super().__init__()65 self.weight = nn.Parameter(torch.ones(hidden_size))66 self.variance_epsilon = eps67 68 def forward(self, hidden_states):69 input_dtype = hidden_states.dtype70 hidden_states = hidden_states.to(torch.float32)71 variance = hidden_states.pow(2).mean(-1, keepdim=True)72 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)73 return self.weight * hidden_states.to(input_dtype)74 75 def extra_repr(self):76 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"77 78 79ALL_LAYERNORM_LAYERS.append(DummyLlamaRMSNorm)80 81 82class DummyLlamaRotaryEmbedding(nn.Module):83 def __init__(self, config: DummyLlamaConfig, device=None):84 super().__init__()85 # BC: "rope_type" was originally "type"86 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:87 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))88 else:89 self.rope_type = "default"90 self.max_seq_len_cached = config.max_position_embeddings91 self.original_max_seq_len = config.max_position_embeddings92 93 self.config = config94 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]95 96 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)97 self.register_buffer("inv_freq", inv_freq, persistent=False)98 self.original_inv_freq = self.inv_freq99 100 def _dynamic_frequency_update(self, position_ids, device):101 """102 dynamic RoPE layers should recompute `inv_freq` in the following situations:103 1 - growing beyond the cached sequence length (allow scaling)104 2 - the current sequence length is in the original scale (avoid losing precision with small sequences)105 """106 seq_len = torch.max(position_ids) + 1107 if seq_len > self.max_seq_len_cached: # growth108 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)109 self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation110 self.max_seq_len_cached = seq_len111 112 if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset113 # This .to() is needed if the model has been moved to a device after being initialized (because114 # the buffer is automatically moved, but not the original copy)115 self.original_inv_freq = self.original_inv_freq.to(device)116 self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)117 self.max_seq_len_cached = self.original_max_seq_len118 119 @torch.no_grad()120 def forward(self, x, position_ids):121 if "dynamic" in self.rope_type:122 self._dynamic_frequency_update(position_ids, device=x.device)123 124 # Core RoPE block125 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)126 position_ids_expanded = position_ids[:, None, :].float()127 # Force float32 (see https://github.com/huggingface/transformers/pull/29285)128 device_type = x.device.type129 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"130 with torch.autocast(device_type=device_type, enabled=False):131 freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2)132 emb = torch.cat((freqs, freqs), dim=-1)133 cos = emb.cos()134 sin = emb.sin()135 136 # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention137 cos = cos * self.attention_scaling138 sin = sin * self.attention_scaling139 140 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)141 142 143def rotate_half(x):144 """Rotates half the hidden dims of the input."""145 x1 = x[..., : x.shape[-1] // 2]146 x2 = x[..., x.shape[-1] // 2 :]147 return torch.cat((-x2, x1), dim=-1)148 149 150def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):151 """Applies Rotary Position Embedding to the query and key tensors.152 153 Args:154 q (`torch.Tensor`): The query tensor.155 k (`torch.Tensor`): The key tensor.156 cos (`torch.Tensor`): The cosine part of the rotary embedding.157 sin (`torch.Tensor`): The sine part of the rotary embedding.158 position_ids (`torch.Tensor`, *optional*):159 Deprecated and unused.160 unsqueeze_dim (`int`, *optional*, defaults to 1):161 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and162 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note163 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and164 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes165 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have166 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.167 Returns:168 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.169 """170 cos = cos.unsqueeze(unsqueeze_dim)171 sin = sin.unsqueeze(unsqueeze_dim)172 q_embed = (q * cos) + (rotate_half(q) * sin)173 k_embed = (k * cos) + (rotate_half(k) * sin)174 return q_embed, k_embed175 176 177class DummyLlamaMLP(nn.Module):178 def __init__(self, config):179 super().__init__()180 self.config = config181 self.hidden_size = config.hidden_size182 self.intermediate_size = config.intermediate_size183 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)184 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)185 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)186 self.act_fn = ACT2FN[config.hidden_act]187 188 def forward(self, x):189 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))190 return down_proj191 192 193def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:194 """195 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,196 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)197 """198 batch, num_key_value_heads, slen, head_dim = hidden_states.shape199 if n_rep == 1:200 return hidden_states201 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)202 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)203 204 205def eager_attention_forward(206 module: nn.Module,207 query: torch.Tensor,208 key: torch.Tensor,209 value: torch.Tensor,210 attention_mask: Optional[torch.Tensor],211 scaling: float,212 dropout: float = 0.0,213 **kwargs,214):215 key_states = repeat_kv(key, module.num_key_value_groups)216 value_states = repeat_kv(value, module.num_key_value_groups)217 218 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling219 if attention_mask is not None:220 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]221 attn_weights = attn_weights + causal_mask222 223 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)224 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)225 attn_output = torch.matmul(attn_weights, value_states)226 attn_output = attn_output.transpose(1, 2).contiguous()227 228 return attn_output, attn_weights229 230 231class DummyLlamaAttention(nn.Module):232 """Multi-headed attention from 'Attention Is All You Need' paper"""233 234 def __init__(self, config: DummyLlamaConfig, layer_idx: int):235 super().__init__()236 self.config = config237 self.layer_idx = layer_idx238 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)239 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads240 self.scaling = self.head_dim**-0.5241 self.attention_dropout = config.attention_dropout242 self.is_causal = True243 244 self.q_proj = nn.Linear(245 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias246 )247 self.k_proj = nn.Linear(248 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias249 )250 self.v_proj = nn.Linear(251 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias252 )253 self.o_proj = nn.Linear(254 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias255 )256 257 def forward(258 self,259 hidden_states: torch.Tensor,260 position_embeddings: Tuple[torch.Tensor, torch.Tensor],261 attention_mask: Optional[torch.Tensor],262 past_key_value: Optional[Cache] = None,263 cache_position: Optional[torch.LongTensor] = None,264 **kwargs: Unpack[FlashAttentionKwargs],265 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:266 input_shape = hidden_states.shape[:-1]267 hidden_shape = (*input_shape, -1, self.head_dim)268 269 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)270 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)271 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)272 273 cos, sin = position_embeddings274 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)275 276 if past_key_value is not None:277 # sin and cos are specific to RoPE models; cache_position needed for the static cache278 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}279 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)280 281 attention_interface: Callable = eager_attention_forward282 if self.config._attn_implementation != "eager":283 if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):284 logger.warning_once(285 "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "286 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'287 )288 else:289 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]290 291 attn_output, attn_weights = attention_interface(292 self,293 query_states,294 key_states,295 value_states,296 attention_mask,297 dropout=0.0 if not self.training else self.attention_dropout,298 scaling=self.scaling,299 **kwargs,300 )301 302 attn_output = attn_output.reshape(*input_shape, -1).contiguous()303 attn_output = self.o_proj(attn_output)304 return attn_output, attn_weights305 306 307class DummyLlamaDecoderLayer(nn.Module):308 def __init__(self, config: DummyLlamaConfig, layer_idx: int):309 super().__init__()310 self.hidden_size = config.hidden_size311 312 self.self_attn = DummyLlamaAttention(config=config, layer_idx=layer_idx)313 314 self.mlp = DummyLlamaMLP(config)315 self.input_layernorm = DummyLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)316 self.post_attention_layernorm = DummyLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)317 318 def forward(319 self,320 hidden_states: torch.Tensor,321 attention_mask: Optional[torch.Tensor] = None,322 position_ids: Optional[torch.LongTensor] = None,323 past_key_value: Optional[Cache] = None,324 output_attentions: Optional[bool] = False,325 use_cache: Optional[bool] = False,326 cache_position: Optional[torch.LongTensor] = None,327 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC328 **kwargs: Unpack[FlashAttentionKwargs],329 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:330 residual = hidden_states331 332 hidden_states = self.input_layernorm(hidden_states)333 334 # Self Attention335 hidden_states, self_attn_weights = self.self_attn(336 hidden_states=hidden_states,337 attention_mask=attention_mask,338 position_ids=position_ids,339 past_key_value=past_key_value,340 output_attentions=output_attentions,341 use_cache=use_cache,342 cache_position=cache_position,343 position_embeddings=position_embeddings,344 **kwargs,345 )346 hidden_states = residual + hidden_states347 348 # Fully Connected349 residual = hidden_states350 hidden_states = self.post_attention_layernorm(hidden_states)351 hidden_states = self.mlp(hidden_states)352 hidden_states = residual + hidden_states353 354 outputs = (hidden_states,)355 if output_attentions:356 outputs += (self_attn_weights,)357 358 return outputs359 360 361class DummyLlamaPreTrainedModel(PreTrainedModel):362 config_class = DummyLlamaConfig363 base_model_prefix = "model"364 supports_gradient_checkpointing = True365 _no_split_modules = ["DummyLlamaDecoderLayer"]366 _skip_keys_device_placement = ["past_key_values"]367 _supports_flash_attn_2 = True368 _supports_sdpa = True369 _supports_flex_attn = True370 _supports_cache_class = True371 _supports_quantized_cache = True372 _supports_static_cache = True373 _supports_attention_backend = True374 375 376 377class DummyLlamaModel(DummyLlamaPreTrainedModel):378 """379 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`DummyLlamaDecoderLayer`]380 381 Args:382 config: DummyLlamaConfig383 """384 385 def __init__(self, config: DummyLlamaConfig):386 super().__init__(config)387 self.padding_idx = config.pad_token_id388 self.vocab_size = config.vocab_size389 390 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)391 self.layers = nn.ModuleList(392 [DummyLlamaDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]393 )394 self.norm = DummyLlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)395 self.rotary_emb = DummyLlamaRotaryEmbedding(config=config)396 self.gradient_checkpointing = False397 398 # Initialize weights and apply final processing399 self.post_init()400 401 def get_input_embeddings(self):402 return self.embed_tokens403 404 def set_input_embeddings(self, value):405 self.embed_tokens = value406 407 def forward(408 self,409 input_ids: torch.LongTensor = None,410 attention_mask: Optional[torch.Tensor] = None,411 position_ids: Optional[torch.LongTensor] = None,412 past_key_values: Optional[Cache] = None,413 inputs_embeds: Optional[torch.FloatTensor] = None,414 use_cache: Optional[bool] = None,415 output_attentions: Optional[bool] = None,416 output_hidden_states: Optional[bool] = None,417 return_dict: Optional[bool] = None,418 cache_position: Optional[torch.LongTensor] = None,419 **flash_attn_kwargs: Unpack[FlashAttentionKwargs],420 ) -> Union[Tuple, BaseModelOutputWithPast]:421 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions422 output_hidden_states = (423 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states424 )425 use_cache = use_cache if use_cache is not None else self.config.use_cache426 return_dict = return_dict if return_dict is not None else self.config.use_return_dict427 428 if (input_ids is None) ^ (inputs_embeds is not None):429 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")430 431 if self.gradient_checkpointing and self.training and use_cache:432 logger.warning_once(433 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."434 )435 use_cache = False436 437 if inputs_embeds is None:438 inputs_embeds = self.embed_tokens(input_ids)439 440 if use_cache and past_key_values is None:441 past_key_values = DynamicCache()442 443 if cache_position is None:444 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0445 cache_position = torch.arange(446 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device447 )448 449 if position_ids is None:450 position_ids = cache_position.unsqueeze(0)451 452 causal_mask = self._update_causal_mask(453 attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions454 )455 456 hidden_states = inputs_embeds457 458 # create position embeddings to be shared across the decoder layers459 position_embeddings = self.rotary_emb(hidden_states, position_ids)460 461 # decoder layers462 all_hidden_states = () if output_hidden_states else None463 all_self_attns = () if output_attentions else None464 465 for decoder_layer in self.layers[: self.config.num_hidden_layers]:466 if output_hidden_states:467 all_hidden_states += (hidden_states,)468 469 if self.gradient_checkpointing and self.training:470 layer_outputs = self._gradient_checkpointing_func(471 decoder_layer.__call__,472 hidden_states,473 causal_mask,474 position_ids,475 past_key_values,476 output_attentions,477 use_cache,478 cache_position,479 position_embeddings,480 )481 else:482 layer_outputs = decoder_layer(483 hidden_states,484 attention_mask=causal_mask,485 position_ids=position_ids,486 past_key_value=past_key_values,487 output_attentions=output_attentions,488 use_cache=use_cache,489 cache_position=cache_position,490 position_embeddings=position_embeddings,491 **flash_attn_kwargs,492 )493 494 hidden_states = layer_outputs[0]495 496 if output_attentions:497 all_self_attns += (layer_outputs[1],)498 499 hidden_states = self.norm(hidden_states)500 501 # add hidden states from the last decoder layer502 if output_hidden_states:503 all_hidden_states += (hidden_states,)504 505 output = BaseModelOutputWithPast(506 last_hidden_state=hidden_states,507 past_key_values=past_key_values if use_cache else None,508 hidden_states=all_hidden_states,509 attentions=all_self_attns,510 )511 return output if return_dict else output.to_tuple()512 513 def _update_causal_mask(514 self,515 attention_mask: torch.Tensor,516 input_tensor: torch.Tensor,517 cache_position: torch.Tensor,518 past_key_values: Cache,519 output_attentions: bool,520 ):521 if self.config._attn_implementation == "flash_attention_2":522 if attention_mask is not None and (attention_mask == 0.0).any():523 return attention_mask524 return None525 526 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in527 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail528 # to infer the attention mask.529 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0530 using_static_cache = isinstance(past_key_values, StaticCache)531 532 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward533 if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:534 if AttentionMaskConverter._ignore_causal_mask_sdpa(535 attention_mask,536 inputs_embeds=input_tensor,537 past_key_values_length=past_seen_tokens,538 is_training=self.training,539 ):540 return None541 542 dtype, device = input_tensor.dtype, input_tensor.device543 sequence_length = input_tensor.shape[1]544 if using_static_cache:545 target_length = past_key_values.get_max_cache_shape()546 else:547 target_length = (548 attention_mask.shape[-1]549 if isinstance(attention_mask, torch.Tensor)550 else past_seen_tokens + sequence_length + 1551 )552 553 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).554 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(555 attention_mask,556 sequence_length=sequence_length,557 target_length=target_length,558 dtype=dtype,559 device=device,560 cache_position=cache_position,561 batch_size=input_tensor.shape[0],562 )563 564 if (565 self.config._attn_implementation == "sdpa"566 and attention_mask is not None567 and attention_mask.device.type == "cuda"568 and not output_attentions569 ):570 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when571 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.572 # Details: https://github.com/pytorch/pytorch/issues/110213573 min_dtype = torch.finfo(dtype).min574 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)575 576 return causal_mask577 578 @staticmethod579 def _prepare_4d_causal_attention_mask_with_cache_position(580 attention_mask: torch.Tensor,581 sequence_length: int,582 target_length: int,583 dtype: torch.dtype,584 device: torch.device,585 cache_position: torch.Tensor,586 batch_size: int,587 **kwargs,588 ):589 """590 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape591 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.592 593 Args:594 attention_mask (`torch.Tensor`):595 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape596 `(batch_size, 1, query_length, key_value_length)`.597 sequence_length (`int`):598 The sequence length being processed.599 target_length (`int`):600 The target length: when generating with static cache, the mask should be as long as the static cache,601 to account for the 0 padding, the part of the cache that is not filled yet.602 dtype (`torch.dtype`):603 The dtype to use for the 4D attention mask.604 device (`torch.device`):605 The device to plcae the 4D attention mask on.606 cache_position (`torch.Tensor`):607 Indices depicting the position of the input sequence tokens in the sequence.608 batch_size (`torch.Tensor`):609 Batch size.610 """611 if attention_mask is not None and attention_mask.dim() == 4:612 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.613 causal_mask = attention_mask614 else:615 min_dtype = torch.finfo(dtype).min616 causal_mask = torch.full(617 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device618 )619 if sequence_length != 1:620 causal_mask = torch.triu(causal_mask, diagonal=1)621 causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)622 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)623 if attention_mask is not None:624 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit625 mask_length = attention_mask.shape[-1]626 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :]627 padding_mask = padding_mask == 0628 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(629 padding_mask, min_dtype630 )631 632 return causal_mask633 634 635class KwargsForCausalLM(FlashAttentionKwargs): ...636 637 638class DummyLlamaForCausalLM(DummyLlamaPreTrainedModel, GenerationMixin):639 _tied_weights_keys = {"lm_head.weight": "model.embed_tokens.weight"}640 _tp_plan = {"lm_head": "colwise_rep"}641 642 def __init__(self, config):643 super().__init__(config)644 self.model = DummyLlamaModel(config)645 self.vocab_size = config.vocab_size646 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)647 648 # Initialize weights and apply final processing649 self.post_init()650 651 def get_input_embeddings(self):652 return self.model.embed_tokens653 654 def set_input_embeddings(self, value):655 self.model.embed_tokens = value656 657 def get_output_embeddings(self):658 return self.lm_head659 660 def set_output_embeddings(self, new_embeddings):661 self.lm_head = new_embeddings662 663 def set_decoder(self, decoder):664 self.model = decoder665 666 def get_decoder(self):667 return self.model668 669 def forward(670 self,671 input_ids: torch.LongTensor = None,672 attention_mask: Optional[torch.Tensor] = None,673 position_ids: Optional[torch.LongTensor] = None,674 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,675 inputs_embeds: Optional[torch.FloatTensor] = None,676 labels: Optional[torch.LongTensor] = None,677 use_cache: Optional[bool] = None,678 output_attentions: Optional[bool] = None,679 output_hidden_states: Optional[bool] = None,680 return_dict: Optional[bool] = None,681 cache_position: Optional[torch.LongTensor] = None,682 logits_to_keep: Union[int, torch.Tensor] = 0,683 **kwargs: Unpack[KwargsForCausalLM],684 ) -> Union[Tuple, CausalLMOutputWithPast]:685 r"""686 Args:687 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):688 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,689 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored690 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.691 692 logits_to_keep (`int` or `torch.Tensor`, *optional*):693 If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all694 `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that695 token can save memory, which becomes pretty significant for long sequences or large vocabulary size.696 If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.697 This is useful when using packed tensor format (single dimension for batch and sequence length).698 699 Returns:700 701 Example:702 703 ```python704 >>> from transformers import AutoTokenizer, DummyLlamaForCausalLM705 706 >>> model = DummyLlamaForCausalLM.from_pretrained("meta-llama/DummyLlama-2-7b-hf")707 >>> tokenizer = AutoTokenizer.from_pretrained("meta-llama/DummyLlama-2-7b-hf")708 709 >>> prompt = "Hey, are you conscious? Can you talk to me?"710 >>> inputs = tokenizer(prompt, return_tensors="pt")711 712 >>> # Generate713 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)714 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]715 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."716 ```"""717 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions718 output_hidden_states = (719 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states720 )721 return_dict = return_dict if return_dict is not None else self.config.use_return_dict722 723 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)724 outputs = self.model(725 input_ids=input_ids,726 attention_mask=attention_mask,727 position_ids=position_ids,728 past_key_values=past_key_values,729 inputs_embeds=inputs_embeds,730 use_cache=use_cache,731 output_attentions=output_attentions,732 output_hidden_states=output_hidden_states,733 return_dict=return_dict,734 cache_position=cache_position,735 **kwargs,736 )737 738 hidden_states = outputs[0]739 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss740 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep741 logits = self.lm_head(hidden_states[:, slice_indices, :])742 743 loss = None744 if labels is not None:745 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)746 747 if not return_dict:748 output = (logits,) + outputs[1:]749 return (loss,) + output if loss is not None else output750 751 return CausalLMOutputWithPast(752 loss=loss,753 logits=logits,754 past_key_values=outputs.past_key_values,755 hidden_states=outputs.hidden_states,756 attentions=outputs.attentions,757 )758 759 760class DummyLlamaForSequenceClassification(DummyLlamaPreTrainedModel):761 def __init__(self, config):762 super().__init__(config)763 self.num_labels = config.num_labels764 self.model = DummyLlamaModel(config)765 self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)766 767 # Initialize weights and apply final processing768 self.post_init()769 770 def get_input_embeddings(self):771 return self.model.embed_tokens772 773 def set_input_embeddings(self, value):774 self.model.embed_tokens = value775 776 def forward(777 self,778 input_ids: Optional[torch.LongTensor] = None,779 attention_mask: Optional[torch.Tensor] = None,780 position_ids: Optional[torch.LongTensor] = None,781 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,782 inputs_embeds: Optional[torch.FloatTensor] = None,783 labels: Optional[torch.LongTensor] = None,784 use_cache: Optional[bool] = None,785 output_attentions: Optional[bool] = None,786 output_hidden_states: Optional[bool] = None,787 return_dict: Optional[bool] = None,788 ) -> Union[Tuple, SequenceClassifierOutputWithPast]:789 r"""790 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):791 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,792 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If793 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).794 """795 return_dict = return_dict if return_dict is not None else self.config.use_return_dict796 797 transformer_outputs = self.model(798 input_ids,799 attention_mask=attention_mask,800 position_ids=position_ids,801 past_key_values=past_key_values,802 inputs_embeds=inputs_embeds,803 use_cache=use_cache,804 output_attentions=output_attentions,805 output_hidden_states=output_hidden_states,806 return_dict=return_dict,807 )808 hidden_states = transformer_outputs[0]809 logits = self.score(hidden_states)810 811 if input_ids is not None:812 batch_size = input_ids.shape[0]813 else:814 batch_size = inputs_embeds.shape[0]815 816 if self.config.pad_token_id is None and batch_size != 1:817 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")818 if self.config.pad_token_id is None:819 sequence_lengths = -1820 else:821 if input_ids is not None:822 # if no pad token found, use modulo instead of reverse indexing for ONNX compatibility823 sequence_lengths = torch.eq(input_ids, self.config.pad_token_id).int().argmax(-1) - 1824 sequence_lengths = sequence_lengths % input_ids.shape[-1]825 sequence_lengths = sequence_lengths.to(logits.device)826 else:827 sequence_lengths = -1828 829 pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]830 831 loss = None832 if labels is not None:833 loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)834 835 if not return_dict:836 output = (pooled_logits,) + transformer_outputs[1:]837 return ((loss,) + output) if loss is not None else output838 839 return SequenceClassifierOutputWithPast(840 loss=loss,841 logits=pooled_logits,842 past_key_values=transformer_outputs.past_key_values,843 hidden_states=transformer_outputs.hidden_states,844 attentions=transformer_outputs.attentions,845 )846 847 848class DummyLlamaForQuestionAnswering(DummyLlamaPreTrainedModel):849 base_model_prefix = "transformer"850 851 # Copied from transformers.models.bloom.modeling_bloom.BloomForQuestionAnswering.__init__ with Bloom->DummyLlama852 def __init__(self, config):853 super().__init__(config)854 self.transformer = DummyLlamaModel(config)855 self.qa_outputs = nn.Linear(config.hidden_size, 2)856 857 # Initialize weights and apply final processing858 self.post_init()859 860 def get_input_embeddings(self):861 return self.transformer.embed_tokens862 863 def set_input_embeddings(self, value):864 self.transformer.embed_tokens = value865 866 def forward(867 self,868 input_ids: Optional[torch.LongTensor] = None,869 attention_mask: Optional[torch.FloatTensor] = None,870 position_ids: Optional[torch.LongTensor] = None,871 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,872 inputs_embeds: Optional[torch.FloatTensor] = None,873 start_positions: Optional[torch.LongTensor] = None,874 end_positions: Optional[torch.LongTensor] = None,875 output_attentions: Optional[bool] = None,876 output_hidden_states: Optional[bool] = None,877 return_dict: Optional[bool] = None,878 **kwargs,879 ) -> Union[Tuple, QuestionAnsweringModelOutput]:880 r"""881 start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):882 Labels for position (index) of the start of the labelled span for computing the token classification loss.883 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence884 are not taken into account for computing the loss.885 end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):886 Labels for position (index) of the end of the labelled span for computing the token classification loss.887 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence888 are not taken into account for computing the loss.889 """890 return_dict = return_dict if return_dict is not None else self.config.use_return_dict891 892 outputs = self.transformer(893 input_ids,894 attention_mask=attention_mask,895 position_ids=position_ids,896 past_key_values=past_key_values,897 inputs_embeds=inputs_embeds,898 output_attentions=output_attentions,899 output_hidden_states=output_hidden_states,900 return_dict=return_dict,901 )902 903 sequence_output = outputs[0]904 905 logits = self.qa_outputs(sequence_output)906 start_logits, end_logits = logits.split(1, dim=-1)907 start_logits = start_logits.squeeze(-1).contiguous()908 end_logits = end_logits.squeeze(-1).contiguous()909 910 loss = None911 if start_positions is not None and end_positions is not None:912 loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs)913 914 if not return_dict:915 output = (start_logits, end_logits) + outputs[2:]916 return ((loss,) + output) if loss is not None else output917 918 return QuestionAnsweringModelOutput(919 loss=loss,920 start_logits=start_logits,921 end_logits=end_logits,922 hidden_states=outputs.hidden_states,923 attentions=outputs.attentions,924 )925 926 927class DummyLlamaForTokenClassification(DummyLlamaPreTrainedModel):928 def __init__(self, config):929 super().__init__(config)930 self.num_labels = config.num_labels931 self.model = DummyLlamaModel(config)932 if getattr(config, "classifier_dropout", None) is not None:933 classifier_dropout = config.classifier_dropout934 elif getattr(config, "hidden_dropout", None) is not None:935 classifier_dropout = config.hidden_dropout936 else:937 classifier_dropout = 0.1938 self.dropout = nn.Dropout(classifier_dropout)939 self.score = nn.Linear(config.hidden_size, config.num_labels)940 941 # Initialize weights and apply final processing942 self.post_init()943 944 def get_input_embeddings(self):945 return self.model.embed_tokens946 947 def set_input_embeddings(self, value):948 self.model.embed_tokens = value949 950 def forward(951 self,952 input_ids: Optional[torch.LongTensor] = None,953 attention_mask: Optional[torch.Tensor] = None,954 position_ids: Optional[torch.LongTensor] = None,955 past_key_values: Optional[List[torch.FloatTensor]] = None,956 inputs_embeds: Optional[torch.FloatTensor] = None,957 labels: Optional[torch.LongTensor] = None,958 use_cache: Optional[bool] = None,959 output_attentions: Optional[bool] = None,960 output_hidden_states: Optional[bool] = None,961 return_dict: Optional[bool] = None,962 ) -> Union[Tuple, TokenClassifierOutput]:963 r"""964 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):965 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,966 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If967 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).968 """969 return_dict = return_dict if return_dict is not None else self.config.use_return_dict970 971 outputs = self.model(972 input_ids,973 attention_mask=attention_mask,974 position_ids=position_ids,975 past_key_values=past_key_values,976 inputs_embeds=inputs_embeds,977 use_cache=use_cache,978 output_attentions=output_attentions,979 output_hidden_states=output_hidden_states,980 return_dict=return_dict,981 )982 sequence_output = outputs[0]983 sequence_output = self.dropout(sequence_output)984 logits = self.score(sequence_output)985 986 loss = None987 if labels is not None:988 loss = self.loss_function(logits, labels, self.config)989 990 if not return_dict:991 output = (logits,) + outputs[2:]992 return ((loss,) + output) if loss is not None else output993 994 return TokenClassifierOutput(995 loss=loss,996 logits=logits,997 hidden_states=outputs.hidden_states,998 attentions=outputs.attentions,999 )1000 1001 1002__all__ = [1003 "DummyLlamaForCausalLM",1004 "DummyLlamaModel",1005 "DummyLlamaPreTrainedModel",1006 "DummyLlamaForSequenceClassification",1007 "DummyLlamaForQuestionAnswering",1008 "DummyLlamaForTokenClassification",1009]1010