model.py

import transformers
import torch
from transformers import AutoModel, AutoConfig, AutoTokenizer, MPNetModel

""" NEW model configuration"""
from transformers.configuration_utils import PretrainedConfig
from transformers.utils import logging

import torch
import torch.utils.checkpoint
from torch import nn
from typing import Optional
import torch.nn.functional as F

from transformers.models.roformer.modeling_roformer import (
    RoFormerEmbeddings,
    RoFormerModel,
    RoFormerEncoder,
    RoFormerLayer,
    RoFormerAttention,
    RoFormerIntermediate,
    RoFormerOutput,
    RoFormerSelfAttention,
    RoFormerPreTrainedModel
)
from transformers import RoFormerConfig

from tokenizer import CodeRangeTokenizer

logger = logging.get_logger(__name__)


class NewConfig(PretrainedConfig):
    r"""
    This is the configuration class to store the configuration of a [`NewModel`] or a [`TFNewModel`]. It is used to
    instantiate a NEW model according to the specified arguments, defining the model architecture. Instantiating a
    configuration with the defaults will yield a similar configuration to that of the NEW
    [izhx/new-base-en](https://huggingface.co/izhx/new-base-en) architecture.

    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
    documentation from [`PretrainedConfig`] for more information.


    Args:
        vocab_size (`int`, *optional*, defaults to 30522):
            Vocabulary size of the NEW model. Defines the number of different tokens that can be represented by the
            `inputs_ids` passed when calling [`NewModel`] or [`TFNewModel`].
        hidden_size (`int`, *optional*, defaults to 768):
            Dimensionality of the encoder layers and the pooler layer.
        num_hidden_layers (`int`, *optional*, defaults to 12):
            Number of hidden layers in the Transformer encoder.
        num_attention_heads (`int`, *optional*, defaults to 12):
            Number of attention heads for each attention layer in the Transformer encoder.
        intermediate_size (`int`, *optional*, defaults to 3072):
            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
        hidden_act (`str` or `Callable`, *optional*, defaults to `"gelu"`):
            The non-linear activation function (function or string) in the encoder and pooler. If string, `"gelu"`,
            `"relu"`, `"silu"` and `"gelu_new"` are supported.
        hidden_dropout_prob (`float`, *optional*, defaults to 0.1):
            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
        attention_probs_dropout_prob (`float`, *optional*, defaults to 0.1):
            The dropout ratio for the attention probabilities.
        max_position_embeddings (`int`, *optional*, defaults to 512):
            The maximum sequence length that this model might ever be used with. Typically set this to something large
            just in case (e.g., 512 or 1024 or 2048).
        type_vocab_size (`int`, *optional*, defaults to 2):
            The vocabulary size of the `token_type_ids` passed when calling [`NewModel`] or [`TFNewModel`].
        initializer_range (`float`, *optional*, defaults to 0.02):
            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
        layer_norm_eps (`float`, *optional*, defaults to 1e-12):
            The epsilon used by the layer normalization layers.
        position_embedding_type (`str`, *optional*, defaults to `"rope"`):
            Type of position embedding. Choose one of `"absolute"`, `"rope"`.
        rope_theta (`float`, *optional*, defaults to 10000.0):
            The base period of the RoPE embeddings.
        rope_scaling (`Dict`, *optional*):
            Dictionary containing the scaling configuration for the RoPE embeddings. Currently supports two scaling
            strategies: linear and dynamic. Their scaling factor must be a float greater than 1. The expected format is
            `{"type": strategy name, "factor": scaling factor}`. When using this flag, don't update
            `max_position_embeddings` to the expected new maximum. See the following thread for more information on how
            these scaling strategies behave:
            https://www.reddit.com/r/LocalLLaMA/comments/14mrgpr/dynamically_scaled_rope_further_increases/. This is an
            experimental feature, subject to breaking API changes in future versions.
        classifier_dropout (`float`, *optional*):
            The dropout ratio for the classification head.

    Examples:

    ```python
    >>> from transformers import NewConfig, NewModel

    >>> # Initializing a NEW izhx/new-base-en style configuration
    >>> configuration = NewConfig()

    >>> # Initializing a model (with random weights) from the izhx/new-base-en style configuration
    >>> model = NewModel(configuration)

    >>> # Accessing the model configuration
    >>> configuration = model.config
    ```"""

    model_type = "new"

    def __init__(
        self,
        vocab_size=30528,
        hidden_size=768,
        num_hidden_layers=12,
        num_attention_heads=12,
        intermediate_size=3072,
        hidden_act="gelu",
        hidden_dropout_prob=0.1,
        attention_probs_dropout_prob=0.0,
        max_position_embeddings=2048,
        type_vocab_size=1,
        initializer_range=0.02,
        layer_norm_type='layer_norm',
        layer_norm_eps=1e-12,
        # pad_token_id=0,
        position_embedding_type="rope",
        rope_theta=10000.0,
        rope_scaling=None,
        classifier_dropout=None,
        pack_qkv=True,
        unpad_inputs=False,
        use_memory_efficient_attention=False,
        logn_attention_scale=False,
        logn_attention_clip1=False,
        **kwargs,
    ):
        super().__init__(**kwargs)

        self.vocab_size = vocab_size
        self.hidden_size = hidden_size
        self.num_hidden_layers = num_hidden_layers
        self.num_attention_heads = num_attention_heads
        self.hidden_act = hidden_act
        self.intermediate_size = intermediate_size
        self.hidden_dropout_prob = hidden_dropout_prob
        self.attention_probs_dropout_prob = attention_probs_dropout_prob
        self.max_position_embeddings = max_position_embeddings
        self.type_vocab_size = type_vocab_size
        self.initializer_range = initializer_range
        self.layer_norm_type = layer_norm_type
        self.layer_norm_eps = layer_norm_eps
        self.position_embedding_type = position_embedding_type
        self.rope_theta = rope_theta
        self.rope_scaling = rope_scaling
        self.classifier_dropout = classifier_dropout

        self.pack_qkv = pack_qkv
        self.unpad_inputs = unpad_inputs
        self.use_memory_efficient_attention = use_memory_efficient_attention
        self.logn_attention_scale = logn_attention_scale
        self.logn_attention_clip1 = logn_attention_clip1


