Aluode/PerceptionLabPortable

1# Copyright 2025 Johns Hopkins University, LightOn, and the HuggingFace Inc. team. All rights reserved.2#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 16import math17from collections.abc import Callable18from typing import Optional, Union19 20import torch21from torch import nn22from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss23 24from ...cache_utils import Cache, DynamicCache25from ...configuration_utils import PretrainedConfig26from ...generation import GenerationMixin27from ...masking_utils import create_causal_mask, create_sliding_window_causal_mask28from ...modeling_layers import GradientCheckpointingLayer29from ...modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast30from ...modeling_utils import ALL_ATTENTION_FUNCTIONS31from ...processing_utils import Unpack32from ...utils import TransformersKwargs, auto_docstring, can_return_tuple, logging33from ...utils.deprecation import deprecate_kwarg34from ...utils.generic import check_model_inputs35from ..modernbert.modeling_modernbert import (36    ModernBertEmbeddings,37    ModernBertMLP,38    ModernBertPredictionHead,39    ModernBertPreTrainedModel,40    ModernBertRotaryEmbedding,41    apply_rotary_pos_emb,42)43 44 45logger = logging.get_logger(__name__)46 47 48class ModernBertDecoderConfig(PretrainedConfig):49    r"""50    This is the configuration class to store the configuration of a [`ModernBertDecoderModel`]. It is used to instantiate a ModernBert51    decoder model according to the specified arguments, defining the model architecture. Instantiating a configuration with the52    defaults will yield a similar configuration to that of the ModernBERT-base decoder.53    e.g. [blab-jhu/test-32m-dec](https://huggingface.co/blab-jhu/test-32m-dec)54 55    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the56    documentation from [`PretrainedConfig`] for more information.57 58    Args:59        vocab_size (`int`, *optional*, defaults to 50368):60            Vocabulary size of the ModernBert decoder model. Defines the number of different tokens that can be represented by the61            `inputs_ids` passed when calling [`ModernBertDecoderModel`]62        hidden_size (`int`, *optional*, defaults to 768):63            Dimension of the hidden representations.64        intermediate_size (`int`, *optional*, defaults to 1152):65            Dimension of the MLP representations.66        num_hidden_layers (`int`, *optional*, defaults to 22):67            Number of hidden layers in the Transformer decoder.68        num_attention_heads (`int`, *optional*, defaults to 12):69            Number of attention heads for each attention layer in the Transformer decoder.70        hidden_activation (`str` or `function`, *optional*, defaults to `"gelu"`):71            The non-linear activation function (function or string) in the decoder. Will default to `"gelu"`72            if not specified.73        max_position_embeddings (`int`, *optional*, defaults to 8192):74            The maximum sequence length that this model might ever be used with.75        initializer_range (`float`, *optional*, defaults to 0.02):76            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.77        initializer_cutoff_factor (`float`, *optional*, defaults to 2.0):78            The cutoff factor for the truncated_normal_initializer for initializing all weight matrices.79        norm_eps (`float`, *optional*, defaults to 1e-05):80            The epsilon used by the rms normalization layers.81        norm_bias (`bool`, *optional*, defaults to `False`):82            Whether to use bias in the normalization layers.83        pad_token_id (`int`, *optional*, defaults to 50283):84            Padding token id.85        eos_token_id (`int`, *optional*, defaults to 50282):86            End of stream token id.87        bos_token_id (`int`, *optional*, defaults to 50281):88            Beginning of stream token id.89        cls_token_id (`int`, *optional*, defaults to 50281):90            Classification token id.91        sep_token_id (`int`, *optional*, defaults to 50282):92            Separation token id.93        global_rope_theta (`float`, *optional*, defaults to 160000.0):94            The base period of the global RoPE embeddings.95        attention_bias (`bool`, *optional*, defaults to `False`):96            Whether to use a bias in the query, key, value and output projection layers during self-attention.97        attention_dropout (`float`, *optional*, defaults to 0.0):98            The dropout ratio for the attention probabilities.99        embedding_dropout (`float`, *optional*, defaults to 0.0):100            The dropout ratio for the embeddings.101        mlp_bias (`bool`, *optional*, defaults to `False`):102            Whether to use bias in the MLP layers.103        mlp_dropout (`float`, *optional*, defaults to 0.0):104            The dropout ratio for the MLP layers.105        decoder_bias (`bool`, *optional*, defaults to `True`):106            Whether to use bias in the decoder layers.107        classifier_dropout (`float`, *optional*, defaults to 0.0):108            The dropout ratio for the classifier.109        classifier_bias (`bool`, *optional*, defaults to `False`):110            Whether to use bias in the classifier.111        classifier_activation (`str`, *optional*, defaults to `"gelu"`):112            The activation function for the classifier.113        use_cache (`bool`, *optional*, defaults to `True`):114            Whether or not the model should return the last key/values attentions (not used by all models). Only115            relevant if `config.is_decoder=True`.116        local_attention (`int`, *optional*, defaults to 128):117            The sliding window size for local attention. Only used for layers that use local attention. Note that for118            the decoder to match ModernBERT this is actually half of the sliding window size, so 128 => 64.119        global_attn_every_n_layers (`int`, *optional*, defaults to 3):120            Every `global_attn_every_n_layers` layers will use global attention instead of local attention.121        local_rope_theta (`float`, *optional*, defaults to 160000.0):122            The base period of the local RoPE embeddings. If not specified, defaults to 160000.0123        layer_types (`list`, *optional*):124            List of layer types, one for each layer. If not specified, will be automatically generated based on125            `global_attn_every_n_layers`. Should contain "full_attention" or "sliding_attention".126 127    Examples:128 129    ```python130    >>> from transformers import ModernBertDecoderModel, ModernBertDecoderConfig131 132    >>> # Initializing a ModernBert decoder style configuration133    >>> configuration = ModernBertDecoderConfig()134 135    >>> # Initializing a model from the modernbert-base decoder style configuration136    >>> model = ModernBertDecoderModel(configuration)137 138    >>> # Accessing the model configuration139    >>> configuration = model.config140    ```"""141 142    model_type = "modernbert-decoder"143    attribute_map = {"rope_theta": "global_rope_theta"}144    keys_to_ignore_at_inference = ["past_key_values"]145 146    def __init__(147        self,148        vocab_size=50368,149        hidden_size=768,150        intermediate_size=1152,151        num_hidden_layers=22,152        num_attention_heads=12,153        hidden_activation="gelu",154        max_position_embeddings=8192,155        initializer_range=0.02,156        initializer_cutoff_factor=2.0,157        norm_eps=1e-5,158        norm_bias=False,159        pad_token_id=50283,160        eos_token_id=50282,161        bos_token_id=50281,162        cls_token_id=50281,163        sep_token_id=50282,164        global_rope_theta=160000.0,165        attention_bias=False,166        attention_dropout=0.0,167        embedding_dropout=0.0,168        mlp_bias=False,169        mlp_dropout=0.0,170        decoder_bias=True,171        classifier_dropout=0.0,172        classifier_bias=False,173        classifier_activation="gelu",174        use_cache=True,175        local_attention=128,176        global_attn_every_n_layers=3,177        local_rope_theta=160000.0,178        layer_types=None,179        **kwargs,180    ):181        super().__init__(182            pad_token_id=pad_token_id,183            bos_token_id=bos_token_id,184            eos_token_id=eos_token_id,185            cls_token_id=cls_token_id,186            sep_token_id=sep_token_id,187            **kwargs,188        )189        self.vocab_size = vocab_size190        self.max_position_embeddings = max_position_embeddings191        self.hidden_size = hidden_size192        self.intermediate_size = intermediate_size193        self.num_hidden_layers = num_hidden_layers194        self.num_attention_heads = num_attention_heads195        self.initializer_range = initializer_range196        self.initializer_cutoff_factor = initializer_cutoff_factor197        self.norm_eps = norm_eps198        self.norm_bias = norm_bias199        self.global_rope_theta = global_rope_theta200        self.attention_bias = attention_bias201        self.attention_dropout = attention_dropout202        self.hidden_activation = hidden_activation203        self.embedding_dropout = embedding_dropout204        self.mlp_bias = mlp_bias205        self.mlp_dropout = mlp_dropout206        self.decoder_bias = decoder_bias207        self.classifier_dropout = classifier_dropout208        self.classifier_bias = classifier_bias209        self.classifier_activation = classifier_activation210        self.use_cache = use_cache211        self.global_attn_every_n_layers = global_attn_every_n_layers212        self.local_rope_theta = local_rope_theta213        # for consistency with ModernBert214        self.reference_compile = False215 216        # Set up layer_types for standardized layer type detection217        self.layer_types = layer_types218        if self.layer_types is None:219            # Create layer_types based on the alternating pattern220            self.layer_types = []221            for layer_id in range(num_hidden_layers):222                if layer_id % global_attn_every_n_layers != 0:223                    self.layer_types.append("sliding_attention")224                else:225                    self.layer_types.append("full_attention")226 227        # NOTE: sliding window numbers matches ModernBERT but is only half of it228        self.sliding_window = local_attention // 2 if local_attention else -1229 230 231class ModernBertDecoderEmbeddings(ModernBertEmbeddings):232    pass233 234 235class ModernBertDecoderMLP(ModernBertMLP):236    pass237 238 239class ModernBertDecoderRotaryEmbedding(ModernBertRotaryEmbedding):240    pass241 242 243def eager_attention_forward(244    module: "ModernBertDecoderAttention",245    query: torch.Tensor,246    key: torch.Tensor,247    value: torch.Tensor,248    attention_mask: Optional[torch.Tensor],249    dropout: float = 0.0,250    scaling: Optional[float] = None,251    sliding_window: Optional[int] = None,252    **kwargs,253) -> tuple[torch.Tensor, Optional[torch.Tensor]]:254    """A simple eager attention implementation for ModernBERT decoder."""255    if scaling is None:256        scaling = module.head_dim**-0.5257 258    # Compute attention scores259    attn_weights = torch.matmul(query, key.transpose(2, 3)) * scaling260 261    # Use the pre-computed attention mask262    causal_mask = attention_mask[:, :, :, : key.shape[-2]]263    attn_weights = attn_weights + causal_mask264 265    # upcast attention to fp32266    attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query.dtype)267    attn_weights = nn.functional.dropout(attn_weights, p=dropout, training=module.training)268    attn_output = torch.matmul(attn_weights, value)269    attn_output = attn_output.transpose(1, 2).contiguous()270    return attn_output, attn_weights271 272 273class ModernBertDecoderAttention(nn.Module):274    """Performs causal multi-headed self attention for ModernBERT decoder.275 276    It supports both local attention (sliding window) and global attention patterns.277    """278 279    def __init__(self, config: ModernBertDecoderConfig, layer_idx: Optional[int] = None):280        super().__init__()281        self.is_sliding = config.layer_types[layer_idx] == "sliding_attention"282        self.config = config283        self.layer_idx = layer_idx284        self.head_dim = config.hidden_size // config.num_attention_heads285        self.num_heads = config.num_attention_heads286        self.all_head_size = self.head_dim * self.num_heads287        self.scaling = self.head_dim**-0.5288        self.attention_dropout = self.config.attention_dropout289        self.is_causal = True290 291        if config.hidden_size % config.num_attention_heads != 0:292            raise ValueError(293                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention heads ({config.num_attention_heads})"294            )295 296        # NOTE: this is different than ModernBERT (separated QKV) so be sure to adapt to this297        self.q_proj = nn.Linear(self.config.hidden_size, self.all_head_size, bias=self.config.attention_bias)298        self.k_proj = nn.Linear(self.config.hidden_size, self.all_head_size, bias=self.config.attention_bias)299        self.v_proj = nn.Linear(self.config.hidden_size, self.all_head_size, bias=self.config.attention_bias)300 301        self.Wo = nn.Linear(config.hidden_size, config.hidden_size, bias=config.attention_bias)302        self.out_drop = nn.Dropout(config.attention_dropout)303 304        self.sliding_window = config.sliding_window if config.layer_types[layer_idx] == "sliding_attention" else None305 306    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")307    def forward(308        self,309        hidden_states: torch.Tensor,310        position_embeddings: torch.Tensor,311        attention_mask: Optional[torch.Tensor],312        past_key_values: Optional[Cache] = None,313        cache_position: Optional[torch.LongTensor] = None,314        **kwargs: Unpack[TransformersKwargs],315    ) -> tuple[torch.Tensor, Optional[torch.Tensor]]:316        input_shape = hidden_states.shape[:-1]317        hidden_shape = (*input_shape, -1, self.head_dim)318 319        query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)320        key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)321        value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)322 323        cos, sin = position_embeddings324        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)325 326        if past_key_values is not None:327            # sin and cos are specific to RoPE models; cache_position needed for the static cache328            cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position}329            key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx, cache_kwargs)330 331        attention_interface: Callable = eager_attention_forward332        if self.config._attn_implementation != "eager":333            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]334 335        attn_output, attn_weights = attention_interface(336            self,337            query_states,338            key_states,339            value_states,340            attention_mask,341            dropout=self.attention_dropout if self.training else 0.0,342            scaling=self.scaling,343            sliding_window=self.sliding_window,344            **kwargs,345        )346 347        attn_output = attn_output.reshape(*input_shape, -1).contiguous()348        attn_output = self.out_drop(self.Wo(attn_output))349        return attn_output, attn_weights350 351 352class ModernBertDecoderLayer(GradientCheckpointingLayer):353    def __init__(self, config: ModernBertDecoderConfig, layer_idx: Optional[int] = None):354        super().__init__()355        self.config = config356        self.layer_idx = layer_idx357        self.attention_type = config.layer_types[layer_idx]358        self.attn_norm = (359            nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)360            if layer_idx != 0361            else nn.Identity()362        )363        self.attn = ModernBertDecoderAttention(config=config, layer_idx=layer_idx)364        self.mlp_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)365        self.mlp = ModernBertDecoderMLP(config)366 367    @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58")368    def forward(369        self,370        hidden_states: torch.Tensor,371        position_embeddings_global: torch.Tensor,372        position_embeddings_local: torch.Tensor,373        attention_mask: Optional[torch.Tensor] = None,374        past_key_values: Optional[Cache] = None,375        use_cache: Optional[bool] = False,376        cache_position: Optional[torch.LongTensor] = None,377        **kwargs,378    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:379        residual = hidden_states380        hidden_states = self.attn_norm(hidden_states)381 382        # apply global RoPE to non-sliding layer only383        if self.attn.is_sliding:384            position_embeddings = position_embeddings_local385        else:386            position_embeddings = position_embeddings_global387 388        # Self Attention389        attn_outputs = self.attn(390            hidden_states=hidden_states,391            position_embeddings=position_embeddings,392            attention_mask=attention_mask,393            past_key_values=past_key_values,394            cache_position=cache_position,395            **kwargs,396        )397        hidden_states = attn_outputs[0]398 399        # Add residual connection400        hidden_states = residual + hidden_states401 402        # MLP403        residual = hidden_states404        hidden_states = self.mlp_norm(hidden_states)405        mlp_output = self.mlp(hidden_states)406        hidden_states = residual + mlp_output407        return hidden_states408 409 410class ModernBertDecoderPredictionHead(ModernBertPredictionHead):411    pass412 413 414@auto_docstring415class ModernBertDecoderPreTrainedModel(ModernBertPreTrainedModel):416    _skip_keys_device_placement = ["past_key_values"]417    _no_split_modules = ["ModernBertDecoderLayer"]418    _supports_flex_attn = True419    _supports_attention_backend = True420    _can_record_outputs = {421        "hidden_states": ModernBertDecoderLayer,422        "attentions": ModernBertDecoderAttention,423    }424 425    def _init_weights(self, module: nn.Module):426        cutoff_factor = self.config.initializer_cutoff_factor427        if cutoff_factor is None:428            cutoff_factor = 3429 430        def init_weight(module: nn.Module, std: float):431            nn.init.trunc_normal_(432                module.weight,433                mean=0.0,434                std=std,435                a=-cutoff_factor * std,436                b=cutoff_factor * std,437            )438 439            if isinstance(module, nn.Linear):440                if module.bias is not None:441                    nn.init.zeros_(module.bias)442 443        stds = {444            "in": self.config.initializer_range,445            "out": self.config.initializer_range / math.sqrt(2.0 * self.config.num_hidden_layers),446            "embedding": self.config.initializer_range,447            "final_out": self.config.hidden_size**-0.5,448        }449 450        if isinstance(module, ModernBertDecoderEmbeddings):451            init_weight(module.tok_embeddings, stds["embedding"])452        elif isinstance(module, ModernBertDecoderMLP):453            init_weight(module.Wi, stds["in"])454            init_weight(module.Wo, stds["out"])455        elif isinstance(module, ModernBertDecoderAttention):456            init_weight(module.q_proj, stds["in"])457            init_weight(module.k_proj, stds["in"])458            init_weight(module.v_proj, stds["in"])459            init_weight(module.Wo, stds["out"])460        elif isinstance(module, ModernBertDecoderPredictionHead):461            init_weight(module.dense, stds["out"])462        elif isinstance(module, ModernBertDecoderForSequenceClassification):463            init_weight(module.classifier, stds["final_out"])464        elif isinstance(module, ModernBertDecoderForCausalLM):465            init_weight(module.decoder, stds["out"])466        elif isinstance(module, nn.LayerNorm):467            module.weight.data.fill_(1.0)468            if module.bias is not None:469                module.bias.data.zero_()470 471    def _check_and_adjust_attn_implementation(self, attn_implementation, is_init_check):472        raise AttributeError("No need to inherit!")473 474    def _maybe_set_compile(self):475        raise AttributeError("No need to inherit!")476 477    def resize_token_embeddings(self, *args, **kwargs):478        raise AttributeError("No need to inherit!")479 480 481@auto_docstring482class ModernBertDecoderModel(ModernBertDecoderPreTrainedModel):483    def __init__(self, config: ModernBertDecoderConfig):484        super().__init__(config)485        self.config = config486        self.embeddings = ModernBertDecoderEmbeddings(config)487        self.layers = nn.ModuleList(488            [ModernBertDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)]489        )490        self.final_norm = nn.LayerNorm(config.hidden_size, eps=config.norm_eps, bias=config.norm_bias)491        self.gradient_checkpointing = False492 493        self.global_rotary_emb = ModernBertDecoderRotaryEmbedding(config=config)494        self.local_rotary_emb = ModernBertDecoderRotaryEmbedding(config=config)495 496        self.post_init()497 498    def get_input_embeddings(self):499        return self.embeddings.tok_embeddings500 501    def set_input_embeddings(self, value):502        self.embeddings.tok_embeddings = value503 504    @check_model_inputs()505    @auto_docstring506    def forward(507        self,508        input_ids: Optional[torch.LongTensor] = None,509        attention_mask: Optional[torch.Tensor] = None,510        position_ids: Optional[torch.LongTensor] = None,511        past_key_values: Optional[Cache] = None,512        inputs_embeds: Optional[torch.Tensor] = None,513        use_cache: Optional[bool] = None,514        cache_position: Optional[torch.LongTensor] = None,515        **kwargs,516    ) -> Union[tuple[torch.Tensor, ...], BaseModelOutputWithPast]:517        if (input_ids is None) == (inputs_embeds is None):518            raise ValueError("You must specify exactly one of input_ids or inputs_embeds")519 520        if input_ids is not None:521            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)522            batch_size, seq_length = input_ids.shape[:2]523        else:524            batch_size, seq_length = inputs_embeds.shape[:2]525 526        # Handle past_key_values and cache setup527        if use_cache and past_key_values is None and not self.training:528            past_key_values = DynamicCache(config=self.config)529 530        if cache_position is None:531            past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0532            cache_position = torch.arange(533                past_seen_tokens,534                past_seen_tokens + seq_length,535                device=input_ids.device if input_ids is not None else inputs_embeds.device,536            )537 538        if position_ids is None:539            position_ids = cache_position.unsqueeze(0).expand(batch_size, -1)540 541        # Calculate embeddings542        hidden_states = self.embeddings(input_ids=input_ids, inputs_embeds=inputs_embeds)543 544        # It may already have been prepared by e.g. `generate`545        if not isinstance(causal_mask_mapping := attention_mask, dict):546            # Prepare mask arguments547            mask_kwargs = {548                "config": self.config,549                "input_embeds": hidden_states,550                "attention_mask": attention_mask,551                "cache_position": cache_position,552                "past_key_values": past_key_values,553                "position_ids": position_ids,554            }555 556            causal_mask_mapping = {557                "full_attention": create_causal_mask(**mask_kwargs),558                "sliding_attention": create_sliding_window_causal_mask(**mask_kwargs),559            }560 561        # create position embeddings to be shared across the decoder layers562        position_embeddings_global = self.global_rotary_emb(hidden_states, position_ids)563        position_embeddings_local = self.local_rotary_emb(hidden_states, position_ids)564 565        for idx, decoder_layer in enumerate(self.layers):566            hidden_states = decoder_layer(567                hidden_states,568                position_embeddings_global=position_embeddings_global,569                position_embeddings_local=position_embeddings_local,570                attention_mask=causal_mask_mapping[decoder_layer.attention_type],571                past_key_values=past_key_values,572                use_cache=use_cache,573                cache_position=cache_position,574                **kwargs,575            )576 577        hidden_states = self.final_norm(hidden_states)578 579        return BaseModelOutputWithPast(580            last_hidden_state=hidden_states,581            past_key_values=past_key_values,582        )583 584 585@auto_docstring(586    custom_intro="""587    The ModernBert Decoder Model with a language modeling head on top for causal language modeling (CLM).588    """589)590class ModernBertDecoderForCausalLM(ModernBertDecoderPreTrainedModel, GenerationMixin):591    _tied_weights_keys = ["decoder.weight"]592 593    def __init__(self, config: ModernBertDecoderConfig):594        super().__init__(config)595        self.config = config596        self.model = ModernBertDecoderModel(config)597        self.lm_head = ModernBertDecoderPredictionHead(config)598        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=config.decoder_bias)599 600        # Initialize weights and apply final processing601        self.post_init()602 603    def get_output_embeddings(self):604        return self.decoder605 606    def set_output_embeddings(self, new_embeddings):607        self.decoder = new_embeddings608 609    @can_return_tuple610    @auto_docstring611    def forward(612        self,613        input_ids: Optional[torch.LongTensor] = None,614        attention_mask: Optional[torch.Tensor] = None,615        position_ids: Optional[torch.LongTensor] = None,616        past_key_values: Optional[Cache] = None,617        inputs_embeds: Optional[torch.Tensor] = None,618        labels: Optional[torch.LongTensor] = None,619        use_cache: Optional[bool] = None,620        **kwargs,621    ) -> Union[tuple, CausalLMOutputWithPast]:622        r"""623        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):624            Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,625            config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored626            (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.627 628        Returns:629            [`~modeling_outputs.CausalLMOutputWithPast`]630            comprising various elements depending on the configuration and inputs.631 632        Example:633 634        ```python635        >>> from transformers import AutoTokenizer, ModernBertDecoderForCausalLM636 637        >>> model = ModernBertDecoderForCausalLM.from_pretrained("blab-jhu/test-32m-dec")638        >>> tokenizer = AutoTokenizer.from_pretrained("blab-jhu/test-32m-dec")639 640        >>> prompt = "The capital of France is"641        >>> inputs = tokenizer(prompt, return_tensors="pt")642 643        >>> # Generate644        >>> generate_ids = model.generate(inputs.input_ids, max_length=1)645        >>> tokenizer.batch_decode(generate_ids, skip_special_tokens=True, clean_up_tokenization_spaces=False)[0]646        "The capital of France is Paris"647        ```648        """649        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)650        outputs = self.model(651            input_ids=input_ids,652            attention_mask=attention_mask,653            position_ids=position_ids,654            past_key_values=past_key_values,655            inputs_embeds=inputs_embeds,656            use_cache=use_cache,657            **kwargs,658        )659 660        hidden_states = outputs[0]661        logits = self.decoder(self.lm_head(hidden_states))662 663        loss = None664        if labels is not None:665            # Shift so that tokens < n predict n666            shift_logits = logits[..., :-1, :].contiguous()667            shift_labels = labels[..., 1:].contiguous()668            # Flatten the tokens669            loss_fct = CrossEntropyLoss()670            shift_logits = shift_logits.view(-1, self.config.vocab_size)671            shift_labels = shift_labels.view(-1)672            # Enable model parallelism673            shift_labels = shift_labels.to(shift_logits.device)674            loss = loss_fct(shift_logits, shift_labels)675 676        return CausalLMOutputWithPast(677            loss=loss,678            logits=logits,679            past_key_values=outputs.past_key_values,680            hidden_states=outputs.hidden_states,681            attentions=outputs.attentions,682        )683 684 685@auto_docstring(686    custom_intro="""687    The ModernBert Decoder Model with a sequence classification head on top (linear layer).688 689    [`ModernBertDecoderForSequenceClassification`] uses the last token in order to do the classification, as other causal models690    (e.g. GPT-1, GPT-2) do.691 692    Since it does classification on the last token, it requires to know the position of the last token. If a693    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If694    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the695    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in696    each row of the batch).697    """698)699class ModernBertDecoderForSequenceClassification(ModernBertDecoderPreTrainedModel):700    def __init__(self, config: ModernBertDecoderConfig):701        super().__init__(config)702        self.num_labels = config.num_labels703        self.model = ModernBertDecoderModel(config)704 705        self.head = ModernBertDecoderPredictionHead(config)706        self.classifier = nn.Linear(config.hidden_size, config.num_labels, bias=config.classifier_bias)707        self.drop = torch.nn.Dropout(config.classifier_dropout)708 709        # Initialize weights and apply final processing710        self.post_init()711 712    @can_return_tuple713    @auto_docstring(checkpoint="blab-jhu/test-32m-dec")714    def forward(715        self,716        input_ids: Optional[torch.LongTensor] = None,717        attention_mask: Optional[torch.Tensor] = None,718        position_ids: Optional[torch.LongTensor] = None,719        past_key_values: Optional[Cache] = None,720        inputs_embeds: Optional[torch.Tensor] = None,721        labels: Optional[torch.LongTensor] = None,722        use_cache: Optional[bool] = None,723        **kwargs,724    ) -> Union[tuple, SequenceClassifierOutputWithPast]:725        r"""726        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):727            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,728            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If729            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).730        """731        transformer_outputs = self.model(732            input_ids,733            attention_mask=attention_mask,734            position_ids=position_ids,735            past_key_values=past_key_values,736            inputs_embeds=inputs_embeds,737            use_cache=use_cache,738            **kwargs,739        )740        hidden_states = transformer_outputs[0]741        hidden_states = self.drop(self.head(hidden_states))742        logits = self.classifier(hidden_states)743 744        if input_ids is not None:745            batch_size, sequence_length = input_ids.shape[:2]746        else:747            batch_size, sequence_length = inputs_embeds.shape[:2]748 749        if self.config.pad_token_id is None and batch_size != 1:750            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")751        if self.config.pad_token_id is None:752            last_non_pad_token = -1753        elif input_ids is not None:754            # To handle both left- and right- padding, we take the rightmost token that is not equal to pad_token_id755            non_pad_mask = (input_ids != self.config.pad_token_id).to(logits.device, torch.int32)756            token_indices = torch.arange(input_ids.shape[-1], device=logits.device, dtype=torch.int32)757            last_non_pad_token = (token_indices * non_pad_mask).argmax(-1)758        else:759            last_non_pad_token = -1760            logger.warning_once(761                f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "762                "unexpected if using padding tokens in conjunction with `inputs_embeds.`"763            )764 765        pooled_logits = logits[torch.arange(batch_size, device=logits.device), last_non_pad_token]766 767        loss = None768        if labels is not None:769            if self.config.problem_type is None:770                if self.num_labels == 1:771                    self.config.problem_type = "regression"772                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):773                    self.config.problem_type = "single_label_classification"774                else:775                    self.config.problem_type = "multi_label_classification"776 777            if self.config.problem_type == "regression":778                loss_fct = MSELoss()779                if self.num_labels == 1:780                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())781                else:782                    loss = loss_fct(pooled_logits, labels)783            elif self.config.problem_type == "single_label_classification":784                loss_fct = CrossEntropyLoss()785                loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))786            elif self.config.problem_type == "multi_label_classification":787                loss_fct = BCEWithLogitsLoss()788                loss = loss_fct(pooled_logits, labels)789 790        return SequenceClassifierOutputWithPast(791            loss=loss,792            logits=pooled_logits,793            past_key_values=transformer_outputs.past_key_values,794            hidden_states=transformer_outputs.hidden_states,795            attentions=transformer_outputs.attentions,796        )797 798 799__all__ = [800    "ModernBertDecoderConfig",801    "ModernBertDecoderModel",802    "ModernBertDecoderPreTrainedModel",803    "ModernBertDecoderForCausalLM",804    "ModernBertDecoderForSequenceClassification",805]806