WisdomShell/CodeShell-7B-Chat-int4
291.1k
1# coding=utf-82# Copyright 2023 WisdomShell Inc. All Rights Reserved.3 4# Licensed under the Apache License, Version 2.0 (the "License");5# you may not use this file except in compliance with the License.6# You may obtain a copy of the License at7#8# http://www.apache.org/licenses/LICENSE-2.09#10# Unless required by applicable law or agreed to in writing, software11# distributed under the License is distributed on an "AS IS" BASIS,12# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.13# See the License for the specific language governing permissions and14# limitations under the License.15 16# This code is based on Bigcode's GPTBigCode model. It has been modified from17# its original forms to accommodate minor architectural differences compared to 18# GPTBigCode model that trained the model.19 20# Copyright 2023 The Bigcode team and HuggingFace Inc. team.21# Licensed under the Apache License, Version 2.0 (the "License");22# you may not use this file except in compliance with the License.23# You may obtain a copy of the License at24#25# http://www.apache.org/licenses/LICENSE-2.026#27# Unless required by applicable law or agreed to in writing, software28# distributed under the License is distributed on an "AS IS" BASIS,29# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.30# See the License for the specific language governing permissions and31# limitations under the License.32"""PyTorch CodeShell model."""33import os34import math35from typing import List, Optional, Tuple, Union, Callable36from threading import Thread37from queue import Queue38 39 40import torch41import torch.utils.checkpoint42from torch import nn43from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss44 45from transformers import LogitsProcessorList, StoppingCriteriaList, StoppingCriteria, PreTrainedModel, PretrainedConfig46from transformers.generation.utils import GenerationConfig47 48from transformers.activations import ACT2FN49from transformers.modeling_outputs import (50 BaseModelOutputWithPastAndCrossAttentions,51 CausalLMOutputWithCrossAttentions,52)53from transformers.modeling_utils import PreTrainedModel54from transformers.utils import (55 add_start_docstrings,56 add_start_docstrings_to_model_forward,57)58from .configuration_codeshell import CodeShellConfig59 60# Fused kernels61# Use separate functions for each case because conditionals prevent kernel fusion.62# TODO: Could have better fused kernels depending on scaling, dropout and head mask.63# Is it doable without writing 32 functions?64@torch.jit.script65def upcast_masked_softmax(66 x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor, scale: float, softmax_dtype: torch.dtype67):68 input_dtype = x.dtype69 x = x.to(softmax_dtype) * scale70 x = torch.where(mask, x, mask_value)71 x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)72 return x73 74 75@torch.jit.script76def upcast_softmax(x: torch.Tensor, scale: float, softmax_dtype: torch.dtype):77 input_dtype = x.dtype78 x = x.to(softmax_dtype) * scale79 x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)80 return x81 82 83@torch.jit.script84def masked_softmax(x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor):85 x = torch.where(mask, x, mask_value)86 x = torch.nn.functional.softmax(x, dim=-1)87 return x88 89 90class CodeShellRotaryEmbedding(torch.nn.Module):91 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):92 super().__init__()93 94 self.dim = dim95 self.max_position_embeddings = max_position_embeddings96 self.base = base97 inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))98 self.register_buffer("inv_freq", inv_freq)99 100 # Build here to make `torch.jit.trace` work.101 self._set_cos_sin_cache(102 seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()103 )104 105 def _set_cos_sin_cache(self, seq_len, device, dtype):106 self.max_seq_len_cached = seq_len107 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)108 109 freqs = torch.einsum("i,j->ij", t, self.inv_freq)110 # Different from paper, but it uses a different permutation in order to obtain the same calculation111 emb = torch.cat((freqs, freqs), dim=-1)112 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)113 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)114 115 def forward(self, x, seq_len=None):116 # x: [bs, num_attention_heads, seq_len, head_size]117 if seq_len > self.max_seq_len_cached:118 self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)119 120 return (121 self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),122 self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),123 )124 125 126class CodeShellLinearScalingRotaryEmbedding(CodeShellRotaryEmbedding):127 """CodeShellRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""128 129 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):130 self.scaling_factor = scaling_factor131 super().__init__(dim, max_position_embeddings, base, device)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 t = t / self.scaling_factor137 138 freqs = torch.einsum("i,j->ij", t, self.inv_freq)139 # Different from paper, but it uses a different permutation in order to obtain the same calculation140 emb = torch.cat((freqs, freqs), dim=-1)141 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)142 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)143 144 145class CodeShellDynamicNTKScalingRotaryEmbedding(CodeShellRotaryEmbedding):146 """ShellRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""147 148 def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):149 self.scaling_factor = scaling_factor150 super().__init__(dim, max_position_embeddings, base, device)151 152 def _set_cos_sin_cache(self, seq_len, device, dtype):153 self.max_seq_len_cached = seq_len154 155 if seq_len > self.max_position_embeddings:156 base = self.base * (157 (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)158 ) ** (self.dim / (self.dim - 2))159 inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))160 self.register_buffer("inv_freq", inv_freq)161 162 t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)163 164 freqs = torch.einsum("i,j->ij", t, self.inv_freq)165 # Different from paper, but it uses a different permutation in order to obtain the same calculation166 emb = torch.cat((freqs, freqs), dim=-1)167 self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)168 self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)169 170def rotate_half(x):171 """Rotates half the hidden dims of the input."""172 x1 = x[..., : x.shape[-1] // 2]173 x2 = x[..., x.shape[-1] // 2 :]174 return torch.cat((-x2, x1), dim=-1)175 176 177def apply_rotary_pos_emb(q, k, cos, sin, position_ids):178 # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.179 cos = cos.squeeze(1).squeeze(0) # [seq_len, dim]180 sin = sin.squeeze(1).squeeze(0) # [seq_len, dim]181 cos = cos[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]182 sin = sin[position_ids].unsqueeze(1) # [bs, 1, seq_len, dim]183 q_embed = (q * cos) + (rotate_half(q) * sin)184 k_embed = (k * cos) + (rotate_half(k) * sin)185 return q_embed, k_embed186 187def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:188 """189 This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,190 num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)191 """192 batch, num_key_value_heads, slen, head_dim = hidden_states.shape193 if n_rep == 1:194 return hidden_states195 hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)196 return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)197 198class CodeShellAttention(nn.Module):199 def __init__(self, config, layer_idx=None):200 super().__init__()201 self.mask_value = None202 203 self.position_embedding_type = config.position_embedding_type204 self.rope_scaling = config.rope_scaling205 self.max_position_embeddings = config.max_position_embeddings206 207 self.group_query_attention = config.group_query_attention208 self.num_query_groups = config.num_query_groups209 self.num_key_value_groups = config.num_attention_heads // config.num_query_groups210 211 self.embed_dim = config.hidden_size212 self.num_heads = config.num_attention_heads213 self.head_dim = self.embed_dim // self.num_heads214 self.kv_heads = config.num_query_groups if self.group_query_attention else self.num_heads215 self.kv_dim = self.kv_heads * self.head_dim216 self.split_size = self.embed_dim217 if self.head_dim * self.num_heads != self.embed_dim:218 raise ValueError(219 f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"220 f" {self.num_heads})."221 )222 223 self.layer_idx = layer_idx224 225 self.c_attn = nn.Linear(self.embed_dim, self.embed_dim + 2 * self.kv_dim)226 self.c_proj = nn.Linear(self.embed_dim, self.embed_dim)227 228 self.attn_dropout = nn.Dropout(config.attn_pdrop)229 self.resid_dropout = nn.Dropout(config.resid_pdrop)230 231 if self.position_embedding_type == "rope":232 self._init_rope()233 234 def _init_rope(self):235 if self.rope_scaling is None:236 self.rotary_emb = CodeShellRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)237 else:238 scaling_type = self.rope_scaling["type"]239 scaling_factor = self.rope_scaling["factor"]240 if scaling_type == "linear":241 self.rotary_emb = CodeShellLinearScalingRotaryEmbedding(242 self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor243 )244 elif scaling_type == "dynamic":245 self.rotary_emb = CodeShellDynamicNTKScalingRotaryEmbedding(246 self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor247 )248 else:249 raise ValueError(f"Unknown RoPE scaling type {scaling_type}")250 251 252 def _get_mask_value(self, device, dtype):253 # torch.where expects a tensor. We use a cache to avoid recreating it every time.254 if self.mask_value is None or self.mask_value.dtype != dtype or self.mask_value.device != device:255 self.mask_value = torch.full([], torch.finfo(dtype).min, dtype=dtype, device=device)256 return self.mask_value257 258 def forward(259 self,260 hidden_states: torch.Tensor,261 layer_past: Optional[torch.Tensor] = None,262 attention_mask: Optional[torch.Tensor] = None,263 position_ids: Optional[torch.LongTensor] = None,264 head_mask: Optional[torch.Tensor] = None,265 use_cache: Optional[bool] = False,266 output_attentions: Optional[bool] = False,267 ) -> Union[268 Tuple[torch.Tensor, Optional[torch.Tensor]],269 Tuple[torch.Tensor, Optional[torch.Tensor], Tuple[torch.Tensor, ...]],270 ]:271 bsz, q_len, _ = hidden_states.size()272 query_states, key_states, value_states = self.c_attn(hidden_states).split((self.embed_dim, self.kv_dim, self.kv_dim), dim=2)273 274 query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)275 key_states = key_states.view(bsz, q_len, self.num_query_groups, self.head_dim).transpose(1, 2)276 value_states = value_states.view(bsz, q_len, self.num_query_groups, self.head_dim).transpose(1, 2)277 278 kv_seq_len = key_states.shape[-2]279 if layer_past is not None:280 kv_seq_len += layer_past[0].shape[-2]281 cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)282 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)283 284 if layer_past is not None:285 # reuse k, v, self_attention286 key_states = torch.cat([layer_past[0], key_states], dim=2)287 value_states = torch.cat([layer_past[1], value_states], dim=2)288 289 layer_past = (key_states, value_states) if use_cache else None290 291 # repeat k/v heads if n_kv_heads < n_heads292 key_states = repeat_kv(key_states, self.num_heads // self.kv_heads)293 value_states = repeat_kv(value_states, self.num_heads // self.kv_heads)294 295 attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)296 297 if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):298 raise ValueError(299 f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"300 f" {attn_weights.size()}"301 )302 303 if attention_mask is not None:304 if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):305 raise ValueError(306 f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"307 )308 mask_value = self._get_mask_value(attn_weights.device, attn_weights.dtype)309 # The fused kernel is very slow when the key length is not a multiple of 8, so we skip fusion.310 attn_weights = torch.where(attention_mask, attn_weights, mask_value)311 312 # upcast attention to fp32313 attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)314 attn_weights = self.attn_dropout(attn_weights)315 attn_output = torch.matmul(attn_weights, value_states)316 317 if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):318 raise ValueError(319 f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"320 f" {attn_output.size()}"321 )322 323 attn_output = attn_output.transpose(1, 2).contiguous()324 attn_output = attn_output.reshape(bsz, q_len, self.embed_dim)325 326 attn_output = self.c_proj(attn_output)327 attn_output = self.resid_dropout(attn_output)328 329 outputs = (attn_output, layer_past)330 if output_attentions:331 outputs += (attn_weights,)332 333 return outputs # a, present, (attentions)334 335 336class CodeShellMLP(nn.Module):337 def __init__(self, intermediate_size, config):338 super().__init__()339 embed_dim = config.hidden_size340 self.c_fc = nn.Linear(embed_dim, intermediate_size)341 self.c_proj = nn.Linear(intermediate_size, embed_dim)342 self.act = ACT2FN[config.activation_function]343 self.dropout = nn.Dropout(config.resid_pdrop)344 345 # Copied from transformers.models.gpt2.modeling_gpt2.GPT2MLP.forward346 def forward(self, hidden_states: Optional[Tuple[torch.Tensor]]) -> torch.Tensor:347 hidden_states = self.c_fc(hidden_states)348 hidden_states = self.act(hidden_states)349 hidden_states = self.c_proj(hidden_states)350 hidden_states = self.dropout(hidden_states)351 return hidden_states352 353 354class CodeShellBlock(nn.Module):355 def __init__(self, config, layer_idx=None):356 super().__init__()357 hidden_size = config.hidden_size358 self.inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size359 360 self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)361 self.attn = CodeShellAttention(config, layer_idx=layer_idx)362 self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)363 364 self.mlp = CodeShellMLP(self.inner_dim, config)365 366 def forward(367 self,368 hidden_states: Optional[Tuple[torch.Tensor]],369 layer_past: Optional[torch.Tensor] = None,370 attention_mask: Optional[torch.Tensor] = None,371 position_ids: Optional[torch.LongTensor] = None,372 head_mask: Optional[torch.Tensor] = None,373 encoder_hidden_states: Optional[torch.Tensor] = None,374 encoder_attention_mask: Optional[torch.Tensor] = None,375 use_cache: Optional[bool] = False,376 output_attentions: Optional[bool] = False,377 ) -> Union[378 Tuple[torch.Tensor], Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor, torch.Tensor, torch.Tensor]379 ]:380 residual = hidden_states381 hidden_states = self.ln_1(hidden_states)382 attn_outputs = self.attn(383 hidden_states,384 layer_past=layer_past,385 attention_mask=attention_mask,386 position_ids=position_ids,387 head_mask=head_mask,388 use_cache=use_cache,389 output_attentions=output_attentions,390 )391 attn_output = attn_outputs[0] # output_attn: a, present, (attentions)392 393 outputs = attn_outputs[1:]394 # residual connection395 hidden_states = attn_output + residual396 397 residual = hidden_states398 hidden_states = self.ln_2(hidden_states)399 feed_forward_hidden_states = self.mlp(hidden_states)400 # residual connection401 hidden_states = residual + feed_forward_hidden_states402 403 if use_cache:404 outputs = (hidden_states,) + outputs405 else:406 outputs = (hidden_states,) + outputs[1:]407 408 return outputs # hidden_states, present, (attentions, cross_attentions)409 410 411class CodeShellPreTrainedModel(PreTrainedModel):412 """413 An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained414 models.415 """416 417 config_class = CodeShellConfig418 base_model_prefix = "transformer"419 supports_gradient_checkpointing = True420 _no_split_modules = ["ShellBlock"]421 _skip_keys_device_placement = "past_key_values"422 423 def __init__(self, *inputs, **kwargs):424 super().__init__(*inputs, **kwargs)425 426 def _init_weights(self, module):427 """Initialize the weights."""428 if isinstance(module, (CodeShellMLP, CodeShellAttention)):429 # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:430 # > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale431 # > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.432 # > -- GPT-2 :: https://openai.com/blog/better-language-models/433 #434 # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py435 module.c_proj.weight.data.normal_(436 mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer))437 )438 module.c_proj._is_hf_initialized = True439 elif isinstance(module, nn.Linear):440 # Slightly different from the TF version which uses truncated_normal for initialization441 # cf https://github.com/pytorch/pytorch/pull/5617442 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)443 if module.bias is not None:444 module.bias.data.zero_()445 elif isinstance(module, nn.Embedding):446 module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)447 if module.padding_idx is not None:448 module.weight.data[module.padding_idx].zero_()449 elif isinstance(module, nn.LayerNorm):450 module.bias.data.zero_()451 module.weight.data.fill_(1.0)452 453 # Copied from transformers.models.gpt2.modeling_gpt2.GPT2PreTrainedModel._set_gradient_checkpointing with GPT2->Shell454 def _set_gradient_checkpointing(self, module, value=False):455 if isinstance(module, CodeShellModel):456 module.gradient_checkpointing = value457 458 459GPT_BIGCODE_START_DOCSTRING = r"""460 This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the461 library implements for all its model (such as downloading or saving, resizing the input embeddings, pruning heads462 etc.)463 This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.464 Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage465 and behavior.466 Parameters:467 config ([`CodeShellConfig`]): Model configuration class with all the parameters of the model.468 Initializing with a config file does not load the weights associated with the model, only the469 configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.470"""471 472GPT_BIGCODE_INPUTS_DOCSTRING = r"""473 Args:474 input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`):475 `input_ids_length` = `sequence_length` if `past_key_values` is `None` else476 `past_key_values[0][0].shape[-2]` (`sequence_length` of input past key value states). Indices of input477 sequence tokens in the vocabulary.478 If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as479 `input_ids`.480 Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and481 [`PreTrainedTokenizer.__call__`] for details.482 [What are input IDs?](../glossary#input-ids)483 past_key_values (`Tuple[torch.Tensor]` of length `config.n_layers`):484 Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see485 `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have486 their past given to this model should not be passed as `input_ids` as they have already been computed.487 attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):488 Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:489 - 1 for tokens that are **not masked**,490 - 0 for tokens that are **masked**.491 If `past_key_values` is used, `attention_mask` needs to contain the masking strategy that was used for492 `past_key_values`. In other words, the `attention_mask` always has to have the length:493 `len(past_key_values) + len(input_ids)`494 [What are attention masks?](../glossary#attention-mask)495 token_type_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`, *optional*):496 Segment token indices to indicate first and second portions of the inputs. Indices are selected in `[0,497 1]`:498 - 0 corresponds to a *sentence A* token,499 - 1 corresponds to a *sentence B* token.500 [What are token type IDs?](../glossary#token-type-ids)501 position_ids (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):502 Indices of positions of each input sequence tokens in the position embeddings. Selected in the range `[0,503 config.max_position_embeddings - 1]`.504 [What are position IDs?](../glossary#position-ids)505 head_mask (`torch.Tensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):506 Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:507 - 1 indicates the head is **not masked**,508 - 0 indicates the head is **masked**.509 inputs_embeds (`torch.Tensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):510 Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This511 is useful if you want more control over how to convert `input_ids` indices into associated vectors than the512 model's internal embedding lookup matrix.513 If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see514 `past_key_values`).515 use_cache (`bool`, *optional*):516 If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see517 `past_key_values`).518 output_attentions (`bool`, *optional*):519 Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned520 tensors for more detail.521 output_hidden_states (`bool`, *optional*):522 Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for523 more detail.524 return_dict (`bool`, *optional*):525 Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple.526"""527 528 529@add_start_docstrings(530 "The bare GPT_BIGCODE Model transformer outputting raw hidden-states without any specific head on top.",531 GPT_BIGCODE_START_DOCSTRING,532)533class CodeShellModel(CodeShellPreTrainedModel):534 def __init__(self, config):535 super().__init__(config)536 self.group_query_attention = config.group_query_attention537 self.num_query_groups = config.num_query_groups538 self.position_embedding_type = config.position_embedding_type539 self.embed_dim = config.hidden_size540 541 self.wte = nn.Embedding(config.vocab_size, self.embed_dim)542 if self.position_embedding_type == "learned_absolute":543 self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)544 else:545 pass546 547 self.drop = nn.Dropout(config.embd_pdrop)548 self.h = nn.ModuleList([CodeShellBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)])549 self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)550 551 max_positions = config.max_position_embeddings552 self.register_buffer(553 "bias", torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)), persistent=False554 )555 556 self.gradient_checkpointing = False557 558 # Initialize weights and apply final processing559 self.post_init()560 561 def get_input_embeddings(self):562 return self.wte563 564 def set_input_embeddings(self, new_embeddings):565 self.wte = new_embeddings566 567 @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING)568 def forward(569 self,570 input_ids: Optional[torch.Tensor] = None,571 past_key_values: Optional[List[torch.Tensor]] = None,572 attention_mask: Optional[torch.Tensor] = None,573 token_type_ids: Optional[torch.Tensor] = None,574 position_ids: Optional[torch.Tensor] = None,575 head_mask: Optional[torch.Tensor] = None,576 inputs_embeds: Optional[torch.Tensor] = None,577 encoder_hidden_states: Optional[torch.Tensor] = None,578 encoder_attention_mask: Optional[torch.Tensor] = None,579 use_cache: Optional[bool] = None,580 output_attentions: Optional[bool] = None,581 output_hidden_states: Optional[bool] = None,582 return_dict: Optional[bool] = None,583 ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:584 output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions585 output_hidden_states = (586 output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states587 )588 use_cache = use_cache if use_cache is not None else self.config.use_cache589 return_dict = return_dict if return_dict is not None else self.config.use_return_dict590 591 if input_ids is not None and inputs_embeds is not None:592 raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")593 elif input_ids is not None:594 input_shape = input_ids.size()595 input_ids = input_ids.reshape(-1, input_shape[-1])596 batch_size = input_ids.shape[0]597 elif inputs_embeds is not None:598 input_shape = inputs_embeds.size()[:-1]599 batch_size = inputs_embeds.shape[0]600 else:601 raise ValueError("You have to specify either input_ids or inputs_embeds")602 603 if batch_size <= 0:604 raise ValueError("batch_size has to be defined and > 0")605 606 device = input_ids.device if input_ids is not None else inputs_embeds.device607 608 if token_type_ids is not None:609 token_type_ids = token_type_ids.reshape(-1, input_shape[-1])610 if position_ids is not None:611 position_ids = position_ids.reshape(-1, input_shape[-1])612 613 if past_key_values is None:614 past_length = 0615 past_key_values = tuple([None] * len(self.h))616 else:617 past_length = past_key_values[0][0].size(-2)618 619 if attention_mask is not None and len(attention_mask.shape) == 2 and position_ids is None:620 # create position_ids on the fly for batch generation621 position_ids = attention_mask.long().cumsum(-1) - 1622 position_ids.masked_fill_(attention_mask == 0, 1)623 if past_length > 0:624 position_ids = position_ids[:, past_length : input_shape[-1] + past_length :]625 elif position_ids is None:626 position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)627 position_ids = position_ids.unsqueeze(0).reshape(-1, input_shape[-1])628 629 # Self-attention mask.630 query_length = input_shape[-1]631 key_length = past_length + query_length632 self_attention_mask = self.bias[None, key_length - query_length : key_length, :key_length]633 634 if attention_mask is not None:635 self_attention_mask = self_attention_mask * attention_mask.reshape(batch_size, 1, -1).to(636 dtype=torch.bool, device=self_attention_mask.device637 )638 639 # MQA models: (batch_size, query_length, n_heads, key_length)640 # MHA models: (batch_size, n_heads, query_length, key_length)641 attention_mask = self_attention_mask.unsqueeze(1)642 643 encoder_attention_mask = None644 645 # Prepare head mask if needed646 # 1.0 in head_mask indicate we keep the head647 # attention_probs has shape bsz x n_heads x N x N648 # head_mask has shape n_layer x batch x n_heads x N x N649 head_mask = self.get_head_mask(head_mask, self.config.n_layer)650 651 if inputs_embeds is None:652 inputs_embeds = self.wte(input_ids)653 654 hidden_states = inputs_embeds655 if self.position_embedding_type == "learned_absolute":656 position_embeds = self.wpe(position_ids)657 hidden_states = hidden_states + position_embeds658 659 if token_type_ids is not None:660 token_type_embeds = self.wte(token_type_ids)661 hidden_states = hidden_states + token_type_embeds662 663 hidden_states = self.drop(hidden_states)664 665 output_shape = input_shape + (hidden_states.size(-1),)666 667 presents = [] if use_cache else None668 all_self_attentions = () if output_attentions else None669 all_hidden_states = () if output_hidden_states else None670 for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):671 if output_hidden_states:672 all_hidden_states = all_hidden_states + (hidden_states,)673 674 if self.gradient_checkpointing and self.training:675 676 def create_custom_forward(module):677 def custom_forward(*inputs):678 # None for past_key_value679 return module(*inputs, use_cache, output_attentions)680 681 return custom_forward682 683 outputs = torch.utils.checkpoint.checkpoint(684 create_custom_forward(block),685 hidden_states,686 None,687 attention_mask,688 position_ids,689 head_mask[i],690 encoder_hidden_states,691 encoder_attention_mask,692 )693 else:694 outputs = block(695 hidden_states,696 layer_past=layer_past,697 attention_mask=attention_mask,698 position_ids=position_ids,699 head_mask=head_mask[i],700 encoder_hidden_states=encoder_hidden_states,701 encoder_attention_mask=encoder_attention_mask,702 use_cache=use_cache,703 output_attentions=output_attentions,704 )705 706 hidden_states = outputs[0]707 if use_cache:708 presents.append(outputs[1])709 710 if output_attentions:711 all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)712 713 hidden_states = self.ln_f(hidden_states)714 hidden_states = hidden_states.reshape(output_shape)715 # Add last hidden state716 if output_hidden_states:717 all_hidden_states = all_hidden_states + (hidden_states,)718 719 720 if not return_dict:721 return tuple(722 v723 for v in [hidden_states, presents, all_hidden_states, all_self_attentions]724 if v is not None725 )726 727 return BaseModelOutputWithPastAndCrossAttentions(728 last_hidden_state=hidden_states,729 past_key_values=presents,730 hidden_states=all_hidden_states,731 attentions=all_self_attentions,732 )733 734class EndOfFunctionCriteria(StoppingCriteria):735 """Custom `StoppingCriteria` which checks if all generated functions in the batch are completed."""736 def __init__(self, input_lengths, eof_strings, tokenizer):737 self.input_lengths = input_lengths738 self.eof_strings = eof_strings739 self.tokenizer = tokenizer740 741 def __call__(self, input_ids, scores, **kwargs):742 """Returns true if all generated sequences contain any of the end-of-function strings."""743 decoded_generations = []744 for _input_ids, input_length in zip(input_ids, self.input_lengths):745 decoded_generations.append(self.tokenizer.decode(_input_ids[input_length:]))746 done = []747 for decoded_generation in decoded_generations:748 done.append(749 any(750 [751 stop_string in decoded_generation752 for stop_string in self.eof_strings753 ]754 )755 )756 return all(done)757 758class TextIterStreamer:759 def __init__(self, tokenizer, skip_prompt=False, skip_special_tokens=False):760 self.tokenizer = tokenizer761 self.skip_prompt = skip_prompt762 self.skip_special_tokens = skip_special_tokens763 self.tokens = []764 self.text_queue = Queue()765 self.next_tokens_are_prompt = True766 767 def put(self, value):768 if self.skip_prompt and self.next_tokens_are_prompt:769 self.next_tokens_are_prompt = False770 else:771 if len(value.shape) > 1:772 value = value[0]773 self.tokens.extend(value.tolist())774 self.text_queue.put(775 self.tokenizer.decode(self.tokens, skip_special_tokens=self.skip_special_tokens))776 777 def end(self):778 self.text_queue.put(None)779 780 def __iter__(self):781 return self782 783 def __next__(self):784 value = self.text_queue.get()785 if value is None:786 raise StopIteration()787 else:788 return value789 790 791@add_start_docstrings(792 """793 The GPT_BIGCODE Model transformer with a language modeling head on top (linear layer with weights tied to the input794 embeddings).795 """,796 GPT_BIGCODE_START_DOCSTRING,797)798class CodeShellForCausalLM(CodeShellPreTrainedModel):799 _tied_weights_keys = ["lm_head.weight"]800 801 def __init__(self, config):802 super().__init__(config)803 self.transformer = CodeShellModel(config)804 self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)805 806 # Initialize weights and apply final processing807 self.post_init()808 809 def quantize(self, bits: int):810 try:811 import bitsandbytes812 from .quantizer import quantize813 except ImportError:814 raise ImportError(f"Needs bitsandbytes to run quantize.")815 return quantize(self, bits)816 817 def get_output_embeddings(self):818 return self.lm_head819 820 def set_output_embeddings(self, new_embeddings):821 self.lm_head = new_embeddings822 823 def prepare_inputs_for_generation(self, input_ids, past_key_values=None, inputs_embeds=None, **kwargs):824 token_type_ids = kwargs.get("token_type_ids", None)825 # only last token for inputs_ids if past is defined in kwargs826 if past_key_values:827 input_ids = input_ids[:, -1].unsqueeze(-1)828 if token_type_ids is not None:829 token_type_ids = token_type_ids[:, -1].unsqueeze(-1)830 831 attention_mask = kwargs.get("attention_mask", None)832 position_ids = kwargs.get("position_ids", None)833 834 if attention_mask is not None and position_ids is None:835 # create position_ids on the fly for batch generation836 position_ids = attention_mask.long().cumsum(-1) - 1837 position_ids.masked_fill_(attention_mask == 0, 1)838 if past_key_values:839 position_ids = position_ids[:, -1].unsqueeze(-1)840 else:841 position_ids = None842 843 # if `inputs_embeds` are passed, we only want to use them in the 1st generation step844 if inputs_embeds is not None and past_key_values is None:845 model_inputs = {"inputs_embeds": inputs_embeds}846 else:847 model_inputs = {"input_ids": input_ids}848 849 model_inputs.update(850 {851 "past_key_values": past_key_values,852 "use_cache": kwargs.get("use_cache"),853 "position_ids": position_ids,854 "attention_mask": attention_mask,855 "token_type_ids": token_type_ids,856 }857 )858 return model_inputs859 860 @add_start_docstrings_to_model_forward(GPT_BIGCODE_INPUTS_DOCSTRING)861 def forward(862 self,863 input_ids: Optional[torch.Tensor] = None,864 past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,865 attention_mask: Optional[torch.Tensor] = None,866 token_type_ids: Optional[torch.Tensor] = None,867 position_ids: Optional[torch.Tensor] = None,868 head_mask: Optional[torch.Tensor] = None,869 inputs_embeds: Optional[torch.Tensor] = None,870 encoder_hidden_states: Optional[torch.Tensor] = None,871 encoder_attention_mask: Optional[torch.Tensor] = None,872 labels: Optional[torch.Tensor] = None,873 use_cache: Optional[bool] = None,874 output_attentions: Optional[bool] = None,875 output_hidden_states: Optional[bool] = None,876 return_dict: Optional[bool] = None,877 ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:878 r"""879 labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):880 Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set881 `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`882 are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`883 """884 return_dict = return_dict if return_dict is not None else self.config.use_return_dict885 886 transformer_outputs = self.transformer(887 input_ids,888 past_key_values=past_key_values,889 attention_mask=attention_mask,890 token_type_ids=token_type_ids,891 position_ids=position_ids,892 head_mask=head_mask,893 inputs_embeds=inputs_embeds,894 encoder_hidden_states=encoder_hidden_states,895 encoder_attention_mask=encoder_attention_mask,896 use_cache=use_cache,897 output_attentions=output_attentions,898 output_hidden_states=output_hidden_states,899 return_dict=return_dict,900 )901 hidden_states = transformer_outputs[0]902 lm_logits = self.lm_head(hidden_states)903 loss = None904 if labels is not None:905 # Shift so that tokens < n predict n906 shift_logits = lm_logits[..., :-1, :].contiguous()907 shift_labels = labels[..., 1:].contiguous().to(shift_logits.device)908 # Flatten the tokens909 loss_fct = CrossEntropyLoss()910 loss = loss_fct(shift_logits.reshape(-1, shift_logits.size(-1)), shift_labels.reshape(-1))911 912 if not return_dict:913 output = (lm_logits,) + transformer_outputs[1:]914 return ((loss,) + output) if loss is not None else output915 916 return CausalLMOutputWithCrossAttentions(917 loss=loss,918 logits=lm_logits,919 past_key_values=transformer_outputs.past_key_values,920 hidden_states=transformer_outputs.hidden_states,921 attentions=transformer_outputs.attentions,922 )923 924 @staticmethod925 def _reorder_cache(past_key_values, beam_idx):926 reordered_past = ()927 for layer_past in past_key_values:928 reordered_past += (929 tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),930 )931 return reordered_past932 933 934 def build_chat_input(self, query, history, tokenizer, max_new_tokens=None):935 user_name = "## human:"936 ai_name = "## assistant: "937 stop = '|<end>|'938 939 prompt = ''940 for q, r in history:941 prompt += f"{user_name}{q}{stop}"942 prompt += f"{ai_name}{r}{stop}"943 prompt += f"{user_name}{query}{stop}"944 prompt += ai_name.rstrip()945 946 max_new_tokens = max_new_tokens or self.generation_config.max_new_tokens947 max_new_tokens = max_new_tokens or 128948 max_input_tokens = self.config.n_positions - max_new_tokens949 950 input_tokens = tokenizer.encode(prompt)951 input_tokens = input_tokens[-max_input_tokens:] # truncate left952 return torch.LongTensor([input_tokens]).to(self.device)953 954 def chat(self, query, history, tokenizer, stream=False,955 generation_config: Optional[GenerationConfig]=None):956 generation_config = generation_config or self.generation_config957 input_ids = self.build_chat_input(query, history, tokenizer, generation_config.max_new_tokens)958 stopping_criteria = StoppingCriteriaList(959 [EndOfFunctionCriteria([len(input_ids[0])], ['|<end>|', '|end|', '<|endoftext|>', '## human'], tokenizer)]960 )961 962 if stream:963 streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)964 Thread(target=self.generate, kwargs=dict(965 inputs=input_ids, streamer=streamer,966 stopping_criteria = stopping_criteria,967 generation_config=generation_config,968 )).start()969 return streamer970 else:971 outputs = self.generate(input_ids, generation_config=generation_config, stopping_criteria = stopping_criteria)972 response = tokenizer.decode(outputs[0][len(input_ids[0]):], skip_special_tokens=True)973 return response974 975 def generate_stream(self, prompt, tokenizer, generation_config=None, **kwargs):976 generation_config = generation_config or self.generation_config977 max_input_tokens = self.config.n_positions - self.generation_config.max_new_tokens978 979 input_ids = tokenizer.encode(prompt)980 input_ids = input_ids[-max_input_tokens:] # truncate left981 982 stopping_criteria = StoppingCriteriaList(983 [EndOfFunctionCriteria([len(input_ids[0])], ['|<end>|', '|end|', '<|endoftext|>', '## human'], tokenizer)]984 )985 986 streamer = TextIterStreamer(tokenizer, skip_prompt=True, skip_special_tokens=True)987 Thread(target=self.generate, kwargs=dict(988 inputs=input_ids, stopping_criteria=stopping_criteria, **kwargs989 )).start()990 return streamer991 992 993class CodeShell4bitForCausalLM(CodeShellForCausalLM):994 def __init__(self, config):995 CodeShellPreTrainedModel.__init__(self, config) 996 self.transformer = CodeShellModel(config)997 self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)998 999 try:1000 import bitsandbytes1001 from .quantizer import quantize_offline1002 quantize_offline(self)1003 except ImportError:1004 raise ImportError(f"Needs bitsandbytes to run quantize.")1005 1006 self.post_init()1007 1008 @classmethod1009 def from_pretrained(1010 cls,1011 pretrained_model_name_or_path: Optional[Union[str, os.PathLike]],1012 *model_args,1013 config: Optional[Union[PretrainedConfig, str, os.PathLike]] = None,1014 cache_dir: Optional[Union[str, os.PathLike]] = None,1015 ignore_mismatched_sizes: bool = False,1016 force_download: bool = False,1017 local_files_only: bool = False,1018 token: Optional[Union[str, bool]] = None,1019 revision: str = "main",1020 use_safetensors: bool = None,1021 **kwargs,1022 ):1023 if not isinstance(config, PretrainedConfig):1024 config_path = config if config is not None else pretrained_model_name_or_path1025 config, _ = cls.config_class.from_pretrained(1026 config_path,1027 cache_dir=cache_dir,1028 return_unused_kwargs=True,1029 force_download=force_download,1030 resume_download=False,1031 proxies=None,1032 local_files_only=local_files_only,1033 token=token,1034 revision=revision,1035 subfolder="",1036 _from_auto=False,1037 _from_pipeline=None,1038 **kwargs,1039 )1040 1041 # Load config if we don't provide a configuration1042 from .quantizer import load_state_dict_for_qunantied_model1043 model = cls(config)1044 state_dict = torch.load(os.path.join(pretrained_model_name_or_path, 'pytorch_model.bin'), map_location="cpu") 1045 model = load_state_dict_for_qunantied_model(model, state_dict)1046 model.eval()1047 1048 # If it is a model with generation capabilities, attempt to load the generation config1049 if model.can_generate():1050 try:1051 model.generation_config = GenerationConfig.from_pretrained(1052 pretrained_model_name_or_path,1053 cache_dir=cache_dir,1054 force_download=force_download,1055 resume_download=False,1056 proxies=None,1057 local_files_only=local_files_only,1058 token=token,1059 revision=revision,1060 subfolder="",1061 _from_auto=False,1062 _from_pipeline=None,1063 **kwargs,1064 )1065 except (OSError, TypeError):1066 pass1067 1068 device_map = kwargs.pop("device_map", None)1069 if device_map is not None:1070 model = model.to(torch.device(device_map))1071 1072 return model