"""PyTorch NEW model."""

import math
from typing import List, Optional, Tuple, Union

import torch
import torch.utils.checkpoint
from torch import nn

from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPooling,
    MaskedLMOutput,
    MultipleChoiceModelOutput,
    QuestionAnsweringModelOutput,
    SequenceClassifierOutput,
    TokenClassifierOutput,
)
from transformers.modeling_utils import PreTrainedModel
from transformers.utils import logging

try:
    import xformers.ops as xops
except ImportError as e:
    xops = None

logger = logging.get_logger(__name__)


# Adapted from https://github.com/HazyResearch/flash-attention/blob/main/flash_attn/bert_padding.py
# Which was adapted from https://github.com/mlcommons/training_results_v1.1/blob/main/NVIDIA/benchmarks/bert/implementations/pytorch/padding.py
class IndexFirstAxis(torch.autograd.Function):
    @staticmethod
    def forward(ctx, input, indices):
        ctx.save_for_backward(indices)
        assert input.ndim >= 2
        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]
        second_dim = other_shape.numel()
        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.
        # return input[indices]
        # return torch.gather(
        #     rearrange(input, "b ... -> b (...)"), 0, repeat(indices, "z -> z d", d=second_dim)
        # ).reshape(-1, *other_shape)
        return torch.gather(
            input.view(ctx.first_axis_dim, second_dim),
            0,
            indices.unsqueeze(-1).expand(indices.size(0), second_dim)
        ).reshape(-1, *other_shape)

    @staticmethod
    def backward(ctx, grad_output):
        (indices,) = ctx.saved_tensors
        assert grad_output.ndim >= 2
        other_shape = grad_output.shape[1:]
        # grad_output = rearrange(grad_output, "b ... -> b (...)")
        grad_output = grad_output.view(grad_output.size(0), other_shape.numel())
        grad_input = torch.zeros(
            [ctx.first_axis_dim, grad_output.shape[1]],
            device=grad_output.device,
            dtype=grad_output.dtype,
        )
        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.
        # grad_input[indices] = grad_output
        # grad_input.scatter_(0, repeat(indices, "z -> z d", d=grad_output.shape[1]), grad_output)
        grad_input.scatter_(
            0, indices.unsqueeze(-1).expand(indices.size(0), grad_output.size(1)), grad_output
        )
        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None


index_first_axis = IndexFirstAxis.apply


def unpad_input(hidden_states, attention_mask=None, indices=None):
    """
    Arguments:
        hidden_states: (batch, seqlen, ...)
        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.
        indices: (total_nnz), the indices of non-masked tokens from the flattened input sequence.
    Return:
        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
    """
    if indices is None:
        assert attention_mask is not None
        indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()

    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the
    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim
    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to
    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,
    # so we write custom forward and backward to make it a bit faster.
    hidden_states = hidden_states.view(-1, *hidden_states.shape[2:])
    return index_first_axis(hidden_states, indices)


class IndexPutFirstAxis(torch.autograd.Function):
    @staticmethod
    def forward(
        ctx,
        values: torch.Tensor,
        indices: torch.Tensor,
        first_axis_dim
    ) -> torch.Tensor:
        ctx.save_for_backward(indices)
        assert indices.ndim == 1
        assert values.ndim >= 2
        output = torch.zeros(
            first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype
        )
        output[indices] = values
        return output

    @staticmethod
    def backward(ctx, grad_output: torch.Tensor) -> Tuple[torch.Tensor, None, None]:
        indices, = ctx.saved_tensors
        grad_values = grad_output[indices]
        return grad_values, None, None


index_put_first_axis = IndexPutFirstAxis.apply


def pad_input(inputs: torch.Tensor, indices: torch.Tensor, batch: int, seqlen: int) -> torch.Tensor:
    """Add padding to sequences.

    Arguments:
        inputs: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.
        indices: (total_nnz), `indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()`
        batch: int batch_size
        seqlen: int max sequence length

    Returns:
        inputs: (batch, seqlen, ...)
    """
    output = index_put_first_axis(inputs, indices, batch * seqlen)
    return output.view(batch, seqlen, *inputs.shape[1:])


def rotate_half(x):
    """Rotates half the hidden dims of the input."""
    x1 = x[..., : x.shape[-1] // 2]
    x2 = x[..., x.shape[-1] // 2 :]
    return torch.cat((-x2, x1), dim=-1)


def apply_rotary_pos_emb(q, k, cos, sin):
    """Applies Rotary Position Embedding to the query and key tensors.

    Args:
        q (`torch.Tensor`): The query tensor.
        k (`torch.Tensor`): The key tensor.
        cos (`torch.Tensor`): The cosine part of the rotary embedding.
        sin (`torch.Tensor`): The sine part of the rotary embedding.
    Returns:
        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.
    """
    cos, sin = cos.to(q.dtype), sin.to(q.dtype)
    q_embed = (q * cos) + (rotate_half(q) * sin)
    k_embed = (k * cos) + (rotate_half(k) * sin)
    return q_embed, k_embed


class RotaryEmbedding(torch.nn.Module):
    def __init__(self, dim, max_position_embeddings=512, base=10000.0, device=None):
        super().__init__()

        self.dim = dim
        self.max_position_embeddings = max_position_embeddings
        self.base = base
        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
        self.register_buffer("inv_freq", inv_freq, persistent=False)

        # Build here to make `torch.jit.trace` work.
        self._set_cos_sin_cache(
            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()
        )

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len
        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)

        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)

    def forward(self, x, seq_len=None):
        # x: [bs, num_attention_heads, seq_len, head_size]
        if seq_len > self.max_seq_len_cached:
            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)

        return (
            self.cos_cached[:seq_len, ...].to(dtype=x.dtype),
            self.sin_cached[:seq_len, ...].to(dtype=x.dtype),
        )


