Aluode/PerceptionLabPortable

1#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨2#           This file was automatically generated from src/transformers/models/biogpt/modular_biogpt.py.3#               Do NOT edit this file manually as any edits will be overwritten by the generation of4#             the file from the modular. If any change should be done, please apply the change to the5#                          modular_biogpt.py file directly. One of our CI enforces this.6#                🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨🚨7# coding=utf-88# Copyright 2022 The HuggingFace Team and Microsoft Research AI4Science All rights reserved.9#10# Licensed under the Apache License, Version 2.0 (the "License");11# you may not use this file except in compliance with the License.12# You may obtain a copy of the License at13#14#     http://www.apache.org/licenses/LICENSE-2.015#16# Unless required by applicable law or agreed to in writing, software17# distributed under the License is distributed on an "AS IS" BASIS,18# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.19# See the License for the specific language governing permissions and20# limitations under the License.21 22import math23from typing import Callable, Optional, Union24 25import torch26import torch.nn as nn27from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss28 29from ...activations import ACT2FN30from ...cache_utils import Cache, DynamicCache, EncoderDecoderCache31from ...generation import GenerationMixin32from ...modeling_attn_mask_utils import AttentionMaskConverter33from ...modeling_flash_attention_utils import FlashAttentionKwargs34from ...modeling_layers import GradientCheckpointingLayer35from ...modeling_outputs import (36    BaseModelOutputWithPastAndCrossAttentions,37    CausalLMOutputWithCrossAttentions,38    SequenceClassifierOutputWithPast,39    TokenClassifierOutput,40)41from ...modeling_utils import ALL_ATTENTION_FUNCTIONS, PreTrainedModel42from ...processing_utils import Unpack43from ...utils import TransformersKwargs, auto_docstring, is_torch_flex_attn_available, logging44from ...utils.deprecation import deprecate_kwarg45from .configuration_biogpt import BioGptConfig46 47 48if is_torch_flex_attn_available():49    from ...integrations.flex_attention import BlockMask, make_flex_block_causal_mask50 51 52logger = logging.get_logger(__name__)53 54 55class BioGptLearnedPositionalEmbedding(nn.Embedding):56    """57    This module learns positional embeddings up to a fixed maximum size.58    """59 60    def __init__(self, num_embeddings: int, embedding_dim: int):61        # BIOGPT is set up so that if padding_idx is specified then offset the embedding ids by 262        # and adjust num_embeddings appropriately. Other models don't have this hack63        self.offset = 264        super().__init__(num_embeddings + self.offset, embedding_dim)65 66    def forward(67        self,68        attention_mask: torch.LongTensor,69        past_key_values_length: int = 0,70        position_ids: Optional[torch.LongTensor] = None,71    ):72        """`input_ids_shape` is expected to be [bsz x seqlen]."""73 74        if position_ids is None:75            position_ids = torch.cumsum(attention_mask, dim=1)76            position_ids = (position_ids * attention_mask - 1).long()77            # cut positions if `past_key_values_length` is > 078            position_ids = position_ids[:, past_key_values_length:]79 80        return super().forward(position_ids + self.offset)81 82 83class BioGptScaledWordEmbedding(nn.Embedding):84    """85    This module overrides nn.Embeddings' forward by multiplying with embeddings scale.86    """87 88    def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0):89        super().__init__(num_embeddings, embedding_dim, padding_idx)90        self.embed_scale = embed_scale91 92    def forward(self, input_ids: torch.Tensor):93        return super().forward(input_ids) * self.embed_scale94 95 96def eager_attention_forward(97    module: nn.Module,98    query: torch.Tensor,99    key: torch.Tensor,100    value: torch.Tensor,101    attention_mask: Optional[torch.Tensor],102    scaling: Optional[float] = None,103    dropout: float = 0.0,104    head_mask: Optional[torch.Tensor] = None,105    **kwargs,106):107    if scaling is None:108        scaling = query.size(-1) ** -0.5109 110    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling111    if attention_mask is not None:112        attn_weights = attn_weights + attention_mask113 114    attn_weights = nn.functional.softmax(attn_weights, dim=-1)115 116    if head_mask is not None:117        attn_weights = attn_weights * head_mask.view(1, -1, 1, 1)118 119    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)120    attn_output = torch.matmul(attn_weights, value)121    attn_output = attn_output.transpose(1, 2).contiguous()122 123    return attn_output, attn_weights124 125 126class BioGptAttention(nn.Module):127    """Multi-headed attention from 'Attention Is All You Need' paper"""128 129    def __init__(130        self,131        embed_dim: int,132        num_heads: int,133        dropout: float = 0.0,134        is_decoder: bool = False,135        bias: bool = True,136        is_causal: bool = False,137        config: Optional[BioGptConfig] = None,138        layer_idx: Optional[int] = None,139    ):140        super().__init__()141        self.embed_dim = embed_dim142        self.num_heads = num_heads143        self.dropout = dropout144        self.head_dim = embed_dim // num_heads145        self.config = config146 147        if (self.head_dim * num_heads) != self.embed_dim:148            raise ValueError(149                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"150                f" and `num_heads`: {num_heads})."151            )152        self.scaling = self.head_dim**-0.5153        self.is_decoder = is_decoder154        self.is_causal = is_causal155        self.layer_idx = layer_idx156        if layer_idx is None and self.is_decoder:157            logger.warning_once(158                f"Instantiating a decoder {self.__class__.__name__} without passing `layer_idx` is not recommended and "159                "will lead to errors during the forward call, if caching is used. Please make sure to provide a `layer_idx` "160                "when creating this class."161            )162 163        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)164        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)165        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)166        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)167 168    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")169    def forward(170        self,171        hidden_states: torch.Tensor,172        key_value_states: Optional[torch.Tensor] = None,173        past_key_values: Optional[Cache] = None,174        attention_mask: Optional[torch.Tensor] = None,175        layer_head_mask: Optional[torch.Tensor] = None,176        output_attentions: bool = False,177        cache_position: Optional[torch.Tensor] = None,178        # TODO: we need a refactor so that the different attention modules can get their specific kwargs179        # ATM, we have mixed things encoder, decoder, and encoder-decoder attn180        **kwargs: Unpack[FlashAttentionKwargs],181    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:182        """Input shape: Batch x Time x Channel"""183 184        # if key_value_states are provided this layer is used as a cross-attention layer185        # for the decoder186        is_cross_attention = key_value_states is not None187 188        # determine input shapes189        bsz, tgt_len = hidden_states.shape[:-1]190        src_len = key_value_states.shape[1] if is_cross_attention else tgt_len191 192        q_input_shape = (bsz, tgt_len, -1, self.head_dim)193        kv_input_shape = (bsz, src_len, -1, self.head_dim)194 195        # get query proj196        query_states = self.q_proj(hidden_states).view(*q_input_shape).transpose(1, 2)197 198        is_updated = False199        if past_key_values is not None:200            if isinstance(past_key_values, EncoderDecoderCache):201                is_updated = past_key_values.is_updated.get(self.layer_idx)202                if is_cross_attention:203                    # after the first generated id, we can subsequently re-use all key/value_states from cache204                    curr_past_key_value = past_key_values.cross_attention_cache205                else:206                    curr_past_key_value = past_key_values.self_attention_cache207            else:208                curr_past_key_value = past_key_values209 210        current_states = key_value_states if is_cross_attention else hidden_states211        if is_cross_attention and past_key_values is not None and is_updated:212            # reuse k,v, cross_attentions213            key_states = curr_past_key_value.layers[self.layer_idx].keys214            value_states = curr_past_key_value.layers[self.layer_idx].values215        else:216            key_states = self.k_proj(current_states)217            value_states = self.v_proj(current_states)218            key_states = key_states.view(*kv_input_shape).transpose(1, 2)219            value_states = value_states.view(*kv_input_shape).transpose(1, 2)220 221            if past_key_values is not None:222                # save all key/value_states to cache to be re-used for fast auto-regressive generation223                cache_position = cache_position if not is_cross_attention else None224                key_states, value_states = curr_past_key_value.update(225                    key_states, value_states, self.layer_idx, {"cache_position": cache_position}226                )227                # set flag that curr layer for cross-attn is already updated so we can re-use in subsequent calls228                if is_cross_attention and isinstance(past_key_values, EncoderDecoderCache):229                    past_key_values.is_updated[self.layer_idx] = True230 231        attention_interface: Callable = eager_attention_forward232        if self.config._attn_implementation != "eager":233            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]234 235        attn_output, attn_weights = attention_interface(236            self,237            query_states,238            key_states,239            value_states,240            attention_mask,241            dropout=0.0 if not self.training else self.dropout,242            scaling=self.scaling,243            output_attentions=output_attentions,244            head_mask=layer_head_mask,245            **kwargs,246        )247 248        attn_output = attn_output.reshape(bsz, tgt_len, -1).contiguous()249        attn_output = self.out_proj(attn_output)250 251        return attn_output, attn_weights252 253 254class BioGptDecoderLayer(GradientCheckpointingLayer):255    def __init__(self, config: BioGptConfig, layer_idx: Optional[int] = None):256        super().__init__()257        self.embed_dim = config.hidden_size258 259        self.self_attn = BioGptAttention(260            embed_dim=self.embed_dim,261            num_heads=config.num_attention_heads,262            dropout=config.attention_probs_dropout_prob,263            is_decoder=True,264            is_causal=True,265            config=config,266            layer_idx=layer_idx,267        )268        self.dropout = config.hidden_dropout_prob269        self.activation_fn = ACT2FN[config.hidden_act]270        self.activation_dropout = config.activation_dropout271 272        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)273 274        self.fc1 = nn.Linear(self.embed_dim, config.intermediate_size)275        self.fc2 = nn.Linear(config.intermediate_size, self.embed_dim)276        self.final_layer_norm = nn.LayerNorm(self.embed_dim)277 278    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")279    def forward(280        self,281        hidden_states: torch.Tensor,282        attention_mask: Optional[torch.Tensor] = None,283        layer_head_mask: Optional[torch.Tensor] = None,284        past_key_values: Optional[Cache] = None,285        output_attentions: Optional[bool] = False,286        use_cache: Optional[bool] = True,287        position_ids: Optional[torch.LongTensor] = None,288        cache_position: Optional[torch.Tensor] = None,289        **kwargs: Unpack[TransformersKwargs],290    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:291        """292        Args:293            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`294            attention_mask (`torch.FloatTensor`): attention mask of size295                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.296            layer_head_mask (`torch.FloatTensor`): mask for attention heads in a given layer of size297                `(encoder_attention_heads,)`.298            past_key_values (`Cache`): cached past key and value projection states299            output_attentions (`bool`, *optional*):300                Whether or not to return the attentions tensors of all attention layers. See `attentions` under301                returned tensors for more detail.302            use_cache (`bool`, *optional*):303                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding304                (see `past_key_values`).305            cache_position (`torch.LongTensor` of shape `(sequence_length)`, *optional*):306                Indices depicting the position of the input sequence tokens in the sequence. It is used to update the307                cache in the correct position and to infer the complete sequence length.308        """309        residual = hidden_states310 311        hidden_states = self.self_attn_layer_norm(hidden_states)312 313        # Self Attention314        hidden_states, self_attn_weights = self.self_attn(315            hidden_states=hidden_states,316            past_key_values=past_key_values,317            attention_mask=attention_mask,318            layer_head_mask=layer_head_mask,319            output_attentions=output_attentions,320            position_ids=position_ids,321            cache_position=cache_position,322            **kwargs,323        )324        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)325        hidden_states = residual + hidden_states326 327        # Fully Connected328        residual = hidden_states329        hidden_states = self.final_layer_norm(hidden_states)330        hidden_states = self.fc1(hidden_states)331        hidden_states = self.activation_fn(hidden_states)332        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)333        hidden_states = self.fc2(hidden_states)334        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)335        hidden_states = residual + hidden_states336 337        outputs = (hidden_states,)338 339        if output_attentions:340            outputs += (self_attn_weights,)341 342        return outputs343 344 345@auto_docstring346class BioGptPreTrainedModel(PreTrainedModel):347    config: BioGptConfig348    base_model_prefix = "biogpt"349    supports_gradient_checkpointing = True350    _supports_flash_attn = True351    _supports_sdpa = True352    _supports_flex_attn = True353 354    _can_compile_fullgraph = True355 356    # Copied from transformers.models.bart.modeling_bart.BartPreTrainedModel._update_causal_mask357    def _update_causal_mask(358        self,359        attention_mask: Optional[Union[torch.Tensor, "BlockMask"]],360        input_tensor: torch.Tensor,361        cache_position: torch.Tensor,362        past_key_values: Cache,363    ):364        if self.config._attn_implementation == "flex_attention":365            if isinstance(attention_mask, torch.Tensor):366                attention_mask = make_flex_block_causal_mask(attention_mask)367            # Other attention flavors support in-built causal (when `mask is None`)368            # while we need to create our specific block mask regardless369            elif attention_mask is None:370                attention_mask = make_flex_block_causal_mask(371                    torch.ones(372                        size=(input_tensor.shape[0], input_tensor.shape[1]),373                        device=attention_mask.device,374                    )375                )376            return attention_mask377 378        if self.config._attn_implementation == "flash_attention_2":379            if attention_mask is not None and (attention_mask == 0.0).any():380                return attention_mask381            return None382 383        # For SDPA, when possible, we will rely on its `is_causal` argument instead of its `attn_mask` argument, in384        # order to dispatch on Flash Attention 2. This feature is not compatible with static cache, as SDPA will fail385        # to infer the attention mask.386        past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0387        using_compilable_cache = past_key_values.is_compileable if past_key_values is not None else False388 389        # When output attentions is True, sdpa implementation's forward method calls the eager implementation's forward390        if self.config._attn_implementation == "sdpa" and not using_compilable_cache:391            if AttentionMaskConverter._ignore_causal_mask_sdpa(392                attention_mask,393                inputs_embeds=input_tensor,394                past_key_values_length=past_seen_tokens,395                is_training=self.training,396            ):397                return None398 399        dtype = input_tensor.dtype400        sequence_length = input_tensor.shape[1]401        if using_compilable_cache:402            target_length = past_key_values.get_max_cache_shape()403        else:404            target_length = (405                attention_mask.shape[-1]406                if isinstance(attention_mask, torch.Tensor)407                else past_seen_tokens + sequence_length + 1408            )409 410        # In case the provided `attention` mask is 2D, we generate a causal mask here (4D).411        causal_mask = self._prepare_4d_causal_attention_mask_with_cache_position(412            attention_mask,413            sequence_length=sequence_length,414            target_length=target_length,415            dtype=dtype,416            cache_position=cache_position,417            batch_size=input_tensor.shape[0],418        )419 420        if (421            self.config._attn_implementation == "sdpa"422            and attention_mask is not None423            and attention_mask.device.type in ["cuda", "xpu", "npu"]424        ):425            # Attend to all tokens in fully masked rows in the causal_mask, for example the relevant first rows when426            # using left padding. This is required by F.scaled_dot_product_attention memory-efficient attention path.427            # Details: https://github.com/pytorch/pytorch/issues/110213428            min_dtype = torch.finfo(dtype).min429            causal_mask = AttentionMaskConverter._unmask_unattended(causal_mask, min_dtype)430 431        return causal_mask432 433    @staticmethod434    # Copied from transformers.models.gptj.modeling_gptj.GPTJModel._prepare_4d_causal_attention_mask_with_cache_position435    def _prepare_4d_causal_attention_mask_with_cache_position(436        attention_mask: torch.Tensor,437        sequence_length: int,438        target_length: int,439        dtype: torch.dtype,440        cache_position: torch.Tensor,441        batch_size: int,442        **kwargs,443    ):444        """445        Creates a causal 4D mask of shape `(batch_size, 1, query_length, key_value_length)` from a 2D mask of shape446        `(batch_size, key_value_length)`, or if the input `attention_mask` is already 4D, do nothing.447 448        Args:449            attention_mask (`torch.Tensor`):450                A 2D attention mask of shape `(batch_size, key_value_length)` or a 4D attention mask of shape451                `(batch_size, 1, query_length, key_value_length)`.452            sequence_length (`int`):453                The sequence length being processed.454            target_length (`int`):455                The target length: when generating with static cache, the mask should be as long as the static cache,456                to account for the 0 padding, the part of the cache that is not filled yet.457            dtype (`torch.dtype`):458                The dtype to use for the 4D attention mask.459            cache_position (`torch.Tensor`):460                Indices depicting the position of the input sequence tokens in the sequence.461            batch_size (`torch.Tensor`):462                Batch size.463        """464        if attention_mask is not None and attention_mask.dim() == 4:465            # In this case we assume that the mask comes already in inverted form and requires no inversion or slicing.466            causal_mask = attention_mask467        else:468            min_dtype = torch.finfo(dtype).min469            causal_mask = torch.full(470                (sequence_length, target_length), fill_value=min_dtype, dtype=dtype, device=cache_position.device471            )472            if sequence_length != 1:473                causal_mask = torch.triu(causal_mask, diagonal=1)474            causal_mask *= torch.arange(target_length, device=cache_position.device) > cache_position.reshape(-1, 1)475            causal_mask = causal_mask[None, None, :, :].expand(batch_size, 1, -1, -1)476            if attention_mask is not None:477                causal_mask = causal_mask.clone()  # copy to contiguous memory for in-place edit478                mask_length = attention_mask.shape[-1]479                padding_mask = causal_mask[:, :, :, :mask_length] + attention_mask[:, None, None, :].to(480                    causal_mask.device481                )482                padding_mask = padding_mask == 0483                causal_mask[:, :, :, :mask_length] = causal_mask[:, :, :, :mask_length].masked_fill(484                    padding_mask, min_dtype485                )486 487        return causal_mask488 489 490@auto_docstring491class BioGptModel(BioGptPreTrainedModel):492    def __init__(self, config: BioGptConfig):493        super().__init__(config)494        self.config = config495        self.layerdrop = config.layerdrop496        self.dropout = config.hidden_dropout_prob497        self.embed_dim = config.hidden_size498        self.padding_idx = config.pad_token_id499        embed_scale = math.sqrt(config.hidden_size) if config.scale_embedding else 1.0500 501        self.embed_tokens = BioGptScaledWordEmbedding(502            config.vocab_size, self.embed_dim, self.padding_idx, embed_scale=embed_scale503        )504        self.embed_positions = BioGptLearnedPositionalEmbedding(config.max_position_embeddings, self.embed_dim)505 506        self.layers = nn.ModuleList([BioGptDecoderLayer(config, layer_idx=i) for i in range(config.num_hidden_layers)])507        self.layer_norm = nn.LayerNorm(self.embed_dim)508 509        self.gradient_checkpointing = False510        # Initialize weights and apply final processing511        self.post_init()512 513    @auto_docstring514    def forward(515        self,516        input_ids: Optional[torch.LongTensor] = None,517        attention_mask: Optional[torch.FloatTensor] = None,518        head_mask: Optional[torch.FloatTensor] = None,519        inputs_embeds: Optional[torch.FloatTensor] = None,520        past_key_values: Optional[Cache] = None,521        use_cache: Optional[bool] = None,522        position_ids: Optional[torch.LongTensor] = None,523        output_attentions: Optional[bool] = None,524        output_hidden_states: Optional[bool] = None,525        return_dict: Optional[bool] = None,526        cache_position: Optional[torch.Tensor] = None,527        **kwargs: Unpack[TransformersKwargs],528    ) -> Union[tuple, BaseModelOutputWithPastAndCrossAttentions]:529        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions530        output_hidden_states = (531            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states532        )533        use_cache = use_cache if use_cache is not None else self.config.use_cache534        return_dict = return_dict if return_dict is not None else self.config.use_return_dict535 536        # retrieve input_ids and inputs_embeds537        if (input_ids is None) ^ (inputs_embeds is not None):538            raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")539        elif input_ids is not None:540            input = input_ids541            input_shape = input.shape542            input_ids = input_ids.view(-1, input_shape[-1])543        elif inputs_embeds is not None:544            input_shape = inputs_embeds.size()[:-1]545            input = inputs_embeds[:, :, -1]546        else:547            raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")548 549        if inputs_embeds is None:550            inputs_embeds = self.embed_tokens(input)551 552        if self.gradient_checkpointing and self.training:553            if use_cache:554                logger.warning_once(555                    "`use_cache=True` is incompatible with gradient checkpointing`. Setting `use_cache=False`..."556                )557                use_cache = False558 559        # initialize past_key_values560        if use_cache and past_key_values is None:561            past_key_values = DynamicCache(config=self.config)562        if use_cache and isinstance(past_key_values, tuple):563            logger.warning_once(564                "Passing a tuple of `past_key_values` is deprecated and will be removed in Transformers v4.58.0. "565                "You should pass an instance of `DynamicCache` instead, e.g. "566                "`past_key_values=DynamicCache.from_legacy_cache(past_key_values)`."567            )568            past_key_values = DynamicCache.from_legacy_cache(past_key_values)569 570        batch_size, seq_length = inputs_embeds.size()[:-1]571        past_key_values_length = past_key_values.get_seq_length() if past_key_values is not None else 0572        if cache_position is None:573            cache_position = torch.arange(574                past_key_values_length, past_key_values_length + seq_length, device=inputs_embeds.device575            )576 577        if attention_mask is None:578            # required mask seq length can be calculated via length of past cache579            mask_seq_length = past_key_values_length + seq_length580            attention_mask = torch.ones(batch_size, mask_seq_length, device=inputs_embeds.device)581 582        self_attn_cache = past_key_values583 584        causal_mask = self._update_causal_mask(585            attention_mask,586            inputs_embeds,587            cache_position,588            self_attn_cache,589        )590 591        # embed positions592        if position_ids is None:593            # position_ids = cache_position.unsqueeze(0)594            position_ids = torch.cumsum(attention_mask, dim=1)595            position_ids = (position_ids * attention_mask - 1).long()596            # cut positions if `past_seen_tokens` is > 0597            position_ids = position_ids[:, past_key_values_length:]598 599        positions = self.embed_positions(attention_mask, past_key_values_length, position_ids=position_ids)600        hidden_states = inputs_embeds + positions601        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)602 603        if self.gradient_checkpointing and self.training:604            if use_cache:605                logger.warning_once(606                    "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."607                )608                use_cache = False609 610        all_hidden_states = () if output_hidden_states else None611        all_self_attns = () if output_attentions else None612        all_cross_attentions = None613 614        for idx, decoder_layer in enumerate(self.layers):615            # add LayerDrop (see https://huggingface.co/papers/1909.11556 for description)616            if output_hidden_states:617                all_hidden_states += (hidden_states,)618            if self.training:619                dropout_probability = torch.rand([])620                if dropout_probability < self.layerdrop:621                    continue622 623            layer_outputs = decoder_layer(624                hidden_states,625                attention_mask=causal_mask,626                layer_head_mask=(head_mask[idx] if head_mask is not None else None),627                past_key_values=past_key_values,628                output_attentions=output_attentions,629                use_cache=use_cache,630                position_ids=position_ids,631                cache_position=cache_position,632                **kwargs,633            )634 635            hidden_states = layer_outputs[0]636 637            if output_attentions:638                all_self_attns += (layer_outputs[1],)639 640        # add hidden states from the last decoder layer641        if output_hidden_states:642            all_hidden_states += (hidden_states,)643 644        hidden_states = self.layer_norm(hidden_states)645 646        if not return_dict:647            return tuple(648                v649                for v in [hidden_states, past_key_values, all_hidden_states, all_self_attns, all_cross_attentions]650                if v is not None651            )652        return BaseModelOutputWithPastAndCrossAttentions(653            last_hidden_state=hidden_states,654            past_key_values=past_key_values,655            hidden_states=all_hidden_states,656            attentions=all_self_attns,657            cross_attentions=all_cross_attentions,658        )659 660 661@auto_docstring(662    custom_intro="""663    BioGPT Model with a `language modeling` head on top for CLM fine-tuning.664    """665)666class BioGptForCausalLM(BioGptPreTrainedModel, GenerationMixin):667    _tied_weights_keys = ["output_projection.weight"]668 669    def __init__(self, config):670        super().__init__(config)671 672        self.biogpt = BioGptModel(config)673        self.output_projection = nn.Linear(config.hidden_size, config.vocab_size, bias=False)674 675        # Initialize weights and apply final processing676        self.post_init()677 678    def get_output_embeddings(self):679        return self.output_projection680 681    def set_output_embeddings(self, new_embeddings):682        self.output_projection = new_embeddings683 684    @auto_docstring685    def forward(686        self,687        input_ids: Optional[torch.LongTensor] = None,688        attention_mask: Optional[torch.FloatTensor] = None,689        head_mask: Optional[torch.FloatTensor] = None,690        inputs_embeds: Optional[torch.FloatTensor] = None,691        past_key_values: Optional[Cache] = None,692        labels: Optional[torch.LongTensor] = None,693        use_cache: Optional[bool] = None,694        position_ids: Optional[torch.LongTensor] = None,695        output_attentions: Optional[bool] = None,696        output_hidden_states: Optional[bool] = None,697        return_dict: Optional[bool] = None,698        cache_position: Optional[torch.Tensor] = None,699        **kwargs: Unpack[TransformersKwargs],700    ) -> Union[tuple, CausalLMOutputWithCrossAttentions]:701        r"""702        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):703            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set704            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`705            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`706        """707        return_dict = return_dict if return_dict is not None else self.config.use_return_dict708 709        outputs = self.biogpt(710            input_ids,711            attention_mask=attention_mask,712            head_mask=head_mask,713            inputs_embeds=inputs_embeds,714            past_key_values=past_key_values,715            use_cache=use_cache,716            position_ids=position_ids,717            output_attentions=output_attentions,718            output_hidden_states=output_hidden_states,719            return_dict=return_dict,720            cache_position=cache_position,721            **kwargs,722        )723 724        sequence_output = outputs[0]725        prediction_scores = self.output_projection(sequence_output)726 727        lm_loss = None728        if labels is not None:729            lm_loss = self.loss_function(730                prediction_scores,731                labels,732                vocab_size=self.config.vocab_size,733                **kwargs,734            )735 736        if not return_dict:737            output = (prediction_scores,) + outputs[1:]738            return ((lm_loss,) + output) if lm_loss is not None else output739 740        return CausalLMOutputWithCrossAttentions(741            loss=lm_loss,742            logits=prediction_scores,743            past_key_values=outputs.past_key_values,744            hidden_states=outputs.hidden_states,745            attentions=outputs.attentions,746            cross_attentions=outputs.cross_attentions,747        )748 749 750@auto_docstring751class BioGptForTokenClassification(BioGptPreTrainedModel):752    def __init__(self, config):753        super().__init__(config)754        self.num_labels = config.num_labels755 756        self.biogpt = BioGptModel(config)757        if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:758            classifier_dropout = config.classifier_dropout759        else:760            classifier_dropout = config.hidden_dropout_prob761        self.dropout = nn.Dropout(classifier_dropout)762        self.classifier = nn.Linear(config.hidden_size, config.num_labels)763 764        self.post_init()765 766    @auto_docstring767    def forward(768        self,769        input_ids: Optional[torch.LongTensor] = None,770        token_type_ids: Optional[torch.LongTensor] = None,771        attention_mask: Optional[torch.FloatTensor] = None,772        head_mask: Optional[torch.FloatTensor] = None,773        past_key_values: Optional[Cache] = None,774        inputs_embeds: Optional[torch.FloatTensor] = None,775        labels: Optional[torch.LongTensor] = None,776        use_cache: Optional[bool] = None,777        position_ids: Optional[torch.LongTensor] = None,778        output_attentions: Optional[bool] = None,779        output_hidden_states: Optional[bool] = None,780        return_dict: Optional[bool] = None,781        cache_position: Optional[torch.Tensor] = None,782    ) -> Union[tuple, TokenClassifierOutput]:783        r"""784        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):785            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,786            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If787            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).788        """789        return_dict = return_dict if return_dict is not None else self.config.use_return_dict790 791        transformer_outputs = self.biogpt(792            input_ids,793            past_key_values=past_key_values,794            attention_mask=attention_mask,795            head_mask=head_mask,796            inputs_embeds=inputs_embeds,797            use_cache=use_cache,798            position_ids=position_ids,799            output_attentions=output_attentions,800            output_hidden_states=output_hidden_states,801            return_dict=return_dict,802            cache_position=cache_position,803        )804 805        hidden_states = transformer_outputs[0]806        hidden_states = self.dropout(hidden_states)807        logits = self.classifier(hidden_states)808 809        loss = None810        if labels is not None:811            loss_fct = CrossEntropyLoss()812            # Only keep active parts of the loss813            if attention_mask is not None:814                active_loss = attention_mask.view(-1) == 1815                active_logits = logits.view(-1, self.num_labels)816                active_labels = torch.where(817                    active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)818                )819                loss = loss_fct(active_logits, active_labels)820            else:821                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))822 823        if not return_dict:824            output = (logits,) + transformer_outputs[2:]825            return ((loss,) + output) if loss is not None else output826 827        return TokenClassifierOutput(828            loss=loss,829            logits=logits,830            hidden_states=transformer_outputs.hidden_states,831            attentions=transformer_outputs.attentions,832        )833 834 835@auto_docstring(836    custom_intro="""837    The BioGpt Model transformer with a sequence classification head on top (linear layer).838 839    [`BioGptForSequenceClassification`] uses the last token in order to do the classification, as other causal models840    (e.g. GPT-2) do.841 842    Since it does classification on the last token, it is required to know the position of the last token. If a843    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If844    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the845    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in846    each row of the batch).847    """848)849class BioGptForSequenceClassification(BioGptPreTrainedModel):850    def __init__(self, config: BioGptConfig):851        super().__init__(config)852        self.num_labels = config.num_labels853        self.biogpt = BioGptModel(config)854        self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)855 856        # Initialize weights and apply final processing857        self.post_init()858 859    @auto_docstring860    def forward(861        self,862        input_ids: Optional[torch.LongTensor] = None,863        attention_mask: Optional[torch.FloatTensor] = None,864        head_mask: Optional[torch.FloatTensor] = None,865        past_key_values: Optional[Cache] = None,866        inputs_embeds: Optional[torch.FloatTensor] = None,867        labels: Optional[torch.LongTensor] = None,868        use_cache: Optional[bool] = None,869        position_ids: Optional[torch.LongTensor] = None,870        output_attentions: Optional[bool] = None,871        output_hidden_states: Optional[bool] = None,872        return_dict: Optional[bool] = None,873        cache_position: Optional[torch.Tensor] = None,874        logits_to_keep: Union[int, torch.Tensor] = 0,875    ) -> Union[tuple, SequenceClassifierOutputWithPast]:876        r"""877        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):878            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,879            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If880            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).881        """882        return_dict = return_dict if return_dict is not None else self.config.use_return_dict883 884        transformer_outputs = self.biogpt(885            input_ids,886            past_key_values=past_key_values,887            attention_mask=attention_mask,888            head_mask=head_mask,889            inputs_embeds=inputs_embeds,890            use_cache=use_cache,891            position_ids=position_ids,892            output_attentions=output_attentions,893            output_hidden_states=output_hidden_states,894            return_dict=return_dict,895            cache_position=cache_position,896        )897        hidden_states = transformer_outputs[0]898        slice_indices = slice(-logits_to_keep, None) if isinstance(logits_to_keep, int) else logits_to_keep899        logits = self.score(hidden_states[:, slice_indices, :])900 901        if input_ids is not None:902            batch_size, sequence_length = input_ids.shape[:2]903        else:904            batch_size, sequence_length = inputs_embeds.shape[:2]905 906        if self.config.pad_token_id is None:907            sequence_length = -1908        else:909            if input_ids is not None:910                sequence_length = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)911            else:912                sequence_length = -1913                logger.warning_once(914                    f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "915                    "unexpected if using padding tokens in conjunction with `inputs_embeds.`"916                )917 918        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_length]919 920        loss = None921        if labels is not None:922            if self.config.problem_type is None:923                if self.num_labels == 1:924                    self.config.problem_type = "regression"925                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):926                    self.config.problem_type = "single_label_classification"927                else:928                    self.config.problem_type = "multi_label_classification"929 930            if self.config.problem_type == "regression":931                loss_fct = MSELoss()932                if self.num_labels == 1:933                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())934                else:935                    loss = loss_fct(pooled_logits, labels)936            elif self.config.problem_type == "single_label_classification":937                loss_fct = CrossEntropyLoss()938                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))939            elif self.config.problem_type == "multi_label_classification":940                loss_fct = BCEWithLogitsLoss()941                loss = loss_fct(pooled_logits, labels)942        if not return_dict:943            output = (pooled_logits,) + transformer_outputs[1:]944            return ((loss,) + output) if loss is not None else output945 946        return SequenceClassifierOutputWithPast(947            loss=loss,948            logits=pooled_logits,949            past_key_values=transformer_outputs.past_key_values,950            hidden_states=transformer_outputs.hidden_states,951            attentions=transformer_outputs.attentions,952        )953 954    def get_input_embeddings(self):955        return self.biogpt.embed_tokens956 957    def set_input_embeddings(self, value):958        self.biogpt.embed_tokens = value959 960 961__all__ = [962    "BioGptForCausalLM",963    "BioGptForTokenClassification",964    "BioGptForSequenceClassification",965    "BioGptModel",966    "BioGptPreTrainedModel",967]968