ServiceNow-AI/Apriel-5B-Base
42305
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_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel39from transformers.processing_utils import Unpack40from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS41from transformers.utils import (42 LossKwargs,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_apriel import AprielConfig51from .configuration_apriel import ROPE_INIT_FUNCTIONS52 53 54logger = logging.get_logger(__name__)55 56_CHECKPOINT_FOR_DOC = "ServiceNow-AI/Apriel-5B-Base"57_CONFIG_FOR_DOC = "AprielConfig"58 59 60class AprielRMSNorm(nn.Module):61 def __init__(self, hidden_size, eps=1e-6):62 """63 AprielRMSNorm is equivalent to T5LayerNorm64 """65 super().__init__()66 self.weight = nn.Parameter(torch.ones(hidden_size))67 self.variance_epsilon = eps68 69 def forward(self, hidden_states):70 input_dtype = hidden_states.dtype71 hidden_states = hidden_states.to(torch.float32)72 variance = hidden_states.pow(2).mean(-1, keepdim=True)73 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)74 return self.weight * hidden_states.to(input_dtype)75 76 def extra_repr(self):77 return f"{tuple(self.weight.shape)}, eps={self.variance_epsilon}"78 79 80ALL_LAYERNORM_LAYERS.append(AprielRMSNorm)81 82 83class AprielRotaryEmbedding(nn.Module):84 def __init__(self, config: AprielConfig, device=None):85 super().__init__()86 # BC: "rope_type" was originally "type"87 if hasattr(config, "rope_scaling") and config.rope_scaling is not None:88 self.rope_type = config.rope_scaling.get("rope_type", config.rope_scaling.get("type"))89 else:90 self.rope_type = "default"91 self.max_seq_len_cached = config.max_position_embeddings92 self.original_max_seq_len = config.max_position_embeddings93 94 self.config = config95 self.rope_init_fn = ROPE_INIT_FUNCTIONS[self.rope_type]96 97 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device)98 self.register_buffer("inv_freq", inv_freq, persistent=False)99 self.original_inv_freq = self.inv_freq100 101 def _dynamic_frequency_update(self, position_ids, device):102 """103 dynamic RoPE layers should recompute `inv_freq` in the following situations:104 1 - growing beyond the cached sequence length (allow scaling)105 2 - the current sequence length is in the original scale (avoid losing precision with small sequences)106 """107 seq_len = torch.max(position_ids) + 1108 if seq_len > self.max_seq_len_cached: # growth109 inv_freq, self.attention_scaling = self.rope_init_fn(self.config, device, seq_len=seq_len)110 self.register_buffer("inv_freq", inv_freq, persistent=False) # TODO joao: may break with compilation111 self.max_seq_len_cached = seq_len112 113 if seq_len < self.original_max_seq_len and self.max_seq_len_cached > self.original_max_seq_len: # reset114 # This .to() is needed if the model has been moved to a device after being initialized (because115 # the buffer is automatically moved, but not the original copy)116 self.original_inv_freq = self.original_inv_freq.to(device)117 self.register_buffer("inv_freq", self.original_inv_freq, persistent=False)118 self.max_seq_len_cached = self.original_max_seq_len119 120 @torch.no_grad()121 def forward(self, x, position_ids):122 if "dynamic" in self.rope_type:123 self._dynamic_frequency_update(position_ids, device=x.device)124 125 # Core RoPE block126 inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1)127 position_ids_expanded = position_ids[:, None, :].float()128 # Force float32 (see https://github.com/huggingface/transformers/pull/29285)129 device_type = x.device.type130 device_type = device_type if isinstance(device_type, str) and device_type != "mps" else "cpu"131 with torch.autocast(device_type=device_type, enabled=False):132 freqs = (inv_freq_expanded.float().to(x.device) @ position_ids_expanded.float()).transpose(1, 2)133 emb = torch.cat((freqs, freqs), dim=-1)134 cos = emb.cos()135 sin = emb.sin()136 137 # Advanced RoPE types (e.g. yarn) apply a post-processing scaling factor, equivalent to scaling attention138 cos = cos * self.attention_scaling139 sin = sin * self.attention_scaling140 141 return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype)142 143 144def rotate_half(x):145 """Rotates half the hidden dims of the input."""146 x1 = x[..., : x.shape[-1] // 2]147 x2 = x[..., x.shape[-1] // 2 :]148 return torch.cat((-x2, x1), dim=-1)149 150 151def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):152 """Applies Rotary Position Embedding to the query and key tensors.153 154 Args:155 q (`torch.Tensor`): The query tensor.156 k (`torch.Tensor`): The key tensor.157 cos (`torch.Tensor`): The cosine part of the rotary embedding.158 sin (`torch.Tensor`): The sine part of the rotary embedding.159 position_ids (`torch.Tensor`, *optional*):160 Deprecated and unused.161 unsqueeze_dim (`int`, *optional*, defaults to 1):162 The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and163 sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note164 that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and165 k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes166 cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have167 the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.168 Returns:169 `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.170 """171 cos = cos.unsqueeze(unsqueeze_dim)172 sin = sin.unsqueeze(unsqueeze_dim)173 q_embed = (q * cos) + (rotate_half(q) * sin)174 k_embed = (k * cos) + (rotate_half(k) * sin)175 return q_embed, k_embed176 177 178class AprielMLP(nn.Module):179 def __init__(self, config):180 super().__init__()181 self.config = config182 self.hidden_size = config.hidden_size183 self.intermediate_size = config.intermediate_size184 self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)185 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=config.mlp_bias)186 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=config.mlp_bias)187 self.act_fn = ACT2FN[config.hidden_act]188 189 def forward(self, x):190 down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))191 return down_proj192 193 194def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:195 """196 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,197 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)198 """199 batch, num_key_value_heads, slen, head_dim = hidden_states.shape200 if n_rep == 1:201 return hidden_states202 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)203 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)204 205 206def eager_attention_forward(207 module: nn.Module,208 query: torch.Tensor,209 key: torch.Tensor,210 value: torch.Tensor,211 attention_mask: Optional[torch.Tensor],212 scaling: float,213 dropout: float = 0.0,214 **kwargs,215):216 key_states = repeat_kv(key, module.num_key_value_groups)217 value_states = repeat_kv(value, module.num_key_value_groups)218 219 attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling220 if attention_mask is not None:221 causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]222 attn_weights = attn_weights + causal_mask223 224 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)225 attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)226 attn_output = torch.matmul(attn_weights, value_states)227 attn_output = attn_output.transpose(1, 2).contiguous()228 229 return attn_output, attn_weights230 231 232class AprielAttention(nn.Module):233 """Multi-headed attention from 'Attention Is All You Need' paper"""234 235 def __init__(self, config: AprielConfig, layer_idx: int):236 super().__init__()237 self.config = config238 self.layer_idx = layer_idx239 self.head_dim = getattr(config, "head_dim", config.hidden_size // config.num_attention_heads)240 self.num_key_value_groups = config.num_attention_heads // config.num_key_value_heads241 self.scaling = self.head_dim**-0.5242 self.attention_dropout = config.attention_dropout243 self.is_causal = True244 245 self.q_proj = nn.Linear(246 config.hidden_size, config.num_attention_heads * self.head_dim, bias=config.attention_bias247 )248 self.k_proj = nn.Linear(249 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias250 )251 self.v_proj = nn.Linear(252 config.hidden_size, config.num_key_value_heads * self.head_dim, bias=config.attention_bias253 )254 self.o_proj = nn.Linear(255 config.num_attention_heads * self.head_dim, config.hidden_size, bias=config.attention_bias256 )257 258 def forward(259 self,260 hidden_states: torch.Tensor,261 position_embeddings: Tuple[torch.Tensor, torch.Tensor],262 attention_mask: Optional[torch.Tensor],263 past_key_value: Optional[Cache] = None,264 cache_position: Optional[torch.LongTensor] = None,265 **kwargs: Unpack[FlashAttentionKwargs],266 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:267 input_shape = hidden_states.shape[:-1]268 hidden_shape = (*input_shape, -1, self.head_dim)269 270 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)271 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)272 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)273 274 cos, sin = position_embeddings275 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)276 277 if past_key_value is not None:278 # sin and cos are specific to RoPE models; cache_position needed for the static cache279 cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}280 key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)281 282 attention_interface: Callable = eager_attention_forward283 if self.config._attn_implementation != "eager":284 if self.config._attn_implementation == "sdpa" and kwargs.get("output_attentions", False):285 logger.warning_once(286 "`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. Falling back to "287 'eager attention. This warning can be removed using the argument `attn_implementation="eager"` when loading the model.'288 )289 else:290 attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]291 292 attn_output, attn_weights = attention_interface(293 self,294 query_states,295 key_states,296 value_states,297 attention_mask,298 dropout=0.0 if not self.training else self.attention_dropout,299 scaling=self.scaling,300 **kwargs,301 )302 303 attn_output = attn_output.reshape(*input_shape, -1).contiguous()304 attn_output = self.o_proj(attn_output)305 return attn_output, attn_weights306 307 308class AprielDecoderLayer(nn.Module):309 def __init__(self, config: AprielConfig, layer_idx: int):310 super().__init__()311 self.hidden_size = config.hidden_size312 313 self.self_attn = AprielAttention(config=config, layer_idx=layer_idx)314 315 self.mlp = AprielMLP(config)316 self.input_layernorm = AprielRMSNorm(config.hidden_size, eps=config.rms_norm_eps)317 self.post_attention_layernorm = AprielRMSNorm(config.hidden_size, eps=config.rms_norm_eps)318 319 def forward(320 self,321 hidden_states: torch.Tensor,322 attention_mask: Optional[torch.Tensor] = None,323 position_ids: Optional[torch.LongTensor] = None,324 past_key_value: Optional[Cache] = None,325 output_attentions: Optional[bool] = False,326 use_cache: Optional[bool] = False,327 cache_position: Optional[torch.LongTensor] = None,328 position_embeddings: Optional[Tuple[torch.Tensor, torch.Tensor]] = None, # necessary, but kept here for BC329 **kwargs: Unpack[FlashAttentionKwargs],330 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:331 residual = hidden_states332 333 hidden_states = self.input_layernorm(hidden_states)334 335 # Self Attention336 hidden_states, self_attn_weights = self.self_attn(337 hidden_states=hidden_states,338 attention_mask=attention_mask,339 position_ids=position_ids,340 past_key_value=past_key_value,341 output_attentions=output_attentions,342 use_cache=use_cache,343 cache_position=cache_position,344 position_embeddings=position_embeddings,345 **kwargs,346 )347 hidden_states = residual + hidden_states348 349 # Fully Connected350 residual = hidden_states351 hidden_states = self.post_attention_layernorm(hidden_states)352 hidden_states = self.mlp(hidden_states)353 hidden_states = residual + hidden_states354 355 outputs = (hidden_states,)356 if output_attentions:357 outputs += (self_attn_weights,)358 359 return outputs360 361 362APRIEL_START_DOCSTRING = r"""363 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the364 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads365 etc.)366 367 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.368 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage369 and behavior.370 371 Parameters:372 config ([`AprielConfig`]):373 Model configuration class with all the parameters of the model. Initializing with a config file does not374 load the weights associated with the model, only the configuration. Check out the375 [`~PreTrainedModel.from_pretrained`] method to load the model weights.376"""377 378 379@add_start_docstrings(380 "The bare Apriel Model outputting raw hidden-states without any specific head on top.",381 APRIEL_START_DOCSTRING,382)383class AprielPreTrainedModel(PreTrainedModel):384 config_class = AprielConfig385 base_model_prefix = "model"386 supports_gradient_checkpointing = True387 _no_split_modules = ["AprielDecoderLayer"]388 _skip_keys_device_placement = ["past_key_values"]389 _supports_flash_attn_2 = True390 _supports_sdpa = True391 _supports_flex_attn = True392 _supports_cache_class = True393 _supports_quantized_cache = True394 _supports_static_cache = True395 _supports_attention_backend = True396 397 def _init_weights(self, module):398 std = self.config.initializer_range399 if isinstance(module, nn.Linear):400 module.weight.data.normal_(mean=0.0, std=std)401 if module.bias is not None:402 module.bias.data.zero_()403 elif isinstance(module, nn.Embedding):404 module.weight.data.normal_(mean=0.0, std=std)405 if module.padding_idx is not None:406 module.weight.data[module.padding_idx].zero_()407 408 409APRIEL_INPUTS_DOCSTRING = r"""410 Args:411 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):412 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide413 it.414 415 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and416 [`PreTrainedTokenizer.__call__`] for details.417 418 [What are input IDs?](../glossary#input-ids)419 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):420 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:421 422 - 1 for tokens that are **not masked**,423 - 0 for tokens that are **masked**.424 425 [What are attention masks?](../glossary#attention-mask)426 427 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and428 [`PreTrainedTokenizer.__call__`] for details.429 430 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see431 `past_key_values`).432 433 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]434 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more435 information on the default strategy.436 437 - 1 indicates the head is **not masked**,438 - 0 indicates the head is **masked**.439 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):440 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,441 config.n_positions - 1]`.442 443 [What are position IDs?](../glossary#position-ids)444 past_key_values (`Cache` or `tuple(tuple(torch.FloatTensor))`, *optional*):445 Pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention446 blocks) that can be used to speed up sequential decoding. This typically consists in the `past_key_values`447 returned by the model at a previous stage of decoding, when `use_cache=True` or `config.use_cache=True`.448 449 Two formats are allowed:450 - a [`~cache_utils.Cache`] instance, see our451 [kv cache guide](https://huggingface.co/docs/transformers/en/kv_cache);452 - Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of453 shape `(batch_size, num_heads, sequence_length, embed_size_per_head)`). This is also known as the legacy454 cache format.455 456 The model will output the same cache format that is fed as input. If no `past_key_values` are passed, the457 legacy cache format will be returned.458 459 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't460 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`461 of shape `(batch_size, sequence_length)`.462 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):463 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This464 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the465 model's internal embedding lookup matrix.466 use_cache (`bool`, *optional*):467 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see468 `past_key_values`).469 output_attentions (`bool`, *optional*):470 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned471 tensors for more detail.472 output_hidden_states (`bool`, *optional*):473 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for474 more detail.475 return_dict (`bool`, *optional*):476 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.477 cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):478 Indices depicting the position of the input sequence tokens in the sequence. Contrarily to `position_ids`,479 this tensor is not affected by padding. It is used to update the cache in the correct position and to infer480 the complete sequence length.481"""482 483 484@add_start_docstrings(485 "The bare Apriel Model outputting raw hidden-states without any specific head on top.",486 APRIEL_START_DOCSTRING,487)488class AprielModel(AprielPreTrainedModel):489 """490 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`AprielDecoderLayer`]491 492 Args:493 config: AprielConfig494 """495 496 def __init__(self, config: AprielConfig):497 super().__init__(config)498 self.padding_idx = config.pad_token_id499 self.vocab_size = config.vocab_size500 501 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)502 self.layers = nn.ModuleList(503 [AprielDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]504 )505 self.norm = AprielRMSNorm(config.hidden_size, eps=config.rms_norm_eps)506 self.rotary_emb = AprielRotaryEmbedding(config=config)507 self.gradient_checkpointing = False508 509 # Initialize weights and apply final processing510 self.post_init()511 512 def get_input_embeddings(self):513 return self.embed_tokens514 515 def set_input_embeddings(self, value):516 self.embed_tokens = value517 518 @add_start_docstrings_to_model_forward(APRIEL_INPUTS_DOCSTRING)519 def forward(520 self,521 input_ids: torch.LongTensor = None,522 attention_mask: Optional[torch.Tensor] = None,523 position_ids: Optional[torch.LongTensor] = None,524 past_key_values: Optional[Cache] = None,525 inputs_embeds: Optional[torch.FloatTensor] = None,526 use_cache: Optional[bool] = None,527 output_attentions: Optional[bool] = None,528 output_hidden_states: Optional[bool] = None,529 return_dict: Optional[bool] = None,530 cache_position: Optional[torch.LongTensor] = None,531 **flash_attn_kwargs: Unpack[FlashAttentionKwargs],532 ) -> Union[Tuple, BaseModelOutputWithPast]:533 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions534 output_hidden_states = (535 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states536 )537 use_cache = use_cache if use_cache is not None else self.config.use_cache538 return_dict = return_dict if return_dict is not None else self.config.use_return_dict539 540 if (input_ids is None) ^ (inputs_embeds is not None):541 raise ValueError("You must specify exactly one of input_ids or inputs_embeds")542 543 if self.gradient_checkpointing and self.training and use_cache:544 logger.warning_once(545 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`."546 )547 use_cache = False548 549 if inputs_embeds is None:550 inputs_embeds = self.embed_tokens(input_ids)551 552 if use_cache and past_key_values is None:553 past_key_values = DynamicCache()554 555 if cache_position is None:556 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0557 cache_position = torch.arange(558 past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device559 )560 561 if position_ids is None:562 position_ids = cache_position.unsqueeze(0)563 564 causal_mask = self._update_causal_mask(565 attention_mask, inputs_embeds, cache_position, past_key_values, output_attentions566 )567 568 hidden_states = inputs_embeds569 570 # create position embeddings to be shared across the decoder layers571 position_embeddings = self.rotary_emb(hidden_states, position_ids)572 573 # decoder layers574 all_hidden_states = () if output_hidden_states else None575 all_self_attns = () if output_attentions else None576 577 for decoder_layer in self.layers[: self.config.num_hidden_layers]:578 if output_hidden_states:579 all_hidden_states += (hidden_states,)580 581 if self.gradient_checkpointing and self.training:582 layer_outputs = self._gradient_checkpointing_func(583 decoder_layer.__call__,584 hidden_states,585 causal_mask,586 position_ids,587 past_key_values,588 output_attentions,589 use_cache,590 cache_position,591 position_embeddings,592 )593 else:594 layer_outputs = decoder_layer(595 hidden_states,596 attention_mask=causal_mask,597 position_ids=position_ids,598 past_key_value=past_key_values,599 output_attentions=output_attentions,600 use_cache=use_cache,601 cache_position=cache_position,602 position_embeddings=position_embeddings,603 **flash_attn_kwargs,604 )605 606 hidden_states = layer_outputs[0]607 608 if output_attentions:609 all_self_attns += (layer_outputs[1],)610 611 hidden_states = self.norm(hidden_states)612 613 # add hidden states from the last decoder layer614 if output_hidden_states:615 all_hidden_states += (hidden_states,)616 617 output = BaseModelOutputWithPast(618 last_hidden_state=hidden_states,619 past_key_values=past_key_values if use_cache else None,620 hidden_states=all_hidden_states,621 attentions=all_self_attns,622 )623 return output if return_dict else output.to_tuple()624 625 def _update_causal_mask(626 self,627 attention_mask: torch.Tensor,628 input_tensor: torch.Tensor,629 cache_position: torch.Tensor,630 past_key_values: Cache,631 output_attentions: bool,632 ):633 if self.config._attn_implementation == "flash_attention_2":634 if attention_mask is not None and (attention_mask == 0.0).any():635 return attention_mask636 return None637 638 # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in639 # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail640 # to infer the attention mask.641 past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0642 using_static_cache = isinstance(past_key_values, StaticCache)643 644 # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward645 if self.config._attn_implementation == "sdpa" and not using_static_cache and not output_attentions:646 if AttentionMaskConverter._ignore_causal_mask_sdpa(647 attention_mask,648 inputs_embeds=input_tensor,649 past_key_values_length=past_seen_tokens,650 is_training=self.training,651 ):652 return None653 654 dtype, device = input_tensor.dtype, input_tensor.device655 sequence_length = input_tensor.shape[1]656 if using_static_cache:657 target_length = past_key_values.get_max_cache_shape()658 else:659 target_length = (660 attention_mask.shape[-1]661 if isinstance(attention_mask, torch.Tensor)662 else past_seen_tokens + sequence_length + 1663 )664 665 # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).666 causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(667 attention_mask,668 sequence_length=sequence_length,669 target_length=target_length,670 dtype=dtype,671 device=device,672 cache_position=cache_position,673 batch_size=input_tensor.shape[0],674 )675 676 if (677 self.config._attn_implementation == "sdpa"678 and attention_mask is not None679 and attention_mask.device.type in ["cuda", "xpu"]680 and not output_attentions681 ):682 # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when683 # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.684 # Details: https://github.com/pytorch/pytorch/issues/110213685 min_dtype = torch.finfo(dtype).min686 causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)687 688 return causal_mask689 690 @staticmethod691 def _prepare_4d_causal_attention_mask_with_cache_position(692 attention_mask: torch.Tensor,693 sequence_length: int,694 target_length: int,695 dtype: torch.dtype,696 device: torch.device,697 cache_position: torch.Tensor,698 batch_size: int,699 **kwargs,700 ):701 """702 Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape703 `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.704 705 Args:706 attention_mask (`torch.Tensor`):707 A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape708 `(batch_size, 1, query_length, key_value_length)`.709 sequence_length (`int`):710 The sequence length being processed.711 target_length (`int`):712 The target length: when generating with static cache, the mask should be as long as the static cache,713 to account for the 0 padding, the part of the cache that is not filled yet.714 dtype (`torch.dtype`):715 The dtype to use for the 4D attention mask.716 device (`torch.device`):717 The device to place the 4D attention mask on.718 cache_position (`torch.Tensor`):719 Indices depicting the position of the input sequence tokens in the sequence.720 batch_size (`torch.Tensor`):721 Batch size.722 """723 if attention_mask is not None and attention_mask.dim() == 4:724 # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.725 causal_mask = attention_mask726 else:727 min_dtype = torch.finfo(dtype).min728 causal_mask = torch.full(729 (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=device730 )731 if sequence_length != 1:732 causal_mask = torch.triu(causal_mask, diagonal=1)733 causal_mask *= torch.arange(target_length, device=device) > cache_position.reshape(-1, 1)734 causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)735 if attention_mask is not None:736 causal_mask = causal_mask.clone() # copy to contiguous memory for in-place edit737 mask_length = attention_mask.shape[-1]738 padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(739 causal_mask.device740 )741 padding_mask = padding_mask == 0742 causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(743 padding_mask, min_dtype744 )745 746 return causal_mask747 748 749class KwargsForCausalLM(FlashAttentionKwargs, LossKwargs): ...750 751 752class AprielForCausalLM(AprielPreTrainedModel, GenerationMixin):753 _tied_weights_keys = ["lm_head.weight"]754 _tp_plan = {"lm_head": "colwise_rep"}755 _pp_plan = {"lm_head": (["hidden_states"], ["logits"])}756 757 def __init__(self, config):758 super().__init__(config)759 self.model = AprielModel(config)760 self.vocab_size = config.vocab_size761 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)762 763 # Initialize weights and apply final processing764 self.post_init()765 766 def get_input_embeddings(self):767 return self.model.embed_tokens768 769 def set_input_embeddings(self, value):770 self.model.embed_tokens = value771 772 def get_output_embeddings(self):773 return self.lm_head774 775 def set_output_embeddings(self, new_embeddings):776 self.lm_head = new_embeddings777 778 def set_decoder(self, decoder):779 self.model = decoder780 781 def get_decoder(self):782 return self.model783 784 @deprecate_kwarg("num_logits_to_keep", version="4.50", new_name="logits_to_keep")785 @add_start_docstrings_to_model_forward(APRIEL_INPUTS_DOCSTRING)786 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)787 def forward(788 self,789 input_ids: torch.LongTensor = None,790 attention_mask: Optional[torch.Tensor] = None,791 position_ids: Optional[torch.LongTensor] = None,792 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,793 inputs_embeds: Optional[torch.FloatTensor] = None,794 labels: Optional[torch.LongTensor] = None,795 use_cache: Optional[bool] = None,796 output_attentions: Optional[bool] = None,797 output_hidden_states: Optional[bool] = None,798 return_dict: Optional[bool] = None,799 cache_position: Optional[torch.LongTensor] = None,800 logits_to_keep: Union[int, torch.Tensor] = 0,801 **kwargs: Unpack[KwargsForCausalLM],802 ) -> Union[Tuple, CausalLMOutputWithPast]:803 r"""804 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):805 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,806 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored807 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.808 809 logits_to_keep (`int` or `torch.Tensor`, *optional*):810 If an `int`, compute logits for the last `logits_to_keep` tokens. If `0`, calculate logits for all811 `input_ids` (special case). Only last token logits are needed for generation, and calculating them only for that812 token can save memory, which becomes pretty significant for long sequences or large vocabulary size.813 If a `torch.Tensor`, must be 1D corresponding to the indices to keep in the sequence length dimension.814 This is useful when using packed tensor format (single dimension for batch and sequence length).815 816 Returns:817 818 Example:819 820 ```python821 >>> from transformers import AutoTokenizer, AprielForCausalLM822 823 >>> model = AprielForCausalLM.from_pretrained("ServiceNow-AI/Apriel-5B-Base")824 >>> tokenizer = AutoTokenizer.from_pretrained("ServiceNow-AI/Apriel-5B-Base")825 826 >>> prompt = "Hey, are you conscious? Can you talk to me?"827 >>> inputs = tokenizer(prompt, return_tensors="pt")828 829 >>> # Generate830 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)831 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]832 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."833 ```"""834 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions835 output_hidden_states = (836 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states837 )838 return_dict = return_dict if return_dict is not None else self.config.use_return_dict839 840 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)841 outputs = self.model(842 input_ids=input_ids,843 attention_mask=attention_mask,844 position_ids=position_ids,845 past_key_values=past_key_values,846 inputs_embeds=inputs_embeds,847 use_cache=use_cache,848 output_attentions=output_attentions,849 output_hidden_states=output_hidden_states,850 return_dict=return_dict,851 cache_position=cache_position,852 **kwargs,853 )854 855 hidden_states = outputs[0]856 # Only compute necessary logits, and do not upcast them to float if we are not computing the loss857 slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep858 logits = self.lm_head(hidden_states[:, slice_indices, :])859 860 loss = None861 if labels is not None:862 loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)863 864 if not return_dict:865 output = (logits,) + outputs[1:]866 return (loss,) + output if loss is not None else output867 868 return CausalLMOutputWithPast(869 loss=loss,870 logits=logits,871 past_key_values=outputs.past_key_values,872 hidden_states=outputs.hidden_states,873 attentions=outputs.attentions,874 )875 876 877@add_start_docstrings(878 """879 The Apriel Model transformer with a sequence classification head on top (linear layer).880 881 [`AprielForSequenceClassification`] uses the last token in order to do the classification, as other causal models882 (e.g. GPT-2) do.883 884 Since it does classification on the last token, it requires to know the position of the last token. If a885 `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If886 no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the887 padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in888 each row of the batch).889 """,890 APRIEL_START_DOCSTRING,891)892class AprielForSequenceClassification(AprielPreTrainedModel):893 def __init__(self, config):894 super().__init__(config)895 self.num_labels = config.num_labels896 self.model = AprielModel(config)897 self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)898 899 # Initialize weights and apply final processing900 self.post_init()901 902 def get_input_embeddings(self):903 return self.model.embed_tokens904 905 def set_input_embeddings(self, value):906 self.model.embed_tokens = value907 908 @add_start_docstrings_to_model_forward(APRIEL_INPUTS_DOCSTRING)909 def forward(910 self,911 input_ids: Optional[torch.LongTensor] = None,912 attention_mask: Optional[torch.Tensor] = None,913 position_ids: Optional[torch.LongTensor] = None,914 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,915 inputs_embeds: Optional[torch.FloatTensor] = None,916 labels: Optional[torch.LongTensor] = None,917 use_cache: Optional[bool] = None,918 output_attentions: Optional[bool] = None,919 output_hidden_states: Optional[bool] = None,920 return_dict: Optional[bool] = None,921 ) -> Union[Tuple, SequenceClassifierOutputWithPast]:922 r"""923 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):924 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,925 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If926 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).927 """928 return_dict = return_dict if return_dict is not None else self.config.use_return_dict929 930 transformer_outputs = self.model(931 input_ids,932 attention_mask=attention_mask,933 position_ids=position_ids,934 past_key_values=past_key_values,935 inputs_embeds=inputs_embeds,936 use_cache=use_cache,937 output_attentions=output_attentions,938 output_hidden_states=output_hidden_states,939 return_dict=return_dict,940 )941 hidden_states = transformer_outputs[0]942 logits = self.score(hidden_states)943 944 if input_ids is not None:945 batch_size = input_ids.shape[0]946 else:947 batch_size = inputs_embeds.shape[0]948 949 if self.config.pad_token_id is None and batch_size != 1:950 raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")951 if self.config.pad_token_id is None:952 last_non_pad_token = -1953 elif input_ids is not None:954 # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id955 non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)956 token_indices = torch.arange(input_ids.shape[-1], device=logits.device)957 last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)958 else:959 last_non_pad_token = -1960 logger.warning_once(961 f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "962 "unexpected if using padding tokens in conjunction with `inputs_embeds.`"963 )964 965 pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]966 967 loss = None968 if labels is not None:969 loss = self.loss_function(logits=logits, labels=labels, pooled_logits=pooled_logits, config=self.config)970 971 if not return_dict:972 output = (pooled_logits,) + transformer_outputs[1:]973 return ((loss,) + output) if loss is not None else output974 975 return SequenceClassifierOutputWithPast(976 loss=loss,977 logits=pooled_logits,978 past_key_values=transformer_outputs.past_key_values,979 hidden_states=transformer_outputs.hidden_states,980 attentions=transformer_outputs.attentions,981 )982 983 984@add_start_docstrings(985 """986The Apriel Model transformer with a span classification head on top for extractive question-answering tasks like987SQuAD (a linear layer on top of the hidden-states output to compute `span start logits` and `span end logits`).988 """,989 APRIEL_START_DOCSTRING,990)991class AprielForQuestionAnswering(AprielPreTrainedModel):992 base_model_prefix = "transformer"993 994 def __init__(self, config):995 super().__init__(config)996 self.transformer = AprielModel(config)997 self.qa_outputs = nn.Linear(config.hidden_size, 2)998 999 # Initialize weights and apply final processing1000 self.post_init()1001 1002 def get_input_embeddings(self):1003 return self.transformer.embed_tokens1004 1005 def set_input_embeddings(self, value):1006 self.transformer.embed_tokens = value1007 1008 @add_start_docstrings_to_model_forward(APRIEL_INPUTS_DOCSTRING)1009 def forward(1010 self,1011 input_ids: Optional[torch.LongTensor] = None,1012 attention_mask: Optional[torch.FloatTensor] = None,1013 position_ids: Optional[torch.LongTensor] = None,1014 past_key_values: Optional[Union[Cache, List[torch.FloatTensor]]] = None,1015 inputs_embeds: Optional[torch.FloatTensor] = None,1016 start_positions: Optional[torch.LongTensor] = None,1017 end_positions: Optional[torch.LongTensor] = None,1018 output_attentions: Optional[bool] = None,1019 output_hidden_states: Optional[bool] = None,1020 return_dict: Optional[bool] = None,1021 **kwargs,1022 ) -> Union[Tuple, QuestionAnsweringModelOutput]:1023 r"""1024 start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1025 Labels for position (index) of the start of the labelled span for computing the token classification loss.1026 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1027 are not taken into account for computing the loss.1028 end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1029 Labels for position (index) of the end of the labelled span for computing the token classification loss.1030 Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence1031 are not taken into account for computing the loss.1032 """1033 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1034 1035 outputs = self.transformer(1036 input_ids,1037 attention_mask=attention_mask,1038 position_ids=position_ids,1039 past_key_values=past_key_values,1040 inputs_embeds=inputs_embeds,1041 output_attentions=output_attentions,1042 output_hidden_states=output_hidden_states,1043 return_dict=return_dict,1044 )1045 1046 sequence_output = outputs[0]1047 1048 logits = self.qa_outputs(sequence_output)1049 start_logits, end_logits = logits.split(1, dim=-1)1050 start_logits = start_logits.squeeze(-1).contiguous()1051 end_logits = end_logits.squeeze(-1).contiguous()1052 1053 loss = None1054 if start_positions is not None and end_positions is not None:1055 loss = self.loss_function(start_logits, end_logits, start_positions, end_positions, **kwargs)1056 1057 if not return_dict:1058 output = (start_logits, end_logits) + outputs[2:]1059 return ((loss,) + output) if loss is not None else output1060 1061 return QuestionAnsweringModelOutput(1062 loss=loss,1063 start_logits=start_logits,1064 end_logits=end_logits,1065 hidden_states=outputs.hidden_states,1066 attentions=outputs.attentions,1067 )1068 1069 1070@add_start_docstrings(1071 """1072 The Apriel Model transformer with a token classification head on top (a linear layer on top of the hidden-states1073 output) e.g. for Named-Entity-Recognition (NER) tasks.1074 """,1075 APRIEL_START_DOCSTRING,1076)1077class AprielForTokenClassification(AprielPreTrainedModel):1078 def __init__(self, config):1079 super().__init__(config)1080 self.num_labels = config.num_labels1081 self.model = AprielModel(config)1082 if getattr(config, "classifier_dropout", None) is not None:1083 classifier_dropout = config.classifier_dropout1084 elif getattr(config, "hidden_dropout", None) is not None:1085 classifier_dropout = config.hidden_dropout1086 else:1087 classifier_dropout = 0.11088 self.dropout = nn.Dropout(classifier_dropout)1089 self.score = nn.Linear(config.hidden_size, config.num_labels)1090 1091 # Initialize weights and apply final processing1092 self.post_init()1093 1094 def get_input_embeddings(self):1095 return self.model.embed_tokens1096 1097 def set_input_embeddings(self, value):1098 self.model.embed_tokens = value1099 1100 @add_start_docstrings_to_model_forward(APRIEL_INPUTS_DOCSTRING)1101 @add_code_sample_docstrings(1102 checkpoint=_CHECKPOINT_FOR_DOC,1103 output_type=TokenClassifierOutput,1104 config_class=_CONFIG_FOR_DOC,1105 )1106 def forward(1107 self,1108 input_ids: Optional[torch.LongTensor] = None,1109 attention_mask: Optional[torch.Tensor] = None,1110 position_ids: Optional[torch.LongTensor] = None,1111 past_key_values: Optional[List[torch.FloatTensor]] = None,1112 inputs_embeds: Optional[torch.FloatTensor] = None,1113 labels: Optional[torch.LongTensor] = None,1114 use_cache: Optional[bool] = None,1115 output_attentions: Optional[bool] = None,1116 output_hidden_states: Optional[bool] = None,1117 return_dict: Optional[bool] = None,1118 ) -> Union[Tuple, TokenClassifierOutput]:1119 r"""1120 labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):1121 Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,1122 config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If1123 `config.num_labels > 1` a classification loss is computed (Cross-Entropy).1124 """1125 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1126 1127 outputs = self.model(1128 input_ids,1129 attention_mask=attention_mask,1130 position_ids=position_ids,1131 past_key_values=past_key_values,1132 inputs_embeds=inputs_embeds,1133 use_cache=use_cache,1134 output_attentions=output_attentions,1135 output_hidden_states=output_hidden_states,1136 return_dict=return_dict,1137 )1138 sequence_output = outputs[0]1139 sequence_output = self.dropout(sequence_output)1140 logits = self.score(sequence_output)1141 1142 loss = None1143 if labels is not None:1144 loss = self.loss_function(logits, labels, self.config)1145 1146 if not return_dict:1147 output = (logits,) + outputs[2:]1148 return ((loss,) + output) if loss is not None else output1149 1150 return TokenClassifierOutput(1151 loss=loss,1152 logits=logits,1153 hidden_states=outputs.hidden_states,1154 attentions=outputs.attentions,1155 )1156 1157 1158__all__ = [1159 "AprielForCausalLM",1160 "AprielModel",1161 "AprielPreTrainedModel",1162 "AprielForSequenceClassification",1163 "AprielForQuestionAnswering",1164 "AprielForTokenClassification",1165]1166