class NTKScalingRotaryEmbedding(RotaryEmbedding):
    """RotaryEmbedding extended with fixed and mixed NTK scaling. https://kexue.fm/archives/9706 """

    def __init__(self, dim, max_position_embeddings=512, base=10000, device=None, scaling_factor=1.0, mixed_b=None):
        self.scaling_factor = scaling_factor
        self.mixed_b = mixed_b
        super().__init__(dim, max_position_embeddings, base, device)
        max_position_embeddings = max_position_embeddings * self.scaling_factor
        self._set_cos_sin_cache(max_position_embeddings, self.inv_freq.device, torch.get_default_dtype())

    def _set_cos_sin_cache(self, seq_len, device, dtype):
        self.max_seq_len_cached = seq_len

        if seq_len > self.max_position_embeddings:
            base = self.base * (self.scaling_factor if self.mixed_b is None else 1)
            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))

            if self.mixed_b is None:
                inv_freq = inv_freq / self.scaling_factor ** (2 / self.dim)  # (6)
            else:
                a = torch.tensor(self.scaling_factor).log() / (self.dim / 2) ** self.mixed_b  # (13)
                lambda_1_m = (a * torch.arange(1, self.dim // 2 + 1).float().to(device) ** self.mixed_b).exp()  # (12)
                inv_freq = inv_freq / lambda_1_m  # (10)

            self.register_buffer("inv_freq", inv_freq, persistent=False)

        t = torch.arange(self.max_seq_len_cached, device=device, dtype=torch.float32)

        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
        # Different from paper, but it uses a different permutation in order to obtain the same calculation
        emb = torch.cat((freqs, freqs), dim=-1)
        self.register_buffer("cos_cached", emb.cos().to(dtype), persistent=False)
        self.register_buffer("sin_cached", emb.sin().to(dtype), persistent=False)


class RMSNorm(nn.Module):
    def __init__(self, hidden_size, eps=1e-6):
        """
        RMSNorm is equivalent to T5LayerNorm
        """
        super().__init__()
        self.weight = nn.Parameter(torch.ones(hidden_size))
        self.variance_epsilon = eps

    def forward(self, hidden_states):
        input_dtype = hidden_states.dtype
        hidden_states = hidden_states.to(torch.float32)
        variance = hidden_states.pow(2).mean(-1, keepdim=True)
        hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)
        return self.weight * hidden_states.to(input_dtype)


LAYER_NORM = {
    'layer_norm': nn.LayerNorm,
    'rms_norm': RMSNorm
}


class NewEmbeddings(nn.Module):
    """
    Embedding and Unpadding.
    """

    def __init__(self, config: NewConfig):
        super().__init__()
        self.padding_idx = config.pad_token_id
        self.word_embeddings = nn.Embedding(
            config.vocab_size, config.hidden_size, padding_idx=self.padding_idx
        )

        self.position_embedding_type = config.position_embedding_type
        if self.position_embedding_type == 'absolute':
            self.position_embeddings = nn.Embedding(
                config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
            )
        elif self.position_embedding_type == 'rope':
            self._init_rope(config)
        else:
            raise ValueError

        self.type_vocab_size = config.type_vocab_size
        if self.type_vocab_size > 0:
            self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        # position_ids is contiguous in memory and excluded when serialized
        self.register_buffer(
            "position_ids", torch.arange(config.max_position_embeddings), persistent=False
        )

    def _init_rope(self, config):
        kwargs = dict(
            dim=int(config.hidden_size / config.num_attention_heads),
            max_position_embeddings=config.max_position_embeddings,
            base=config.rope_theta
        )
        if config.rope_scaling is None:
            self.rotary_emb = RotaryEmbedding(**kwargs)
        else:
            kwargs.update(scaling_factor=config.rope_scaling["factor"])
            scaling_type = config.rope_scaling["type"]
            if scaling_type == 'ntk':
                kwargs.update(mixed_b=config.rope_scaling.get('mixed_b', None))
                self.rotary_emb = NTKScalingRotaryEmbedding(**kwargs)
            # elif scaling_type == "linear":
            #     self.rotary_emb = LinearScalingRotaryEmbedding(**kwargs)
            # elif scaling_type == "dynamic":
            #     self.rotary_emb = DynamicNTKScalingRotaryEmbedding(**kwargs)
            else:
                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")

    def forward(
        self,
        unpad_inputs: bool,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        length: Optional[List[int]] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, torch.Tensor, Optional[Tuple], Optional[List[int]]]:
        """
        """
        if inputs_embeds is None:
            device, input_shape = input_ids.device, input_ids.shape
        else:
            device, input_shape = inputs_embeds.device, inputs_embeds.shape[:2]
        batch_size, seq_length = input_shape

        # Set attention_mask if it's None
        if attention_mask is None:
            attention_mask = torch.ones(input_shape, device=device)
            if length is not None:
                for i, l in enumerate(length):
                    attention_mask[i, l:] = 0

        # Set attention_mask_bool for unpadding
        if unpad_inputs:
            attention_mask_bool = attention_mask.bool()
            if length is None:
                length = attention_mask.sum(-1).tolist()

        # Get word embeddings
        if inputs_embeds is None:
            if unpad_inputs:
                input_ids = input_ids[attention_mask_bool].unsqueeze(0)
            inputs_embeds = self.word_embeddings(input_ids)
        else:
            if unpad_inputs:
                inputs_embeds = inputs_embeds[attention_mask_bool].unsqueeze(0)
        embeddings = inputs_embeds

        # Set and unpad position_ids
        if position_ids is None:
            if seq_length > self.position_ids.size(0):
                self.register_buffer(
                    "position_ids", torch.arange(seq_length, device=embeddings.device), persistent=False
                )
            if unpad_inputs:
                # [1, cumsum_seq_len]
                position_ids = torch.cat([self.position_ids[:l] for l in length]).unsqueeze(0)
            else:
                # [bs, seq_len]
                position_ids = self.position_ids[:seq_length].expand(batch_size, -1)
        elif unpad_inputs:
            position_ids = position_ids[attention_mask_bool].unsqueeze(0)  # [1, cumsum_seq_len]

        # Compute rotary embedding
        if self.position_embedding_type == 'rope':
            rope_cos, rope_sin = self.rotary_emb(inputs_embeds, seq_len=seq_length)
            rope_cos = rope_cos[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
            rope_sin = rope_sin[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
            rope_embeds = rope_cos, rope_sin
        else:
            rope_embeds = None

        if self.type_vocab_size > 0:
            if token_type_ids is None:
                token_type_ids = position_ids.mul(0)
            else:
                if self.type_vocab_size < 2:
                    token_type_ids.mul_(0)
                if unpad_inputs:
                    token_type_ids = token_type_ids[attention_mask_bool].unsqueeze(0)

            token_type_embeddings = self.token_type_embeddings(token_type_ids)
            embeddings = embeddings + token_type_embeddings

        # BERT position
        if self.position_embedding_type == "absolute":
            position_embeddings = self.position_embeddings(position_ids)
            embeddings = embeddings + position_embeddings

        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)

        return embeddings, attention_mask, rope_embeds, length


class NewAttention(nn.Module):
    def __init__(self, config: NewConfig, pack_qkv=None, use_memory_efficient_attention=None):
        super().__init__()
        self.config = config
        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
            raise ValueError(
                f"The hidden size ({config.hidden_size}) is not a multiple of the number of attention "
                f"heads ({config.num_attention_heads})"
            )

        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        if pack_qkv is None:
            pack_qkv = config.pack_qkv
        self.pack_qkv = pack_qkv

        if self.pack_qkv:
            self.qkv_proj = nn.Linear(config.hidden_size, self.all_head_size * 3, bias=True)
        else:
            self.q_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
            self.k_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)
            self.v_proj = nn.Linear(config.hidden_size, self.all_head_size, bias=True)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.o_proj = nn.Linear(config.hidden_size, config.hidden_size, bias=True)

        if use_memory_efficient_attention is None:
            use_memory_efficient_attention = self.config.use_memory_efficient_attention
        self.use_memory_efficient_attention = use_memory_efficient_attention
        self.memory_efficient_attention = None if xops is None else xops.memory_efficient_attention
        if self.use_memory_efficient_attention:
            assert self.memory_efficient_attention is not None, 'please install xformers'

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_bias: torch.FloatTensor,
        rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
        attention_scale: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        qkv_inputs: Optional[Tuple] = None,  # For RetroMAE
    ) -> Tuple[torch.Tensor, ...]:
        shape_hd = (self.num_attention_heads, self.attention_head_size)
        # qkv
        if self.pack_qkv and qkv_inputs is None:
            qkv_pack = self.qkv_proj(hidden_states).split(self.all_head_size, dim=-1)
        else:
            if qkv_inputs is None:
                qkv_inputs = (hidden_states, hidden_states, hidden_states)
            qkv_pack = [
                getattr(self, n + '_proj')(s) for s, n in zip(qkv_inputs, 'qkv')
            ]
        query_states, key_states, value_states = [t.view(t.shape[:-1] + shape_hd) for t in qkv_pack]

        if self.config.position_embedding_type == 'rope':
            query_states, key_states = apply_rotary_pos_emb(query_states, key_states, *rope_embeds)

        dtype = query_states.dtype

        if self.config.logn_attention_scale and attention_scale is not None:
            # https://kexue.fm/archives/8823
            query_states = query_states * attention_scale.to(dtype)

        if padding_inputs is not None:
            query_states = pad_input(query_states.squeeze(), *padding_inputs)
            key_states = pad_input(key_states.squeeze(), *padding_inputs)
            value_states = pad_input(value_states.squeeze(), *padding_inputs)

        if self.use_memory_efficient_attention:
            assert self.memory_efficient_attention is not None, "xformers is not loaded"
            assert output_attentions is False, "memory_efficient_attention do not output attentions"
            assert head_mask is None, "Not support yet"
            attention_probs = None
            if torch.is_tensor(attention_bias):
                attention_bias = attention_bias.to(dtype)
            context_layer = self.memory_efficient_attention(
                query_states,
                key_states,
                value_states,
                attn_bias=attention_bias,
                p=self.dropout.p
            )
        else:
            if output_attentions and isinstance(self, NewSdpaAttention):
                raise RuntimeError("SDPA do not output attentions")
            context_layer, attention_probs = self._attention(
                query_states, key_states, value_states, attention_bias, head_mask
            )

        if padding_inputs is not None:
            context_layer = unpad_input(context_layer, indices=padding_inputs[0])

        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(new_context_layer_shape)

        # output proj
        attn_output = self.o_proj(context_layer)

        # add attentions if we output them
        outputs = (attn_output, attention_probs) if output_attentions else (attn_output,)
        return outputs

    def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
        """
        Args:
            q/k/v: (B, L, n_head, head_dim),
        Returns:
            attn_output: (B L, n_head, head_dim)
        """
        query_states = query_states.transpose(1, 2)
        key_states = key_states.transpose(1, 2)
        value_states = value_states.transpose(1, 2)
        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_states, key_states.transpose(-1, -2))

        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_bias is not None:
            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
            attention_scores = attention_scores + attention_bias

        # Normalize the attention scores to probabilities.
        attention_probs = nn.functional.softmax(attention_scores, dim=-1)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        if self.dropout.p > 0:
            attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_states)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        return context_layer, attention_probs


