Open-Foundation-Models/PolyNorm_1B
017
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.20""" PyTorch LLaMA model."""21import math22from typing import List, Optional, Tuple, Union23 24import torch25import torch.nn.functional as F26import torch.utils.checkpoint27from torch import nn28from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss29 30from transformers.activations import ACT2FN31from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast32from transformers.modeling_utils import PreTrainedModel33from transformers.pytorch_utils import ALL_LAYERNORM_LAYERS34from transformers.utils import (35 add_start_docstrings,36 add_start_docstrings_to_model_forward,37 # is_flash_attn_available,38 logging,39 replace_return_docstrings,40)41from .configuration_polyllama import PolyLlamaConfig42 43# if is_flash_attn_available():44# from flash_attn import flash_attn_func, flash_attn_varlen_func45# from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input # noqa46 47 48logger = logging.get_logger(__name__)49 50_CONFIG_FOR_DOC = "PolyLlamaConfig"51 52 53def _get_unpad_data(padding_mask):54 seqlens_in_batch = padding_mask.sum(dim=-1, dtype=torch.int32)55 indices = torch.nonzero(padding_mask.flatten(), as_tuple=False).flatten()56 max_seqlen_in_batch = seqlens_in_batch.max().item()57 cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.torch.int32), (1, 0))58 return (59 indices,60 cu_seqlens,61 max_seqlen_in_batch,62 )63 64 65# Copied from transformers.models.bart.modeling_bart._make_causal_mask66def _make_causal_mask(67 input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device, past_key_values_length: int = 068):69 """70 Make causal mask used for bi-directional self-attention.71 """72 bsz, tgt_len = input_ids_shape73 mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)74 mask_cond = torch.arange(mask.size(-1), device=device)75 mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)76 mask = mask.to(dtype)77 78 if past_key_values_length > 0:79 mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)80 return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)81 82 83# Copied from transformers.models.bart.modeling_bart._expand_mask84def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):85 """86 Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.87 """88 bsz, src_len = mask.size()89 tgt_len = tgt_len if tgt_len is not None else src_len90 91 expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)92 93 inverted_mask = 1.0 - expanded_mask94 95 return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)96 97 98class LlamaRMSNorm(nn.Module):99 def __init__(self, hidden_size, eps=1e-6):100 """101 LlamaRMSNorm is equivalent to T5LayerNorm102 """103 super().__init__()104 self.weight = nn.Parameter(torch.ones(hidden_size))105 self.variance_epsilon = eps106 107 def forward(self, hidden_states):108 input_dtype = hidden_states.dtype109 hidden_states = hidden_states.to(torch.float32)110 variance = hidden_states.pow(2).mean(-1, keepdim=True)111 hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)112 return self.weight * hidden_states.to(input_dtype)113 114 115ALL_LAYERNORM_LAYERS.append(LlamaRMSNorm)116 117 118class LlamaRotaryEmbedding(nn.Module):119 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):120 super().__init__()121 122 self.dim = dim123 self.max_position_embeddings = max_position_embeddings124 self.base = base125 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))126 self.register_buffer("inv_freq", inv_freq, persistent=False)127 128 # Build here to make `torch.jit.trace` work.129 self._set_cos_sin_cache(130 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()131 )132 133 def _set_cos_sin_cache(self, seq_len, device, dtype):134 self.max_seq_len_cached = seq_len135 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)136 137 freqs = torch.einsum("i,j->ij", t, self.inv_freq)138 # Different from paper, but it uses a different permutation in order to obtain the same calculation139 emb = torch.cat((freqs, freqs), dim=-1)140 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)141 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)142 143 def forward(self, x, seq_len=None):144 # x: [bs, num_attention_heads, seq_len, head_size]145 if seq_len > self.max_seq_len_cached:146 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)147 148 return (149 self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),150 self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),151 )152 153 154class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):155 """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""156 157 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):158 self.scaling_factor = scaling_factor159 super().__init__(dim, max_position_embeddings, base, device)160 161 def _set_cos_sin_cache(self, seq_len, device, dtype):162 self.max_seq_len_cached = seq_len163 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)164 t = t / self.scaling_factor165 166 freqs = torch.einsum("i,j->ij", t, self.inv_freq)167 # Different from paper, but it uses a different permutation in order to obtain the same calculation168 emb = torch.cat((freqs, freqs), dim=-1)169 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)170 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)171 172 173class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):174 """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""175 176 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):177 self.scaling_factor = scaling_factor178 super().__init__(dim, max_position_embeddings, base, device)179 180 def _set_cos_sin_cache(self, seq_len, device, dtype):181 self.max_seq_len_cached = seq_len182 183 if seq_len > self.max_position_embeddings:184 base = self.base * (185 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)186 ) ** (self.dim / (self.dim - 2))187 inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))188 self.register_buffer("inv_freq", inv_freq, persistent=False)189 190 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)191 192 freqs = torch.einsum("i,j->ij", t, self.inv_freq)193 # Different from paper, but it uses a different permutation in order to obtain the same calculation194 emb = torch.cat((freqs, freqs), dim=-1)195 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)196 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)197 198 199def rotate_half(x):200 """Rotates half the hidden dims of the input."""201 x1 = x[..., : x.shape[-1] // 2]202 x2 = x[..., x.shape[-1] // 2 :]203 return torch.cat((-x2, x1), dim=-1)204 205 206def apply_rotary_pos_emb(q, k, cos, sin, position_ids):207 # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.208 cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]209 sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]210 cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]211 sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]212 q_embed = (q * cos) + (rotate_half(q) * sin)213 k_embed = (k * cos) + (rotate_half(k) * sin)214 return q_embed, k_embed215 216 217def norm(x, eps: float = 1e-6):218 return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + eps)219 220 221def polynorm(x,weight,bias):222 return weight[0] * norm(x**3) + weight[1] * norm(x**2) + weight[2] * norm(x) + bias223 224 225def poly(x,weight,bias):226 return weight[0] * (x**3) + weight[1] * (x**2) + weight[2] * (x) + bias227 228 229 230class ACT_PolyNorm(nn.Module):231 def __init__(self, inplace: bool = False):232 super(ACT_PolyNorm, self).__init__()233 self.weight = nn.Parameter(torch.ones(3)/3)234 self.bias = nn.Parameter(torch.zeros(1))235 236 def forward(self,x):237 return polynorm(x,self.weight,self.bias)238 239 240 241class ACT_PolyReLU(nn.Module):242 def __init__(self, inplace: bool = False):243 super(ACT_PolyReLU, self).__init__()244 self.weight = nn.Parameter(torch.ones(3)/3)245 self.bias = nn.Parameter(torch.zeros(1))246 247 def forward(self,x):248 return poly(F.relu(x),self.weight,self.bias)249 250ACT_POLY = {251 "PolyNorm": ACT_PolyNorm,252 "PolyReLU": ACT_PolyReLU,253}254 255class PolyLlamaMLP(nn.Module):256 def __init__(self, config):257 super().__init__()258 self.config = config259 self.hidden_size = config.hidden_size260 self.intermediate_size = config.intermediate_size261 # self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)262 self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)263 self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)264 self.act_fn = ACT_POLY[config.hidden_act]() if config.hidden_act in ACT_POLY else ACT2FN[config.hidden_act]265 266 def forward(self, x):267 if self.config.pretraining_tp > 1:268 slice = self.intermediate_size // self.config.pretraining_tp269 # gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)270 up_proj_slices = self.up_proj.weight.split(slice, dim=0)271 down_proj_slices = self.down_proj.weight.split(slice, dim=1)272 273 up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)274 275 # intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)276 intermediate_states = self.act_fn(up_proj).split(slice, dim=2)277 down_proj = [F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.pretraining_tp)]278 down_proj = sum(down_proj)279 else:280 # down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))281 down_proj = self.down_proj(self.act_fn(self.up_proj(x)))282 283 return down_proj284 285 286def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:287 """288 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,289 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)290 """291 batch, num_key_value_heads, slen, head_dim = hidden_states.shape292 if n_rep == 1:293 return hidden_states294 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)295 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)296 297 298class LlamaAttention(nn.Module):299 """Multi-headed attention from 'Attention Is All You Need' paper"""300 301 def __init__(self, config: PolyLlamaConfig):302 super().__init__()303 self.config = config304 self.hidden_size = config.hidden_size305 self.num_heads = config.num_attention_heads306 self.head_dim = self.hidden_size // self.num_heads307 self.num_key_value_heads = config.num_key_value_heads308 self.num_key_value_groups = self.num_heads // self.num_key_value_heads309 self.max_position_embeddings = config.max_position_embeddings310 self.rope_theta = config.rope_theta311 312 if (self.head_dim * self.num_heads) != self.hidden_size:313 raise ValueError(314 f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"315 f" and `num_heads`: {self.num_heads})."316 )317 self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=config.attention_bias)318 self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)319 self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=config.attention_bias)320 self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=config.attention_bias)321 self._init_rope()322 323 def _init_rope(self):324 if self.config.rope_scaling is None:325 self.rotary_emb = LlamaRotaryEmbedding(326 self.head_dim,327 max_position_embeddings=self.max_position_embeddings,328 base=self.rope_theta,329 )330 else:331 scaling_type = self.config.rope_scaling["type"]332 scaling_factor = self.config.rope_scaling["factor"]333 if scaling_type == "linear":334 self.rotary_emb = LlamaLinearScalingRotaryEmbedding(335 self.head_dim,336 max_position_embeddings=self.max_position_embeddings,337 scaling_factor=scaling_factor,338 base=self.rope_theta,339 )340 elif scaling_type == "dynamic":341 self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(342 self.head_dim,343 max_position_embeddings=self.max_position_embeddings,344 scaling_factor=scaling_factor,345 base=self.rope_theta,346 )347 else:348 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")349 350 def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):351 return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()352 353 def forward(354 self,355 hidden_states: torch.Tensor,356 attention_mask: Optional[torch.Tensor] = None,357 position_ids: Optional[torch.LongTensor] = None,358 past_key_value: Optional[Tuple[torch.Tensor]] = None,359 output_attentions: bool = False,360 use_cache: bool = False,361 padding_mask: Optional[torch.LongTensor] = None,362 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:363 bsz, q_len, _ = hidden_states.size()364 365 if self.config.pretraining_tp > 1:366 key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp367 query_slices = self.q_proj.weight.split(368 (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0369 )370 key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)371 value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)372 373 query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]374 query_states = torch.cat(query_states, dim=-1)375 376 key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]377 key_states = torch.cat(key_states, dim=-1)378 379 value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]380 value_states = torch.cat(value_states, dim=-1)381 382 else:383 query_states = self.q_proj(hidden_states)384 key_states = self.k_proj(hidden_states)385 value_states = self.v_proj(hidden_states)386 387 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)388 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)389 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)390 391 kv_seq_len = key_states.shape[-2]392 if past_key_value is not None:393 kv_seq_len += past_key_value[0].shape[-2]394 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)395 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)396 397 if past_key_value is not None:398 # reuse k, v, self_attention399 key_states = torch.cat([past_key_value[0], key_states], dim=2)400 value_states = torch.cat([past_key_value[1], value_states], dim=2)401 402 past_key_value = (key_states, value_states) if use_cache else None403 404 key_states = repeat_kv(key_states, self.num_key_value_groups)405 value_states = repeat_kv(value_states, self.num_key_value_groups)406 407 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)408 409 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):410 raise ValueError(411 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"412 f" {attn_weights.size()}"413 )414 415 if attention_mask is not None:416 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):417 raise ValueError(418 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"419 )420 attn_weights = attn_weights + attention_mask421 422 # upcast attention to fp32423 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)424 attn_output = torch.matmul(attn_weights, value_states)425 426 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):427 raise ValueError(428 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"429 f" {attn_output.size()}"430 )431 432 attn_output = attn_output.transpose(1, 2).contiguous()433 434 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)435 436 if self.config.pretraining_tp > 1:437 attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)438 o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)439 attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])440 else:441 attn_output = self.o_proj(attn_output)442 443 if not output_attentions:444 attn_weights = None445 446 return attn_output, attn_weights, past_key_value447 448 449class LlamaFlashAttention2(LlamaAttention):450 """451 Llama flash attention module. This module inherits from `LlamaAttention` as the weights of the module stays452 untouched. The only required change would be on the forward pass where it needs to correctly call the public API of453 flash attention and deal with padding tokens in case the input contains any of them.454 """455 456 def forward(457 self,458 hidden_states: torch.Tensor,459 attention_mask: Optional[torch.Tensor] = None,460 position_ids: Optional[torch.LongTensor] = None,461 past_key_value: Optional[Tuple[torch.Tensor]] = None,462 output_attentions: bool = False,463 use_cache: bool = False,464 padding_mask: Optional[torch.LongTensor] = None,465 ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:466 # LlamaFlashAttention2 attention does not support output_attentions467 output_attentions = False468 469 bsz, q_len, _ = hidden_states.size()470 471 query_states = self.q_proj(hidden_states)472 key_states = self.k_proj(hidden_states)473 value_states = self.v_proj(hidden_states)474 475 # Flash attention requires the input to have the shape476 # batch_size x seq_length x head_dime x hidden_dim477 # therefore we just need to keep the original shape478 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)479 key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)480 value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)481 482 kv_seq_len = key_states.shape[-2]483 if past_key_value is not None:484 kv_seq_len += past_key_value[0].shape[-2]485 486 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)487 488 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)489 490 if past_key_value is not None:491 # reuse k, v, self_attention492 key_states = torch.cat([past_key_value[0], key_states], dim=2)493 value_states = torch.cat([past_key_value[1], value_states], dim=2)494 495 past_key_value = (key_states, value_states) if use_cache else None496 497 query_states = query_states.transpose(1, 2)498 key_states = key_states.transpose(1, 2)499 value_states = value_states.transpose(1, 2)500 501 # TODO: llama does not have dropout in the config??502 # It is recommended to use dropout with FA according to the docs503 # when training.504 dropout_rate = 0.0 # if not self.training else self.attn_dropout505 506 # In PEFT, usually we cast the layer norms in float32 for training stability reasons507 # therefore the input hidden states gets silently casted in float32. Hence, we need508 # cast them back in float16 just to be sure everything works as expected.509 # This might slowdown training & inference so it is recommended to not cast the LayerNorms510 # in fp32. (LlamaRMSNorm handles it correctly)511 input_dtype = query_states.dtype512 if input_dtype == torch.float32:513 logger.warning_once(514 "The input hidden states seems to be silently casted in float32, this might be related to"515 " the fact you have upcasted embedding or layer norm layers in float32. We will cast back the input in"516 " float16."517 )518 519 query_states = query_states.to(torch.float16)520 key_states = key_states.to(torch.float16)521 value_states = value_states.to(torch.float16)522 523 attn_output = self._flash_attention_forward(524 query_states, key_states, value_states, padding_mask, q_len, dropout=dropout_rate525 )526 527 attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()528 attn_output = self.o_proj(attn_output)529 530 if not output_attentions:531 attn_weights = None532 533 return attn_output, attn_weights, past_key_value534 535 def _flash_attention_forward(536 self, query_states, key_states, value_states, padding_mask, query_length, dropout=0.0, softmax_scale=None537 ):538 """539 Calls the forward method of Flash Attention - if the input hidden states contain at least one padding token540 first unpad the input, then computes the attention scores and pad the final attention scores.541 542 Args:543 query_states (`torch.Tensor`):544 Input query states to be passed to Flash Attention API545 key_states (`torch.Tensor`):546 Input key states to be passed to Flash Attention API547 value_states (`torch.Tensor`):548 Input value states to be passed to Flash Attention API549 padding_mask (`torch.Tensor`):550 The padding mask - corresponds to a tensor of size `(batch_size, seq_len)` where 0 stands for the551 position of padding tokens and 1 for the position of non-padding tokens.552 dropout (`int`, *optional*):553 Attention dropout554 softmax_scale (`float`, *optional*):555 The scaling of QK^T before applying softmax. Default to 1 / sqrt(head_dim)556 """557 # Contains at least one padding token in the sequence558 if padding_mask is not None:559 batch_size = query_states.shape[0]560 query_states, key_states, value_states, indices_q, cu_seq_lens, max_seq_lens = self._upad_input(561 query_states, key_states, value_states, padding_mask, query_length562 )563 564 cu_seqlens_q, cu_seqlens_k = cu_seq_lens565 max_seqlen_in_batch_q, max_seqlen_in_batch_k = max_seq_lens566 567 attn_output_unpad = flash_attn_varlen_func(568 query_states,569 key_states,570 value_states,571 cu_seqlens_q=cu_seqlens_q,572 cu_seqlens_k=cu_seqlens_k,573 max_seqlen_q=max_seqlen_in_batch_q,574 max_seqlen_k=max_seqlen_in_batch_k,575 dropout_p=dropout,576 softmax_scale=softmax_scale,577 causal=True,578 )579 580 attn_output = pad_input(attn_output_unpad, indices_q, batch_size, query_length)581 else:582 attn_output = flash_attn_func(583 query_states, key_states, value_states, dropout, softmax_scale=softmax_scale, causal=True584 )585 586 return attn_output587 588 def _upad_input(self, query_layer, key_layer, value_layer, padding_mask, query_length):589 indices_k, cu_seqlens_k, max_seqlen_in_batch_k = _get_unpad_data(padding_mask)590 batch_size, kv_seq_len, num_key_value_heads, head_dim = key_layer.shape591 592 key_layer = index_first_axis(593 key_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k594 )595 value_layer = index_first_axis(596 value_layer.reshape(batch_size * kv_seq_len, num_key_value_heads, head_dim), indices_k597 )598 if query_length == kv_seq_len:599 query_layer = index_first_axis(600 query_layer.reshape(batch_size * kv_seq_len, self.num_heads, head_dim), indices_k601 )602 cu_seqlens_q = cu_seqlens_k603 max_seqlen_in_batch_q = max_seqlen_in_batch_k604 indices_q = indices_k605 elif query_length == 1:606 max_seqlen_in_batch_q = 1607 cu_seqlens_q = torch.arange(608 batch_size + 1, dtype=torch.int32, device=query_layer.device609 ) # There is a memcpy here, that is very bad.610 indices_q = cu_seqlens_q[:-1]611 query_layer = query_layer.squeeze(1)612 else:613 # The -q_len: slice assumes left padding.614 padding_mask = padding_mask[:, -query_length:]615 query_layer, indices_q, cu_seqlens_q, max_seqlen_in_batch_q = unpad_input(query_layer, padding_mask)616 617 return (618 query_layer,619 key_layer,620 value_layer,621 indices_q,622 (cu_seqlens_q, cu_seqlens_k),623 (max_seqlen_in_batch_q, max_seqlen_in_batch_k),624 )625 626 627class PolyLlamaDecoderLayer(nn.Module):628 def __init__(self, config: PolyLlamaConfig):629 super().__init__()630 self.hidden_size = config.hidden_size631 self.self_attn = (632 LlamaAttention(config=config)633 if not getattr(config, "_flash_attn_2_enabled", False)634 else LlamaFlashAttention2(config=config)635 )636 self.mlp = PolyLlamaMLP(config)637 self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)638 self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)639 640 def forward(641 self,642 hidden_states: torch.Tensor,643 attention_mask: Optional[torch.Tensor] = None,644 position_ids: Optional[torch.LongTensor] = None,645 past_key_value: Optional[Tuple[torch.Tensor]] = None,646 output_attentions: Optional[bool] = False,647 use_cache: Optional[bool] = False,648 padding_mask: Optional[torch.LongTensor] = None,649 ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:650 """651 Args:652 hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`653 attention_mask (`torch.FloatTensor`, *optional*): attention mask of size654 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.655 output_attentions (`bool`, *optional*):656 Whether or not to return the attentions tensors of all attention layers. See `attentions` under657 returned tensors for more detail.658 use_cache (`bool`, *optional*):659 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding660 (see `past_key_values`).661 past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states662 """663 664 residual = hidden_states665 666 hidden_states = self.input_layernorm(hidden_states)667 668 # Self Attention669 hidden_states, self_attn_weights, present_key_value = self.self_attn(670 hidden_states=hidden_states,671 attention_mask=attention_mask,672 position_ids=position_ids,673 past_key_value=past_key_value,674 output_attentions=output_attentions,675 use_cache=use_cache,676 padding_mask=padding_mask,677 )678 hidden_states = residual + hidden_states679 680 # Fully Connected681 residual = hidden_states682 hidden_states = self.post_attention_layernorm(hidden_states)683 hidden_states = self.mlp(hidden_states)684 hidden_states = residual + hidden_states685 686 outputs = (hidden_states,)687 688 if output_attentions:689 outputs += (self_attn_weights,)690 691 if use_cache:692 outputs += (present_key_value,)693 694 return outputs695 696 697LLAMA_START_DOCSTRING = r"""698 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the699 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads700 etc.)701 702 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.703 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage704 and behavior.705 706 Parameters:707 config ([`LlamaConfig`]):708 Model configuration class with all the parameters of the model. Initializing with a config file does not709 load the weights associated with the model, only the configuration. Check out the710 [`~PreTrainedModel.from_pretrained`] method to load the model weights.711"""712 713 714@add_start_docstrings(715 "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",716 LLAMA_START_DOCSTRING,717)718class PolyLlamaPreTrainedModel(PreTrainedModel):719 config_class = PolyLlamaConfig720 base_model_prefix = "model"721 supports_gradient_checkpointing = True722 _no_split_modules = ["PolyLlamaDecoderLayer"]723 _skip_keys_device_placement = "past_key_values"724 _supports_flash_attn_2 = True725 726 def _init_weights(self, module):727 std = self.config.initializer_range728 if isinstance(module, nn.Linear):729 module.weight.data.normal_(mean=0.0, std=std)730 if module.bias is not None:731 module.bias.data.zero_()732 elif isinstance(module, nn.Embedding):733 module.weight.data.normal_(mean=0.0, std=std)734 if module.padding_idx is not None:735 module.weight.data[module.padding_idx].zero_()736 737 def _set_gradient_checkpointing(self, module, value=False):738 if isinstance(module, PolyLlamaModel):739 module.gradient_checkpointing = value740 741 742LLAMA_INPUTS_DOCSTRING = r"""743 Args:744 input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):745 Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide746 it.747 748 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and749 [`PreTrainedTokenizer.__call__`] for details.750 751 [What are input IDs?](../glossary#input-ids)752 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):753 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:754 755 - 1 for tokens that are **not masked**,756 - 0 for tokens that are **masked**.757 758 [What are attention masks?](../glossary#attention-mask)759 760 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and761 [`PreTrainedTokenizer.__call__`] for details.762 763 If `past_key_values` is used, optionally only the last `input_ids` have to be input (see764 `past_key_values`).765 766 If you want to change padding behavior, you should read [`modeling_opt._prepare_decoder_attention_mask`]767 and modify to your needs. See diagram 1 in [the paper](https://arxiv.org/abs/1910.13461) for more768 information on the default strategy.769 770 - 1 indicates the head is **not masked**,771 - 0 indicates the head is **masked**.772 position_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):773 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,774 config.n_positions - 1]`.775 776 [What are position IDs?](../glossary#position-ids)777 past_key_values (`tuple(tuple(torch.FloatTensor))`, *optional*, returned when `use_cache=True` is passed or when `config.use_cache=True`):778 Tuple of `tuple(torch.FloatTensor)` of length `config.n_layers`, with each tuple having 2 tensors of shape779 `(batch_size, num_heads, sequence_length, embed_size_per_head)`) and 2 additional tensors of shape780 `(batch_size, num_heads, encoder_sequence_length, embed_size_per_head)`.781 782 Contains pre-computed hidden-states (key and values in the self-attention blocks and in the cross-attention783 blocks) that can be used (see `past_key_values` input) to speed up sequential decoding.784 785 If `past_key_values` are used, the user can optionally input only the last `input_ids` (those that don't786 have their past key value states given to this model) of shape `(batch_size, 1)` instead of all `input_ids`787 of shape `(batch_size, sequence_length)`.788 inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):789 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This790 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the791 model's internal embedding lookup matrix.792 use_cache (`bool`, *optional*):793 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see794 `past_key_values`).795 output_attentions (`bool`, *optional*):796 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned797 tensors for more detail.798 output_hidden_states (`bool`, *optional*):799 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for800 more detail.801 return_dict (`bool`, *optional*):802 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.803"""804 805 806@add_start_docstrings(807 "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",808 LLAMA_START_DOCSTRING,809)810class PolyLlamaModel(PolyLlamaPreTrainedModel):811 """812 Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`PolyLlamaDecoderLayer`]813 814 Args:815 config: LlamaConfig816 """817 818 def __init__(self, config: PolyLlamaConfig):819 super().__init__(config)820 self.padding_idx = config.pad_token_id821 self.vocab_size = config.vocab_size822 823 self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)824 self.layers = nn.ModuleList([PolyLlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])825 self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)826 827 self.gradient_checkpointing = False828 # Initialize weights and apply final processing829 self.post_init()830 831 def get_input_embeddings(self):832 return self.embed_tokens833 834 def set_input_embeddings(self, value):835 self.embed_tokens = value836 837 # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask838 def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):839 # create causal mask840 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]841 combined_attention_mask = None842 if input_shape[-1] > 1:843 combined_attention_mask = _make_causal_mask(844 input_shape,845 inputs_embeds.dtype,846 device=inputs_embeds.device,847 past_key_values_length=past_key_values_length,848 )849 850 if attention_mask is not None:851 # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]852 expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(853 inputs_embeds.device854 )855 combined_attention_mask = (856 expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask857 )858 859 return combined_attention_mask860 861 @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)862 def forward(863 self,864 input_ids: torch.LongTensor = None,865 attention_mask: Optional[torch.Tensor] = None,866 position_ids: Optional[torch.LongTensor] = None,867 past_key_values: Optional[List[torch.FloatTensor]] = None,868 inputs_embeds: Optional[torch.FloatTensor] = None,869 use_cache: Optional[bool] = None,870 output_attentions: Optional[bool] = None,871 output_hidden_states: Optional[bool] = None,872 return_dict: Optional[bool] = None,873 ) -> Union[Tuple, BaseModelOutputWithPast]:874 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions875 output_hidden_states = (876 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states877 )878 use_cache = use_cache if use_cache is not None else self.config.use_cache879 880 return_dict = return_dict if return_dict is not None else self.config.use_return_dict881 882 # retrieve input_ids and inputs_embeds883 if input_ids is not None and inputs_embeds is not None:884 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")885 elif input_ids is not None:886 batch_size, seq_length = input_ids.shape887 elif inputs_embeds is not None:888 batch_size, seq_length, _ = inputs_embeds.shape889 else:890 raise ValueError("You have to specify either input_ids or inputs_embeds")891 892 seq_length_with_past = seq_length893 past_key_values_length = 0894 895 if past_key_values is not None:896 past_key_values_length = past_key_values[0][0].shape[2]897 seq_length_with_past = seq_length_with_past + past_key_values_length898 899 if position_ids is None:900 device = input_ids.device if input_ids is not None else inputs_embeds.device901 position_ids = torch.arange(902 past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device903 )904 position_ids = position_ids.unsqueeze(0).view(-1, seq_length)905 else:906 position_ids = position_ids.view(-1, seq_length).long()907 908 if inputs_embeds is None:909 inputs_embeds = self.embed_tokens(input_ids)910 # embed positions911 if attention_mask is None:912 attention_mask = torch.ones(913 (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device914 )915 padding_mask = None916 else:917 if 0 in attention_mask:918 padding_mask = attention_mask919 else:920 padding_mask = None921 922 attention_mask = self._prepare_decoder_attention_mask(923 attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length924 )925 926 hidden_states = inputs_embeds927 928 if self.gradient_checkpointing and self.training:929 if use_cache:930 logger.warning_once(931 "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."932 )933 use_cache = False934 935 # decoder layers936 all_hidden_states = () if output_hidden_states else None937 all_self_attns = () if output_attentions else None938 next_decoder_cache = () if use_cache else None939 940 for idx, decoder_layer in enumerate(self.layers):941 if output_hidden_states:942 all_hidden_states += (hidden_states,)943 944 past_key_value = past_key_values[idx] if past_key_values is not None else None945 946 if self.gradient_checkpointing and self.training:947 948 def create_custom_forward(module):949 def custom_forward(*inputs):950 # None for past_key_value951 return module(*inputs, past_key_value, output_attentions, padding_mask=padding_mask)952 953 return custom_forward954 955 layer_outputs = torch.utils.checkpoint.checkpoint(956 create_custom_forward(decoder_layer), hidden_states, attention_mask, position_ids957 )958 else:959 layer_outputs = decoder_layer(960 hidden_states,961 attention_mask=attention_mask,962 position_ids=position_ids,963 past_key_value=past_key_value,964 output_attentions=output_attentions,965 use_cache=use_cache,966 padding_mask=padding_mask,967 )968 969 hidden_states = layer_outputs[0]970 971 if use_cache:972 next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)973 974 if output_attentions:975 all_self_attns += (layer_outputs[1],)976 977 hidden_states = self.norm(hidden_states)978 979 # add hidden states from the last decoder layer980 if output_hidden_states:981 all_hidden_states += (hidden_states,)982 983 next_cache = next_decoder_cache if use_cache else None984 if not return_dict:985 return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)986 return BaseModelOutputWithPast(987 last_hidden_state=hidden_states,988 past_key_values=next_cache,989 hidden_states=all_hidden_states,990 attentions=all_self_attns,991 )992 993 994class PolyLlamaForCausalLM(PolyLlamaPreTrainedModel):995 _tied_weights_keys = ["lm_head.weight"]996 997 def __init__(self, config):998 super().__init__(config)999 self.model = PolyLlamaModel(config)1000 self.vocab_size = config.vocab_size1001 self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)1002 1003 # Initialize weights and apply final processing1004 self.post_init()1005 1006 def get_input_embeddings(self):1007 return self.model.embed_tokens1008 1009 def set_input_embeddings(self, value):1010 self.model.embed_tokens = value1011 1012 def get_output_embeddings(self):1013 return self.lm_head1014 1015 def set_output_embeddings(self, new_embeddings):1016 self.lm_head = new_embeddings1017 1018 def set_decoder(self, decoder):1019 self.model = decoder1020 1021 def get_decoder(self):1022 return self.model1023 1024 @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)1025 @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)1026 def forward(1027 self,1028 input_ids: torch.LongTensor = None,1029 attention_mask: Optional[torch.Tensor] = None,1030 position_ids: Optional[torch.LongTensor] = None,1031 past_key_values: Optional[List[torch.FloatTensor]] = None,1032 inputs_embeds: Optional[torch.FloatTensor] = None,1033 labels: Optional[torch.LongTensor] = None,1034 use_cache: Optional[bool] = None,1035 output_attentions: Optional[bool] = None,1036 output_hidden_states: Optional[bool] = None,1037 return_dict: Optional[bool] = None,1038 ) -> Union[Tuple, CausalLMOutputWithPast]:1039 r"""1040 Args:1041 labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):1042 Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,1043 config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored1044 (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.1045 1046 Returns:1047 1048 Example:1049 1050 ```python1051 >>> from transformers import AutoTokenizer, LlamaForCausalLM1052 1053 >>> model = LlamaForCausalLM.from_pretrained(PATH_TO_CONVERTED_WEIGHTS)1054 >>> tokenizer = AutoTokenizer.from_pretrained(PATH_TO_CONVERTED_TOKENIZER)1055 1056 >>> prompt = "Hey, are you conscious? Can you talk to me?"1057 >>> inputs = tokenizer(prompt, return_tensors="pt")1058 1059 >>> # Generate1060 >>> generate_ids = model.generate(inputs.input_ids, max_length=30)1061 >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]1062 "Hey, are you conscious? Can you talk to me?\nI'm not conscious, but I can talk to you."1063 ```"""1064 1065 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions1066 output_hidden_states = (1067 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states1068 )1069 return_dict = return_dict if return_dict is not None else self.config.use_return_dict1070 1071 # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)1072 outputs = self.model(1073 input_ids=input_ids,1074 attention_mask=attention_mask,1075 position_ids=position_ids,1076 past_key_values=past_key_values,1077 inputs_embeds=inputs_embeds,1078 use_cache=use_cache,1079 output_attentions=output_attentions,1080 output_hidden_states=output_hidden_states,1081 return_dict=return_dict,1082 )1083 1084 hidden_states = outputs[0]1085 if self.config.pretraining_tp > 1:1086 lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)1087 logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]1088 logits = torch.cat(logits, dim=-1)1089 else:1090 logits = self.lm_head(hidden_states)1091 logits = logits.float()1092 1093 loss = None1094 if labels is not None:1095 # Shift so that tokens < n predict n1096 shift_logits = logits[..., :-1, :].contiguous()1097 shift_labels = labels[..., 1:].contiguous()1098 # Flatten the tokens1099 loss_fct = CrossEntropyLoss()1100 shift_logits = shift_logits.view(-1, self.config.vocab_size)1101 shift_labels = shift_labels.view(-1)1102 # Enable model parallelism1103 shift_labels = shift_labels.to(shift_logits.device)1104 loss = loss_fct(shift_logits, shift_labels)1105 1106 if not return_dict:1107 output = (logits,) + outputs[1:]1108 return (loss,) + output if loss is not None else output1109 1110 return CausalLMOutputWithPast(1111 loss=loss,1112 logits=logits,1113 past_key_values=outputs.past_key_values,1114 hidden_states=outputs.hidden_states,1115 attentions=outputs.attentions,1116 )1117 1118 def prepare_inputs_for_generation(1119 self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs1120 ):1121 if past_key_values:1122 input_ids = input_ids[:, -1:]1123 1124 position_ids = kwargs.get("position_ids", None)1125 if attention_mask is not None and position_ids is None:1126 # create position_ids on the fly for batch generation1127 position_ids = attention_mask.long().cumsum(-1) - 11128 position_ids.masked_fill_(attention_mask == 0, 1)1129 if past_key_values:1130 position_ids = position_ids[:, -1].unsqueeze(-1)1131 1132 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step1133 if inputs_embeds is not None and past_key_values is None:1134 model_inputs = {"inputs_embeds": inputs_embeds}1135 else:1136 model_inputs = {"input_ids": input_ids}1137 1138 model_inputs.update(1139 {1140 "position_ids": position_ids,1141 "past_key_values": past_key_values,1142 "use_cache": kwargs.get("use_cache"),1143 "attention_mask": attention_mask,1144 }1145 )1146 return model_inputs1147 1148 @staticmethod1149 def _reorder_cache(past_key_values, beam_idx):1150 reordered_past = ()1151 for layer_past in past_key_values:1152 reordered_past += (1153 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),1154 )1155 return reordered_past1156 1157 1158@add_start_docstrings(1159 """1160 The LLaMa Model transformer with a sequence classification head on top (linear layer).1161 1162 [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models1163 (e.g. GPT-2) do.1164 1165 Since it does classification on the last token, it requires to know the position of the last token. If a1166 `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If1167 no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the1168 padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in1169 each row of the batch).1170 """,1171 LLAMA_START_DOCSTRING,1172)1173class PolyLlamaForSequenceClassification(PolyLlamaPreTrainedModel):1174 def __init__(self, config):1175 super().__init__(config)1176 self.num_labels = config.num_labels1177 self.model = PolyLlamaModel(config)1178 self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)1179 1180 # Initialize weights and apply final processing1181 self.post_init()1182 1183 def get_input_embeddings(self):1184 return self.model.embed_tokens1185 1186 def set_input_embeddings(self, value):1187 self.model.embed_tokens = value1188 1189 @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)1190 def forward(1191 self,1192 input_ids: torch.LongTensor = None,1193 attention_mask: Optional[torch.Tensor] = None,1194 position_ids: Optional[torch.LongTensor] = None,1195 past_key_values: Optional[List[torch.FloatTensor]] = None,1196 inputs_embeds: Optional[torch.FloatTensor] = None,1197 labels: Optional[torch.LongTensor] = None,1198 use_cache: Optional[bool] = None,1199 output_attentions: Optional[bool] = None,1200 output_hidden_states: Optional[bool] = None,