Aluode/PerceptionLabPortable

1# coding=utf-82# Copyright 2023 The Bigcode team and HuggingFace Inc. team.3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14"""PyTorch GPTBigCode model."""15 16import math17from typing import Callable, Optional, Union18 19import torch20from torch import nn21from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss22 23from ...activations import ACT2FN24from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache25from ...generation import GenerationMixin26from ...masking_utils import create_causal_mask27from ...modeling_flash_attention_utils import is_flash_attn_available28from ...modeling_outputs import (29    BaseModelOutputWithPastAndCrossAttentions,30    CausalLMOutputWithCrossAttentions,31    SequenceClassifierOutputWithPast,32    TokenClassifierOutput,33)34from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel35from ...utils import (36    auto_docstring,37    can_return_tuple,38    logging,39)40from .configuration_gpt_bigcode import GPTBigCodeConfig41 42 43if is_flash_attn_available():44    pass45 46 47logger = logging.get_logger(__name__)48 49 50# Fused kernels51# Use separate functions for each case because conditionals prevent kernel fusion.52# TODO: Could have better fused kernels depending on scaling, dropout and head mask.53#  Is it doable without writing 32 functions?54@torch.jit.script55def upcast_masked_softmax(56    x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor, scale: float, softmax_dtype: torch.dtype57):58    input_dtype = x.dtype59    x = x.to(softmax_dtype) * scale60    x = torch.where(mask, x, mask_value)61    x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)62    return x63 64 65@torch.jit.script66def upcast_softmax(x: torch.Tensor, scale: float, softmax_dtype: torch.dtype):67    input_dtype = x.dtype68    x = x.to(softmax_dtype) * scale69    x = torch.nn.functional.softmax(x, dim=-1).to(input_dtype)70    return x71 72 73@torch.jit.script74def masked_softmax(x: torch.Tensor, mask: torch.Tensor, mask_value: torch.Tensor):75    x = torch.where(mask, x, mask_value)76    x = torch.nn.functional.softmax(x, dim=-1)77    return x78 79 80def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:81    """82    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,83    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)84    """85    batch, num_key_value_heads, slen, head_dim = hidden_states.shape86    if n_rep == 1:87        return hidden_states88    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)89    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)90 91 92def eager_attention_forward(93    module: nn.Module,94    query: torch.Tensor,95    key: torch.Tensor,96    value: torch.Tensor,97    attention_mask: Optional[torch.Tensor],98    scaling: float,99    dropout: float = 0.0,100    head_mask: Optional[torch.Tensor] = None,101    **kwargs,102):103    key_states = repeat_kv(key, module.num_key_value_groups)104    value_states = repeat_kv(value, module.num_key_value_groups)105 106    attn_weights = torch.matmul(query, key_states.transpose(2, 3)) * scaling107    if attention_mask is not None:108        causal_mask = attention_mask[:, :, :, : key_states.shape[-2]]109        attn_weights = attn_weights + causal_mask110 111    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)112    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)113 114    if head_mask is not None:115        attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)116 117    attn_output = torch.matmul(attn_weights, value_states)118    attn_output = attn_output.transpose(1, 2).contiguous()119 120    return attn_output, attn_weights121 122 123class GPTBigCodeAttention(nn.Module):124    def __init__(self, config, is_cross_attention=False, layer_idx=None):125        super().__init__()126        self.config = config127 128        self.mask_value = None129        self.multi_query = config.multi_query130        self.embed_dim = config.hidden_size131        self.num_heads = config.num_attention_heads132        self.head_dim = self.embed_dim // self.num_heads133        self.kv_heads = 1 if self.multi_query else self.num_heads134        self.kv_dim = self.kv_heads * self.head_dim135        self.num_key_value_groups = self.num_heads // self.kv_heads136        self.split_size = self.embed_dim137        self.is_causal = True138 139        if self.head_dim * self.num_heads != self.embed_dim:140            raise ValueError(141                f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"142                f" {self.num_heads})."143            )144 145        self.scale_attn_weights = config.scale_attn_weights146        self.scaling = self.head_dim**-0.5 if config.scale_attn_weights else 1.0147        self.is_cross_attention = is_cross_attention148 149        self.layer_idx = layer_idx150        self.attention_softmax_in_fp32 = config.attention_softmax_in_fp32151        self.scale_attention_softmax_in_fp32 = (152            config.scale_attention_softmax_in_fp32 and config.attention_softmax_in_fp32153        )154        self.attn_pdrop = config.attn_pdrop155 156        if self.is_cross_attention:157            if self.multi_query:158                raise NotImplementedError("Multi-Query Attention not supported for cross_attention")159 160            self.c_attn = nn.Linear(self.embed_dim, 2 * self.embed_dim)161            self.q_attn = nn.Linear(self.embed_dim, self.embed_dim)162        else:163            self.c_attn = nn.Linear(self.embed_dim, self.embed_dim + 2 * self.kv_dim)164 165        self.c_proj = nn.Linear(self.embed_dim, self.embed_dim)166 167        self.attn_dropout = config.attn_pdrop168        self.resid_dropout = nn.Dropout(config.resid_pdrop)169 170    def forward(171        self,172        hidden_states: torch.Tensor,173        layer_past: Optional[Cache] = None,174        attention_mask: Optional[torch.Tensor] = None,175        head_mask: Optional[torch.Tensor] = None,176        encoder_hidden_states: Optional[torch.Tensor] = None,177        encoder_attention_mask: Optional[torch.Tensor] = None,178        use_cache: Optional[bool] = False,179        output_attentions: Optional[bool] = False,180        cache_position: Optional[torch.Tensor] = None,181        **kwargs,182    ) -> Union[183        tuple[torch.Tensor, Optional[torch.Tensor]],184        tuple[torch.Tensor, Optional[torch.Tensor], tuple[torch.Tensor, ...]],185    ]:186        input_shape = hidden_states.shape[:-1]187 188        if layer_past is not None:189            if isinstance(layer_past, EncoderDecoderCache):190                is_updated = layer_past.is_updated.get(self.layer_idx)191                if self.is_cross_attention:192                    # after the first generated id, we can subsequently re-use all key/value_states from cache193                    curr_past_key_value = layer_past.cross_attention_cache194                else:195                    curr_past_key_value = layer_past.self_attention_cache196            else:197                curr_past_key_value = layer_past198 199        if self.is_cross_attention:200            if not hasattr(self, "q_attn") or not self.is_cross_attention:201                raise ValueError(202                    "If class is used as cross attention, the weights `q_attn` have to be defined. "203                    "Please make sure to instantiate class with `GPTBigCodeAttention(..., is_cross_attention=True)`."204                )205            if layer_past is not None and is_updated:206                # reuse k,v, cross_attentions207                key = curr_past_key_value.layers[self.layer_idx].keys208                value = curr_past_key_value.layers[self.layer_idx].values209            else:210                query = self.q_attn(hidden_states).view(*input_shape, -1, self.head_dim).transpose(1, 2)211                key, value = self.c_attn(encoder_hidden_states).split((self.head_dim, self.head_dim), dim=-1)212        else:213            if self.multi_query:214                query, key, value = (215                    self.c_attn(hidden_states).unsqueeze(1).split((self.embed_dim, self.kv_dim, self.kv_dim), dim=3)216                )217                query = query.view(*input_shape, -1, self.head_dim).transpose(1, 2)218            else:219                query, key, value = (220                    self.c_attn(hidden_states)221                    .view(*hidden_states.shape[:2], self.num_heads, 3 * self.head_dim)222                    .transpose(1, 2)223                    .split(3 * [self.head_dim], dim=3)224                )225 226        if layer_past is not None:227            # save all key/value_states to cache to be re-used for fast auto-regressive generation228            cache_position = cache_position if not self.is_cross_attention else None229            key, value = curr_past_key_value.update(key, value, self.layer_idx, {"cache_position": cache_position})230            # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls231            if self.is_cross_attention:232                layer_past.is_updated[self.layer_idx] = True233 234        attention_interface: Callable = eager_attention_forward235        if self.config._attn_implementation != "eager":236            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]237 238        attn_output, attn_weights = attention_interface(239            self,240            query,241            key,242            value,243            attention_mask,244            dropout=0.0 if not self.training else self.attn_dropout,245            scaling=self.scaling,246            head_mask=head_mask,247            **kwargs,248        )249 250        attn_output = attn_output.reshape(*input_shape, -1).contiguous()251        attn_output = self.c_proj(attn_output)252        attn_output = self.resid_dropout(attn_output)253        return attn_output, attn_weights254 255 256class GPTBigCodeMLP(nn.Module):257    def __init__(self, intermediate_size, config):258        super().__init__()259        embed_dim = config.hidden_size260        self.c_fc = nn.Linear(embed_dim, intermediate_size)261        self.c_proj = nn.Linear(intermediate_size, embed_dim)262        self.act = ACT2FN[config.activation_function]263        self.dropout = nn.Dropout(config.resid_pdrop)264 265    # Copied from transformers.models.gpt2.modeling_gpt2.GPT2MLP.forward266    def forward(self, hidden_states: Optional[tuple[torch.FloatTensor]]) -> torch.FloatTensor:267        hidden_states = self.c_fc(hidden_states)268        hidden_states = self.act(hidden_states)269        hidden_states = self.c_proj(hidden_states)270        hidden_states = self.dropout(hidden_states)271        return hidden_states272 273 274class GPTBigCodeBlock(nn.Module):275    def __init__(self, config, layer_idx=None):276        super().__init__()277        hidden_size = config.hidden_size278        self.inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size279 280        self.ln_1 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)281 282        self.attn = GPTBigCodeAttention(config, layer_idx=layer_idx)283 284        self.ln_2 = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)285 286        if config.add_cross_attention:287            if config.multi_query:288                raise NotImplementedError("Cross-attention not implemented for MQA")289 290            self.crossattention = GPTBigCodeAttention(config, is_cross_attention=True, layer_idx=layer_idx)291 292            self.ln_cross_attn = nn.LayerNorm(hidden_size, eps=config.layer_norm_epsilon)293 294        self.mlp = GPTBigCodeMLP(self.inner_dim, config)295 296    def forward(297        self,298        hidden_states: Optional[tuple[torch.Tensor]],299        layer_past: Optional[Cache] = None,300        attention_mask: Optional[torch.Tensor] = None,301        head_mask: Optional[torch.Tensor] = None,302        encoder_hidden_states: Optional[torch.Tensor] = None,303        encoder_attention_mask: Optional[torch.Tensor] = None,304        use_cache: Optional[bool] = False,305        output_attentions: Optional[bool] = False,306        cache_position: Optional[torch.Tensor] = None,307        **kwargs,308    ) -> Union[309        tuple[torch.Tensor], tuple[torch.Tensor, torch.Tensor], tuple[torch.Tensor, torch.Tensor, torch.Tensor]310    ]:311        residual = hidden_states312        hidden_states = self.ln_1(hidden_states)313        attn_outputs = self.attn(314            hidden_states,315            layer_past=layer_past,316            attention_mask=attention_mask,317            head_mask=head_mask,318            use_cache=use_cache,319            output_attentions=output_attentions,320            cache_position=cache_position,321            **kwargs,322        )323        attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)324        outputs = attn_outputs[1:]325        # residual connection326        hidden_states = attn_output + residual327 328        if encoder_hidden_states is not None:329            # add one self-attention block for cross-attention330            if not hasattr(self, "crossattention"):331                raise ValueError(332                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "333                    "cross-attention layers by setting `config.add_cross_attention=True`"334                )335            residual = hidden_states336            hidden_states = self.ln_cross_attn(hidden_states)337            cross_attn_outputs = self.crossattention(338                hidden_states,339                attention_mask=attention_mask,340                head_mask=head_mask,341                encoder_hidden_states=encoder_hidden_states,342                encoder_attention_mask=encoder_attention_mask,343                output_attentions=output_attentions,344                cache_position=cache_position,345                **kwargs,346            )347            attn_output = cross_attn_outputs[0]348            # residual connection349            hidden_states = residual + attn_output350            outputs = outputs + cross_attn_outputs[1:]  # add cross attentions if we output attention weights351 352        residual = hidden_states353        hidden_states = self.ln_2(hidden_states)354        feed_forward_hidden_states = self.mlp(hidden_states)355        hidden_states = residual + feed_forward_hidden_states356        return (hidden_states,) + outputs357 358 359@auto_docstring360class GPTBigCodePreTrainedModel(PreTrainedModel):361    config: GPTBigCodeConfig362    base_model_prefix = "transformer"363    supports_gradient_checkpointing = True364    _no_split_modules = ["GPTBigCodeBlock"]365    _skip_keys_device_placement = "past_key_values"366    _supports_flash_attn = True367    _supports_sdpa = True368 369    def __init__(self, *inputs, **kwargs):370        super().__init__(*inputs, **kwargs)371 372    def _init_weights(self, module):373        """Initialize the weights."""374        if isinstance(module, (GPTBigCodeMLP, GPTBigCodeAttention)):375            # Reinitialize selected weights subject to the OpenAI GPT-2 Paper Scheme:376            #   > A modified initialization which accounts for the accumulation on the residual path with model depth. Scale377            #   > the weights of residual layers at initialization by a factor of 1/√N where N is the # of residual layers.378            #   >   -- GPT-2 :: https://openai.com/blog/better-language-models/379            #380            # Reference (Megatron-LM): https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/model/gpt_model.py381            module.c_proj.weight.data.normal_(382                mean=0.0, std=(self.config.initializer_range / math.sqrt(2 * self.config.n_layer))383            )384            module.c_proj._is_hf_initialized = True385        elif isinstance(module, nn.Linear):386            # Slightly different from the TF version which uses truncated_normal for initialization387            # cf https://github.com/pytorch/pytorch/pull/5617388            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)389            if module.bias is not None:390                module.bias.data.zero_()391        elif isinstance(module, nn.Embedding):392            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)393            if module.padding_idx is not None:394                module.weight.data[module.padding_idx].zero_()395        elif isinstance(module, nn.LayerNorm):396            module.bias.data.zero_()397            module.weight.data.fill_(1.0)398 399 400@auto_docstring401class GPTBigCodeModel(GPTBigCodePreTrainedModel):402    def __init__(self, config):403        super().__init__(config)404        self.multi_query = config.multi_query405        self.embed_dim = config.hidden_size406 407        self.wte = nn.Embedding(config.vocab_size, self.embed_dim)408        self.wpe = nn.Embedding(config.max_position_embeddings, self.embed_dim)409 410        self.drop = nn.Dropout(config.embd_pdrop)411        self.h = nn.ModuleList([GPTBigCodeBlock(config, layer_idx=i) for i in range(config.num_hidden_layers)])412        self.ln_f = nn.LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)413 414        max_positions = config.max_position_embeddings415        self.register_buffer(416            "bias", torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)), persistent=False417        )418 419        self.gradient_checkpointing = False420 421        # Initialize weights and apply final processing422        self.post_init()423 424    def get_input_embeddings(self):425        return self.wte426 427    def set_input_embeddings(self, new_embeddings):428        self.wte = new_embeddings429 430    @can_return_tuple431    @auto_docstring432    def forward(433        self,434        input_ids: Optional[torch.Tensor] = None,435        past_key_values: Optional[Cache] = None,436        attention_mask: Optional[torch.Tensor] = None,437        token_type_ids: Optional[torch.Tensor] = None,438        position_ids: Optional[torch.Tensor] = None,439        head_mask: Optional[torch.Tensor] = None,440        inputs_embeds: Optional[torch.Tensor] = None,441        encoder_hidden_states: Optional[torch.Tensor] = None,442        encoder_attention_mask: Optional[torch.Tensor] = None,443        use_cache: Optional[bool] = None,444        output_attentions: Optional[bool] = None,445        output_hidden_states: Optional[bool] = None,446        return_dict: Optional[bool] = None,447        cache_position: Optional[torch.Tensor] = None,448        **kwargs,449    ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:450        r"""451        input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`):452            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else453            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input454            sequence tokens in the vocabulary.455 456            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as457            `input_ids`.458 459            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and460            [`PreTrainedTokenizer.__call__`] for details.461 462            [What are input IDs?](../glossary#input-ids)463        """464        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions465        output_hidden_states = (466            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states467        )468        use_cache = use_cache if use_cache is not None else self.config.use_cache469        return_dict = return_dict if return_dict is not None else self.config.use_return_dict470 471        if (input_ids is None) ^ (inputs_embeds is not None):472            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")473        elif input_ids is not None:474            input_shape = input_ids.size()475            input_ids = input_ids.view(-1, input_shape[-1])476            batch_size = input_ids.shape[0]477        elif inputs_embeds is not None:478            input_shape = inputs_embeds.size()[:-1]479            batch_size = inputs_embeds.shape[0]480        else:481            raise ValueError("You have to specify either input_ids or inputs_embeds")482 483        if batch_size <= 0:484            raise ValueError("batch_size has to be defined and > 0")485 486        if use_cache and past_key_values is None:487            past_key_values = EncoderDecoderCache(DynamicCache(config=self.config), DynamicCache(config=self.config))488        if use_cache and isinstance(past_key_values, tuple):489            logger.warning_once(490                "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. "491                "You should pass an instance of `EncoderDecoderCache` instead, e.g. "492                "`past_key_values=EncoderDecoderCache.from_legacy_cache(past_key_values)`."493            )494            past_key_values = EncoderDecoderCache.from_legacy_cache(past_key_values)495 496        if inputs_embeds is None:497            inputs_embeds = self.wte(input_ids)498 499        if cache_position is None:500            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0501            cache_position = torch.arange(502                past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device503            )504 505        if position_ids is None:506            position_ids = cache_position.unsqueeze(0)507 508        causal_mask = create_causal_mask(509            config=self.config,510            input_embeds=inputs_embeds,511            attention_mask=attention_mask,512            cache_position=cache_position,513            position_ids=position_ids,514            past_key_values=past_key_values,515        )516 517        if self.config._attn_implementation == "flash_attention_2":518            encoder_attention_mask = (519                encoder_attention_mask.bool()520                if (encoder_attention_mask is not None and 0 in encoder_attention_mask)521                else None522            )523        else:524            # If a 2D or 3D attention mask is provided for the cross-attention525            # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]526            if (527                self.config.add_cross_attention528                and encoder_hidden_states is not None529                and encoder_attention_mask is not None530            ):531                if encoder_attention_mask.dim() == 2:532                    encoder_attention_mask.unsqueeze(1)533                assert encoder_attention_mask.dim() == 3534                encoder_attention_mask = encoder_attention_mask.bool().unsqueeze(2 if self.multi_query else 1)535            else:536                encoder_attention_mask = None537 538        # Prepare head mask if needed539        # 1.0 in head_mask indicate we keep the head540        # attention_probs has shape bsz x n_heads x N x N541        # head_mask has shape n_layer x batch x n_heads x N x N542        head_mask = self.get_head_mask(head_mask, self.config.n_layer)543 544        position_embeds = self.wpe(position_ids)545        hidden_states = inputs_embeds + position_embeds.to(inputs_embeds.device)546 547        if token_type_ids is not None:548            token_type_ids = token_type_ids.view(-1, input_shape[-1])549            token_type_embeds = self.wte(token_type_ids)550            hidden_states = hidden_states + token_type_embeds551 552        hidden_states = self.drop(hidden_states)553        output_shape = input_shape + (hidden_states.size(-1),)554 555        all_self_attentions = () if output_attentions else None556        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None557        all_hidden_states = () if output_hidden_states else None558        for i, block in enumerate(self.h):559            if output_hidden_states:560                all_hidden_states = all_hidden_states + (hidden_states,)561 562            outputs = block(563                hidden_states,564                past_key_values,565                causal_mask,566                head_mask[i],567                encoder_hidden_states,  # as a positional argument for gradient checkpointing568                encoder_attention_mask=encoder_attention_mask,569                use_cache=use_cache,570                output_attentions=output_attentions,571                cache_position=cache_position,572                **kwargs,573            )574 575            hidden_states = outputs[0]576            if output_attentions:577                all_self_attentions = all_self_attentions + (outputs[1],)578                if self.config.add_cross_attention:579                    all_cross_attentions = all_cross_attentions + (outputs[2],)580 581        hidden_states = self.ln_f(hidden_states)582 583        hidden_states = hidden_states.view(output_shape)584        # Add last hidden state585        if output_hidden_states:586            all_hidden_states = all_hidden_states + (hidden_states,)587 588        return BaseModelOutputWithPastAndCrossAttentions(589            last_hidden_state=hidden_states,590            past_key_values=past_key_values,591            hidden_states=all_hidden_states,592            attentions=all_self_attentions,593            cross_attentions=all_cross_attentions,594        )595 596 597@auto_docstring(598    custom_intro="""599    The GPT_BIGCODE Model transformer with a language modeling head on top (linear layer with weights tied to the input600    embeddings).601    """602)603class GPTBigCodeForCausalLM(GPTBigCodePreTrainedModel, GenerationMixin):604    _tied_weights_keys = ["lm_head.weight"]605 606    def __init__(self, config):607        super().__init__(config)608        self.transformer = GPTBigCodeModel(config)609        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)610 611        # Initialize weights and apply final processing612        self.post_init()613 614    @auto_docstring615    def forward(616        self,617        input_ids: Optional[torch.Tensor] = None,618        past_key_values: Optional[Cache] = None,619        attention_mask: Optional[torch.Tensor] = None,620        token_type_ids: Optional[torch.Tensor] = None,621        position_ids: Optional[torch.Tensor] = None,622        head_mask: Optional[torch.Tensor] = None,623        inputs_embeds: Optional[torch.Tensor] = None,624        encoder_hidden_states: Optional[torch.Tensor] = None,625        encoder_attention_mask: Optional[torch.Tensor] = None,626        labels: Optional[torch.Tensor] = None,627        use_cache: Optional[bool] = None,628        output_attentions: Optional[bool] = None,629        output_hidden_states: Optional[bool] = None,630        return_dict: Optional[bool] = None,631        cache_position: Optional[torch.Tensor] = None,632        **kwargs,633    ) -> Union[tuple, CausalLMOutputWithCrossAttentions]:634        r"""635        input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`):636            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else637            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input638            sequence tokens in the vocabulary.639 640            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as641            `input_ids`.642 643            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and644            [`PreTrainedTokenizer.__call__`] for details.645 646            [What are input IDs?](../glossary#input-ids)647        labels (`torch.Tensor` of shape `(batch_size, input_ids_length)`, *optional*):648            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set649            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`650            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`651        """652        return_dict = return_dict if return_dict is not None else self.config.use_return_dict653 654        transformer_outputs = self.transformer(655            input_ids,656            past_key_values=past_key_values,657            attention_mask=attention_mask,658            token_type_ids=token_type_ids,659            position_ids=position_ids,660            head_mask=head_mask,661            inputs_embeds=inputs_embeds,662            encoder_hidden_states=encoder_hidden_states,663            encoder_attention_mask=encoder_attention_mask,664            use_cache=use_cache,665            output_attentions=output_attentions,666            output_hidden_states=output_hidden_states,667            return_dict=return_dict,668            cache_position=cache_position,669        )670        hidden_states = transformer_outputs[0]671 672        lm_logits = self.lm_head(hidden_states)673 674        loss = None675        if labels is not None:676            loss = self.loss_function(677                lm_logits,678                labels,679                vocab_size=self.config.vocab_size,680                **kwargs,681            )682 683        if not return_dict:684            output = (lm_logits,) + transformer_outputs[1:]685            return ((loss,) + output) if loss is not None else output686 687        return CausalLMOutputWithCrossAttentions(688            loss=loss,689            logits=lm_logits,690            past_key_values=transformer_outputs.past_key_values,691            hidden_states=transformer_outputs.hidden_states,692            attentions=transformer_outputs.attentions,693            cross_attentions=transformer_outputs.cross_attentions,694        )695 696 697@auto_docstring(698    custom_intro="""699    The GPTBigCode Model transformer with a sequence classification head on top (linear layer).700 701    [`GPTBigCodeForSequenceClassification`] uses the last token in order to do the classification, as other causal702    models (e.g. GPT-1) do.703 704    Since it does classification on the last token, it requires to know the position of the last token. If a705    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If706    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the707    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in708    each row of the batch).709    """710)711class GPTBigCodeForSequenceClassification(GPTBigCodePreTrainedModel):712    def __init__(self, config):713        super().__init__(config)714        self.num_labels = config.num_labels715        self.transformer = GPTBigCodeModel(config)716        self.score = nn.Linear(config.n_embd, self.num_labels, bias=False)717 718        # Initialize weights and apply final processing719        self.post_init()720 721    @auto_docstring722    def forward(723        self,724        input_ids: Optional[torch.Tensor] = None,725        past_key_values: Optional[Cache] = None,726        attention_mask: Optional[torch.Tensor] = None,727        token_type_ids: Optional[torch.Tensor] = None,728        position_ids: Optional[torch.Tensor] = None,729        head_mask: Optional[torch.Tensor] = None,730        inputs_embeds: Optional[torch.Tensor] = None,731        labels: Optional[torch.Tensor] = None,732        use_cache: Optional[bool] = None,733        output_attentions: Optional[bool] = None,734        output_hidden_states: Optional[bool] = None,735        return_dict: Optional[bool] = None,736        **kwargs,737    ) -> Union[tuple, SequenceClassifierOutputWithPast]:738        r"""739        input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`):740            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else741            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input742            sequence tokens in the vocabulary.743 744            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as745            `input_ids`.746 747            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and748            [`PreTrainedTokenizer.__call__`] for details.749 750            [What are input IDs?](../glossary#input-ids)751        labels (`torch.Tensor` of shape `(batch_size,)`, *optional*):752            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,753            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If754            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).755        """756        return_dict = return_dict if return_dict is not None else self.config.use_return_dict757 758        transformer_outputs = self.transformer(759            input_ids,760            past_key_values=past_key_values,761            attention_mask=attention_mask,762            token_type_ids=token_type_ids,763            position_ids=position_ids,764            head_mask=head_mask,765            inputs_embeds=inputs_embeds,766            use_cache=use_cache,767            output_attentions=output_attentions,768            output_hidden_states=output_hidden_states,769            return_dict=return_dict,770            **kwargs,771        )772        hidden_states = transformer_outputs[0]773        logits = self.score(hidden_states)774 775        if input_ids is not None:776            batch_size, sequence_length = input_ids.shape[:2]777        else:778            batch_size, sequence_length = inputs_embeds.shape[:2]779 780        if self.config.pad_token_id is None and batch_size != 1:781            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")782        if self.config.pad_token_id is None:783            last_non_pad_token = -1784        elif input_ids is not None:785            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id786            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)787            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)788            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)789        else:790            last_non_pad_token = -1791            logger.warning_once(792                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "793                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"794            )795 796        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]797 798        loss = None799        if labels is not None:800            labels = labels.to(logits.device)801 802            if self.config.problem_type is None:803                if self.num_labels == 1:804                    self.config.problem_type = "regression"805                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):806                    self.config.problem_type = "single_label_classification"807                else:808                    self.config.problem_type = "multi_label_classification"809 810            if self.config.problem_type == "regression":811                loss_fct = MSELoss()812                if self.num_labels == 1:813                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())814                else:815                    loss = loss_fct(pooled_logits, labels)816            elif self.config.problem_type == "single_label_classification":817                loss_fct = CrossEntropyLoss()818                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))819            elif self.config.problem_type == "multi_label_classification":820                loss_fct = BCEWithLogitsLoss()821                loss = loss_fct(pooled_logits, labels)822        if not return_dict:823            output = (pooled_logits,) + transformer_outputs[1:]824            return ((loss,) + output) if loss is not None else output825 826        return SequenceClassifierOutputWithPast(827            loss=loss,828            logits=pooled_logits,829            past_key_values=transformer_outputs.past_key_values,830            hidden_states=transformer_outputs.hidden_states,831            attentions=transformer_outputs.attentions,832        )833 834 835@auto_docstring836class GPTBigCodeForTokenClassification(GPTBigCodePreTrainedModel):837    def __init__(self, config):838        super().__init__(config)839        self.num_labels = config.num_labels840 841        self.transformer = GPTBigCodeModel(config)842        if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:843            classifier_dropout = config.classifier_dropout844        elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:845            classifier_dropout = config.hidden_dropout846        else:847            classifier_dropout = 0.1848        self.dropout = nn.Dropout(classifier_dropout)849        self.classifier = nn.Linear(config.hidden_size, config.num_labels)850 851        # Initialize weights and apply final processing852        self.post_init()853 854    @auto_docstring855    def forward(856        self,857        input_ids: Optional[torch.Tensor] = None,858        past_key_values: Optional[Cache] = None,859        attention_mask: Optional[torch.Tensor] = None,860        token_type_ids: Optional[torch.Tensor] = None,861        position_ids: Optional[torch.Tensor] = None,862        head_mask: Optional[torch.Tensor] = None,863        inputs_embeds: Optional[torch.Tensor] = None,864        labels: Optional[torch.Tensor] = None,865        use_cache: Optional[bool] = None,866        output_attentions: Optional[bool] = None,867        output_hidden_states: Optional[bool] = None,868        return_dict: Optional[bool] = None,869    ) -> Union[tuple, TokenClassifierOutput]:870        r"""871        input_ids (`torch.Tensor` of shape `(batch_size, input_ids_length)`):872            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else873            `past_key_values.get_seq_length()` (`sequence_length` of input past key value states). Indices of input874            sequence tokens in the vocabulary.875 876            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as877            `input_ids`.878 879            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and880            [`PreTrainedTokenizer.__call__`] for details.881 882            [What are input IDs?](../glossary#input-ids)883        labels (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):884            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,885            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If886            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).887        """888        return_dict = return_dict if return_dict is not None else self.config.use_return_dict889 890        transformer_outputs = self.transformer(891            input_ids,892            past_key_values=past_key_values,893            attention_mask=attention_mask,894            token_type_ids=token_type_ids,895            position_ids=position_ids,896            head_mask=head_mask,897            inputs_embeds=inputs_embeds,898            use_cache=use_cache,899            output_attentions=output_attentions,900            output_hidden_states=output_hidden_states,901            return_dict=return_dict,902        )903 904        hidden_states = transformer_outputs[0]905        hidden_states = self.dropout(hidden_states)906        logits = self.classifier(hidden_states)907 908        loss = None909        if labels is not None:910            loss_fct = CrossEntropyLoss()911            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1).to(logits.device))912 913        if not return_dict:914            output = (logits,) + transformer_outputs[2:]915            return ((loss,) + output) if loss is not None else output916 917        return TokenClassifierOutput(918            loss=loss,919            logits=logits,920            hidden_states=transformer_outputs.hidden_states,921            attentions=transformer_outputs.attentions,922        )923 924 925__all__ = [926    "GPTBigCodeForSequenceClassification",927    "GPTBigCodeForTokenClassification",928    "GPTBigCodeForCausalLM",929    "GPTBigCodeModel",930    "GPTBigCodePreTrainedModel",931]932