class NewSdpaAttention(NewAttention):
    """
    New attention module using torch.nn.functional.scaled_dot_product_attention. This module inherits from
    `NewAttention` as the weights of the module stays untouched. The only changes are on the forward pass to adapt to
    SDPA API.
    """
    def __init__(self, config: NewConfig, **kwargs):
        super().__init__(config, **kwargs)
        # torch.backends.cuda.enable_mem_efficient_sdp(False)
        # logger.warning(
        #     "Disable memory efficient attention kernel for `NewSdpaAttention`, you can set "
        #     "`use_memory_efficient_attention=True` if it expected to use."
        # )

    def _attention(self, query_states, key_states, value_states, attention_bias, head_mask):
        attn_output = torch.nn.functional.scaled_dot_product_attention(
            query_states.transpose(1, 2),
            key_states.transpose(1, 2),
            value_states.transpose(1, 2),
            attn_mask=attention_bias,
            dropout_p=self.dropout.p if self.training else 0.0,
        )
        attn_output = attn_output.permute(0, 2, 1, 3).contiguous()
        return attn_output, None


NEW_ATTENTION_CLASSES = {
    "eager": NewAttention,
    # "flash_attention_2": ,  # TODO
    "sdpa": NewSdpaAttention,
}


class NewGatedMLP(nn.Module):
    """
    GLU Variants Improve Transformer.
    """

    def __init__(self, config: NewConfig):
        super().__init__()
        self.intermediate_size = config.intermediate_size
        self.up_gate_proj = nn.Linear(config.hidden_size, self.intermediate_size * 2, bias=False)
        self.down_proj = nn.Linear(self.intermediate_size, config.hidden_size, bias=True)
        self.act_fn = ACT2FN[config.hidden_act]
        if config.hidden_dropout_prob > 0:
            self.hidden_dropout = nn.Dropout(config.hidden_dropout_prob)
        else:
            self.hidden_dropout = None

    def forward(self, hidden_states):
        up_gate = self.up_gate_proj(hidden_states)
        up_states, gate = torch.split(up_gate, self.intermediate_size, dim=-1)
        gate = self.act_fn(gate)
        gated_states = gate * up_states
        if self.hidden_dropout is not None:
            gated_states = self.hidden_dropout(gated_states)
        down_states = self.down_proj(gated_states)
        return down_states


class NewLayer(nn.Module):
    def __init__(
        self,
        config: NewConfig,
        pack_qkv=None,
        use_memory_efficient_attention=None,
        attn_implementation=None
    ):
        super().__init__()
        if attn_implementation is None:
            attn_implementation = config._attn_implementation
        if use_memory_efficient_attention is None:
            use_memory_efficient_attention = config.use_memory_efficient_attention
        if use_memory_efficient_attention:
            if attn_implementation != 'eager':
                logger.warning_once(f"Override {attn_implementation=} to 'eager' as {use_memory_efficient_attention=}")
                attn_implementation = 'eager'  # Since it will be SDPA by default for torch>=2.1.1
        self.attention = NEW_ATTENTION_CLASSES[attn_implementation](
            config, pack_qkv=pack_qkv, use_memory_efficient_attention=use_memory_efficient_attention
        )
        self.mlp = NewGatedMLP(config)

        ln_class = LAYER_NORM[config.layer_norm_type]
        self.attn_ln = ln_class(config.hidden_size, eps=config.layer_norm_eps)
        self.mlp_ln = ln_class(config.hidden_size, eps=config.layer_norm_eps)

        if config.hidden_dropout_prob > 0:
            self.hidden_dropout = nn.Dropout(config.hidden_dropout_prob)
        else:
            self.hidden_dropout = None

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_bias: torch.FloatTensor,
        rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
        attention_scale: Optional[torch.FloatTensor] = None,
        subset_indices: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        qkv_inputs: Optional[Tuple] = None,  # For RetroMAE
    ) -> Tuple[torch.Tensor, ...]:
        # Multi head self attention
        residual = hidden_states if qkv_inputs is None else qkv_inputs[0]
        attention_outputs = self.attention(
            hidden_states,
            attention_bias,
            rope_embeds,
            padding_inputs,
            attention_scale,
            head_mask,
            output_attentions=output_attentions,
            qkv_inputs=qkv_inputs,
        )
        hidden_states = attention_outputs[0]
        if self.hidden_dropout is not None:
            hidden_states = self.hidden_dropout(hidden_states)
        hidden_states = residual + hidden_states

        # In pretraining, after the attention of last layer, we only need the masked tokens.
        if subset_indices is not None:
            hidden_states = hidden_states[subset_indices]

        hidden_states = self.attn_ln(hidden_states)

        # Fully Connected
        residual = hidden_states
        hidden_states = self.mlp(hidden_states)
        if self.hidden_dropout is not None:
            hidden_states = self.hidden_dropout(hidden_states)
        hidden_states = residual + hidden_states
        hidden_states = self.mlp_ln(hidden_states)

        # add self attentions if we output attention weights
        outputs = (hidden_states,) + attention_outputs[1:]
        return outputs


class NewEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.config = config
        self.layer = nn.ModuleList([NewLayer(config) for _ in range(config.num_hidden_layers)])
        self.gradient_checkpointing = False

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_bias: Optional[torch.FloatTensor] = None,
        rope_embeds: Optional[Tuple[torch.FloatTensor, torch.FloatTensor]] = None,
        padding_inputs: Optional[Tuple] = None,  # indices, batch, seqlen
        attention_scale: Optional[torch.FloatTensor] = None,
        subset_indices: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = False,
        output_hidden_states: Optional[bool] = False,
        return_dict: Optional[bool] = True,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutput]:
        all_hidden_states = () if output_hidden_states else None
        all_self_attentions = () if output_attentions else None

        for i, layer_module in enumerate(self.layer):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if i >= len(self.layer) - 1:
                layer_subset_indices = subset_indices
            else:
                layer_subset_indices = None

            layer_head_mask = head_mask[i] if head_mask is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    layer_module.__call__,
                    hidden_states,
                    attention_bias,
                    rope_embeds,
                    padding_inputs,
                    attention_scale,
                    layer_subset_indices,
                    layer_head_mask,
                )
            else:
                layer_outputs = layer_module(
                    hidden_states,
                    attention_bias,
                    rope_embeds,
                    padding_inputs,
                    attention_scale,
                    layer_subset_indices,
                    layer_head_mask,
                    output_attentions,
                )

            hidden_states = layer_outputs[0]
            if output_attentions:
                all_self_attentions = all_self_attentions + (layer_outputs[1],)

        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [
                    hidden_states,
                    all_hidden_states,
                    all_self_attentions,
                ]
                if v is not None
            )
        return BaseModelOutput(
            last_hidden_state=hidden_states,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
        )


# Copied from transformers.models.bert.modeling_bert.BertPooler with Bert->New
class NewPooler(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class NewPreTrainedModel(PreTrainedModel):
    """
    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
    models.
    """

    config_class = NewConfig
    base_model_prefix = "new"
    supports_gradient_checkpointing = True
    _supports_sdpa = True

    def _init_weights(self, module):
        """Initialize the weights"""
        if isinstance(module, nn.Linear):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


class NewModel(NewPreTrainedModel):
    """
    The bare New Model transformer outputting raw hidden-states without any specific head on top.
    """

    def __init__(self, config: NewConfig, add_pooling_layer=False):
        super().__init__(config)
        self.config = config

        self.embeddings = NewEmbeddings(config)
        self.encoder = NewEncoder(config)

        self.pooler = NewPooler(config) if add_pooling_layer else None

        # Initialize weights and apply final processing
        self.post_init()

    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        length: Optional[List[int]] = None,
        subset_indices: Optional[torch.LongTensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        unpad_inputs: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], BaseModelOutputWithPooling]:
        r"""
        length  (`list` of length `batch_size`, *optional*):
            If is `None`, return padded `last_hidden_state`.
        subset_indices  ():
            pass
        unpad_inputs  (`bool`, *optional*):
            pass
        """
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        unpad_inputs = unpad_inputs if unpad_inputs is not None else self.config.unpad_inputs
        output_padded = length is None

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            self.warn_if_padding_and_no_attention_mask(input_ids, attention_mask)
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        # TODO: not used
        # # Prepare head mask if needed
        # # 1.0 in head_mask indicate we keep the head
        # # attention_probs has shape bsz x n_heads x N x N
        # # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
        # head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)

        # Get embeddings, may unpad them
        (embedding_output, attention_mask, rope_embeds, length) = self.embeddings(
            unpad_inputs,
            input_ids=input_ids,
            attention_mask=attention_mask,
            length=length,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds
        )

        batch_size, seq_length = input_shape
        if unpad_inputs and self.config.use_memory_efficient_attention:
            attention_bias = xops.fmha.attn_bias.BlockDiagonalMask.from_seqlens(length)
        else:
            # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
            # ourselves in which case we just need to make it broadcastable to all heads.
            attention_bias = self.get_extended_attention_mask(attention_mask, input_shape)
            if self.config.use_memory_efficient_attention:
                # Invalid shape for attention bias: torch.Size([48, 1, 1, 512]) (expected (48, 12, 512, 512))
                attention_bias = attention_bias.expand(-1, self.config.num_attention_heads, seq_length, -1)

        padding_inputs = None
        if unpad_inputs and (output_padded or not self.config.use_memory_efficient_attention):
            indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()
            if not self.config.use_memory_efficient_attention:
                padding_inputs = (indices, *input_shape)

        attention_scale = None
        if self.config.logn_attention_scale:
            logger.warning_once("TODO: logn_attention_scale")
        #     # attention scale log_512(input_len)
        #     attention_scale = attention_mask.sum(1).log() / torch.tensor(self.config.max_position_embeddings).log()
        #     # inference-time logn scale need clip 1
        #     if self.config.logn_attention_clip1:
        #         attention_scale.clip_(1)
        #     attention_scale = attention_scale[:, None, None, None]
        # else:
        #     attention_scale = None

        encoder_outputs = self.encoder(
            embedding_output,
            attention_bias=attention_bias,
            rope_embeds=rope_embeds,
            padding_inputs=padding_inputs,
            attention_scale=attention_scale,
            subset_indices=subset_indices,
            head_mask=head_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
        if unpad_inputs and output_padded:
            sequence_output = pad_input(
                sequence_output.squeeze(), indices, batch_size, seq_length
            )

        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]

        return BaseModelOutputWithPooling(
            last_hidden_state=sequence_output,
            pooler_output=pooled_output,
            hidden_states=encoder_outputs.hidden_states,
            attentions=encoder_outputs.attentions,
        )


class NewLMPredictionHead(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.transform_act_fn = ACT2FN[config.hidden_act]
        self.norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(config.hidden_size, config.vocab_size)

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.norm(hidden_states)
        hidden_states = self.decoder(hidden_states)
        return hidden_states


class NewForMaskedLM(NewPreTrainedModel):
    _tied_weights_keys = ["lm_head.decoder.bias", "lm_head.decoder.weight"]

    def __init__(self, config: NewConfig):
        super().__init__(config)
        self.new = NewModel(config, add_pooling_layer=False)
        self.lm_head = NewLMPredictionHead(config)
        self.loss_fct = nn.CrossEntropyLoss()

        # Initialize weights and apply final processing
        self.post_init()

    def get_output_embeddings(self):
        return self.lm_head.decoder

    def set_output_embeddings(self, new_embeddings):
        self.lm_head.decoder = new_embeddings

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        unpad_inputs: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], MaskedLMOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the masked language modeling loss. Indices should be in `[-100, 0, ...,
            config.vocab_size]` (see `input_ids` docstring) Tokens with indices set to `-100` are ignored (masked), the
            loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`
        """

        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        if labels is None or not self.new.config.unpad_inputs:
            length = None
            subset_indices = None
        else:
            length = attention_mask.sum(-1).tolist()
            labels = labels[attention_mask.bool()].unsqueeze(0)
            subset_indices = labels > -100

        outputs = self.new(
            input_ids,
            attention_mask=attention_mask,
            length=length,
            subset_indices=subset_indices,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            unpad_inputs=unpad_inputs,
        )

        sequence_output = outputs[0]
        prediction_scores = self.lm_head(sequence_output)

        masked_lm_loss = None
        if labels is not None:
            if subset_indices is None:
                mask = attention_mask.bool()
                prediction_scores = prediction_scores[mask]
                labels = labels[mask]
            else:
                labels = labels[subset_indices]
            masked_lm_loss = self.loss_fct(prediction_scores, labels)

        if not return_dict:
            output = (prediction_scores,) + outputs[2:]
            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output

        return MaskedLMOutput(
            loss=masked_lm_loss,
            logits=prediction_scores,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class NewForSequenceClassification(NewPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        self.config = config

        self.new = NewModel(config, add_pooling_layer=True)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        unpad_inputs: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.new(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            unpad_inputs=unpad_inputs,
        )

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        loss = None
        if labels is not None:
            if self.config.problem_type is None:
                if self.num_labels == 1:
                    self.config.problem_type = "regression"
                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
                    self.config.problem_type = "single_label_classification"
                else:
                    self.config.problem_type = "multi_label_classification"

            if self.config.problem_type == "regression":
                loss_fct = nn.MSELoss()
                if self.num_labels == 1:
                    loss = loss_fct(logits.squeeze(), labels.squeeze())
                else:
                    loss = loss_fct(logits, labels)
            elif self.config.problem_type == "single_label_classification":
                loss_fct = nn.CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            elif self.config.problem_type == "multi_label_classification":
                loss_fct = nn.BCEWithLogitsLoss()
                loss = loss_fct(logits, labels)

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return SequenceClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class NewForMultipleChoice(NewPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)

        self.new = NewModel(config, add_pooling_layer=True)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, 1)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        unpad_inputs: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], MultipleChoiceModelOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for computing the multiple choice classification loss. Indices should be in `[0, ...,
            num_choices-1]` where `num_choices` is the size of the second dimension of the input tensors. (See
            `input_ids` above)
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
        num_choices = input_ids.shape[1] if input_ids is not None else inputs_embeds.shape[1]

        input_ids = input_ids.view(-1, input_ids.size(-1)) if input_ids is not None else None
        attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
        token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
        position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
        inputs_embeds = (
            inputs_embeds.view(-1, inputs_embeds.size(-2), inputs_embeds.size(-1))
            if inputs_embeds is not None
            else None
        )

        outputs = self.new(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            unpad_inputs=unpad_inputs,
        )

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
        reshaped_logits = logits.view(-1, num_choices)

        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(reshaped_logits, labels)

        if not return_dict:
            output = (reshaped_logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return MultipleChoiceModelOutput(
            loss=loss,
            logits=reshaped_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class NewForTokenClassification(NewPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.new = NewModel(config, add_pooling_layer=False)
        classifier_dropout = (
            config.classifier_dropout if config.classifier_dropout is not None else config.hidden_dropout_prob
        )
        self.dropout = nn.Dropout(classifier_dropout)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        labels: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        unpad_inputs: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for computing the token classification loss. Indices should be in `[0, ..., config.num_labels - 1]`.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.new(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            unpad_inputs=unpad_inputs,
        )

        sequence_output = outputs[0]

        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)

        loss = None
        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
            loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))

        if not return_dict:
            output = (logits,) + outputs[2:]
            return ((loss,) + output) if loss is not None else output

        return TokenClassifierOutput(
            loss=loss,
            logits=logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


class NewForQuestionAnswering(NewPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

        self.new = NewModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        # Initialize weights and apply final processing
        self.post_init()

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        token_type_ids: Optional[torch.Tensor] = None,
        position_ids: Optional[torch.Tensor] = None,
        head_mask: Optional[torch.Tensor] = None,
        inputs_embeds: Optional[torch.Tensor] = None,
        start_positions: Optional[torch.Tensor] = None,
        end_positions: Optional[torch.Tensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        unpad_inputs: Optional[bool] = None,
    ) -> Union[Tuple[torch.Tensor], QuestionAnsweringModelOutput]:
        r"""
        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
            are not taken into account for computing the loss.
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.new(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
            unpad_inputs=unpad_inputs,
        )

        sequence_output = outputs[0]

        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1).contiguous()
        end_logits = end_logits.squeeze(-1).contiguous()

        total_loss = None
        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, split add a dimension
            if len(start_positions.size()) > 1:
                start_positions = start_positions.squeeze(-1)
            if len(end_positions.size()) > 1:
                end_positions = end_positions.squeeze(-1)
            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            ignored_index = start_logits.size(1)
            start_positions = start_positions.clamp(0, ignored_index)
            end_positions = end_positions.clamp(0, ignored_index)

            loss_fct = nn.CrossEntropyLoss(ignore_index=ignored_index)
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2

        if not return_dict:
            output = (start_logits, end_logits) + outputs[2:]
            return ((total_loss,) + output) if total_loss is not None else output

        return QuestionAnsweringModelOutput(
            loss=total_loss,
            start_logits=start_logits,
            end_logits=end_logits,
            hidden_states=outputs.hidden_states,
            attentions=outputs.attentions,
        )


def forward(
    self,
    hidden_states,
    attention_mask=None,
    sinusoidal_pos=None,
    head_mask=None,
    encoder_hidden_states=None,
    encoder_attention_mask=None,
    past_key_value=None,
    output_attentions=False,
):
    mixed_query_layer = self.query(hidden_states)
    query_layer = self.transpose_for_scores(mixed_query_layer)
    # If this is instantiated as a cross-attention module, the keys
    # and values come from an encoder; the attention mask needs to be
    # such that the encoder's padding tokens are not attended to.
    is_cross_attention = encoder_hidden_states is not None

    if is_cross_attention and past_key_value is not None:
        # reuse k,v, cross_attentions
        key_layer = past_key_value[0]
        value_layer = past_key_value[1]
        attention_mask = encoder_attention_mask
    elif is_cross_attention:
        key_layer = self.transpose_for_scores(self.key(encoder_hidden_states))
        value_layer = self.transpose_for_scores(self.value(encoder_hidden_states))
        attention_mask = encoder_attention_mask
    else:
        key_layer = self.transpose_for_scores(self.key(hidden_states))
        value_layer = self.transpose_for_scores(self.value(hidden_states))
        if sinusoidal_pos is not None:
            if self.rotary_value:
                query_layer, key_layer, value_layer = (
                    self.apply_rotary_position_embeddings(
                        sinusoidal_pos, query_layer, key_layer, value_layer
                    )
                )
            else:
                query_layer, key_layer = self.apply_rotary_position_embeddings(
                    sinusoidal_pos, query_layer, key_layer
                )
        if past_key_value is not None:
            key_layer = torch.cat([past_key_value[0], key_layer], dim=2)
            value_layer = torch.cat([past_key_value[1], value_layer], dim=2)
    if self.is_decoder:
        # if cross_attention save Tuple(torch.Tensor, torch.Tensor) of all cross attention key/value_states.
        # Further calls to cross_attention layer can then reuse all cross-attention
        # key/value_states (first "if" case)
        # if uni-directional self-attention (decoder) save Tuple(torch.Tensor, torch.Tensor) of
        # all previous decoder key/value_states. Further calls to uni-directional self-attention
        # can concat previous decoder key/value_states to current projected key/value_states (third "elif" case)
        # if encoder bi-directional self-attention `past_key_value` is always `None`
        past_key_value = (key_layer, value_layer)

    "Start Old Self Attention Implementation"
    # Take the dot product between "query" and "key" to get the raw attention scores.
    # attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))

    # attention_scores = attention_scores / math.sqrt(self.attention_head_size)
    # if attention_mask is not None:
    #     # Apply the attention mask is (precomputed for all layers in RoFormerModel forward() function)
    #     attention_scores = attention_scores + attention_mask

    # # Normalize the attention scores to probabilities.
    # attention_probs = nn.functional.softmax(attention_scores, dim=-1)

    # # This is actually dropping out entire tokens to attend to, which might
    # # seem a bit unusual, but is taken from the original Transformer paper.
    # attention_probs = self.dropout(attention_probs)

    # # Mask heads if we want to
    # if head_mask is not None:
    #     attention_probs = attention_probs * head_mask

    # context_layer = torch.matmul(attention_probs, value_layer)
    "End Old Self Attention Implementation"

    "Start Flash Attention Implementation"
    if self.training:
        p = self.dropout.p
    else:
        p = 0.0
    context_layer = torch.nn.functional.scaled_dot_product_attention(
        query_layer, key_layer, value_layer, attention_mask, p, is_causal=False
    )
    "End Flash Attention Implementation"

    context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
    new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
    context_layer = context_layer.view(*new_context_layer_shape)
    if output_attentions:
        raise Exception(
            f"output_attentions is not supported in MRoFormerSelfAttention (Flash Attention)"
        )
    outputs = (context_layer,)
    if self.is_decoder:
        outputs = outputs + (past_key_value,)
    return outputs


class JRoFormerEmbeddings(RoFormerEmbeddings):
    """Construct the embeddings from word and token_type embeddings."""

    def __init__(self, config):
        super().__init__(config)
        self.word_embeddings = nn.Embedding(
            config.vocab_size, config.embedding_size, padding_idx=config.pad_token_id
        )
        self.token_type_embeddings = self.word_embeddings


class JRoFormerSelfAttention(RoFormerSelfAttention):
    def __init__(self, config):
        super().__init__(config)
        self.query = nn.Linear(
            config.hidden_size, self.all_head_size, bias=config.use_bias
        )
        self.key = nn.Linear(
            config.hidden_size, self.all_head_size, bias=config.use_bias
        )
        self.value = nn.Linear(
            config.hidden_size, self.all_head_size, bias=config.use_bias
        )


class JRoFormerAttention(RoFormerAttention):
    def __init__(self, config):
        super().__init__(config)
        self.self = JRoFormerSelfAttention(config)


class JRoFormerLayer(RoFormerLayer):
    def __init__(self, config):
        super().__init__(config)
        self.attention = JRoFormerAttention(config)
        self.is_decoder = config.is_decoder
        self.add_cross_attention = config.add_cross_attention
        if self.add_cross_attention:
            if not self.is_decoder:
                raise ValueError(
                    f"{self} should be used as a decoder model if cross attention is added"
                )
            self.crossattention = RoFormerAttention(config)
        self.intermediate = RoFormerIntermediate(config)
        self.output = RoFormerOutput(config)


class JRoFormerEncoder(RoFormerEncoder):
    def __init__(self, config):
        super().__init__(config)
        self.layer = nn.ModuleList(
            [JRoFormerLayer(config) for _ in range(config.num_hidden_layers)]
        )


class JRoFormerModel(RoFormerModel):
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.embeddings = JRoFormerEmbeddings(config)

        if config.embedding_size != config.hidden_size:
            self.embeddings_project = nn.Linear(
                config.embedding_size, config.hidden_size
            )

        self.encoder = JRoFormerEncoder(config)

        # Initialize weights and apply final processing
        self.post_init()

class AsmEncoderConfig(RoFormerConfig):
    model_type = "asm_encoder"

class AsmEncoder(RoFormerPreTrainedModel):
    config_class = AsmEncoderConfig
    base_model_prefix = "asm_encoder"
    supports_gradient_checkpointing = True
    
    def __init__(self, config):
        super().__init__(config)
        self.config = config
        self.jroformer = JRoFormerModel(config)
        self.projection = nn.Linear(config.hidden_size, config.hidden_size)

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.FloatTensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ):
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        outputs = self.jroformer(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )    
        outputs.last_hidden_state = self.projection(outputs.last_hidden_state)
        return outputs       

def patch_model():
    AutoConfig.register("new", NewConfig)
    AutoModel.register(NewConfig, NewModel)
    AutoConfig.register("asm_encoder", AsmEncoderConfig)
    AutoModel.register(AsmEncoderConfig, AsmEncoder)
    AutoTokenizer.register("CodeRangeTokenizer", fast_tokenizer_class=CodeRangeTokenizer)

def patch_roformer():
    transformers.models.roformer.modeling_roformer.RoFormerSelfAttention.forward = forward