modeling_mitre.py

# coding=utf-8

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

import torch
from torch import nn
from .configuration_mitre import MitreConfig
from transformers.utils import logging

from transformers.generation import GenerationMixin
from transformers.modeling_utils import PreTrainedModel
from transformers.activations import ACT2FN
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
    Seq2SeqLMOutput,
    Seq2SeqModelOutput,
)
from transformers.generation.configuration_utils import GenerationConfig
from transformers.generation.beam_search import BeamSearchScorer
from transformers.generation.logits_process import LogitsProcessorList
from transformers.generation.stopping_criteria import StoppingCriteriaList

logger = logging.get_logger(__name__)

def create_position_ids_from_input_ids(input_ids, padding_idx, past_key_values_length=0):
    """
    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
    are ignored. This is modified from fairseq's `utils.make_positions`.
    """
    mask = input_ids.ne(padding_idx).int()
    incremental_indices = (torch.cumsum(mask, dim=1).type_as(mask) + past_key_values_length) * mask
    return incremental_indices.long() + padding_idx


# Modified from transformers.models.m2m_100.modeling_m2m_100.M2M100Attention
# and transformers.models.m2m_100.modeling_m2m_100.M2M100SdpaAttention
class MitreSdpaAttention(nn.Module):

    def __init__(
        self,
        embed_dim: int,
        num_heads: int,
        dropout: float = 0.0,
        bias: bool = True,
        config: Optional[MitreConfig] = None,
    ):
        super().__init__()
        self.embed_dim = embed_dim
        self.num_heads = num_heads
        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        self.config = config

        if (self.head_dim * num_heads) != self.embed_dim:
            raise ValueError(
                f"embed_dim must be divisible by num_heads (got `embed_dim`: {self.embed_dim}"
                f" and `num_heads`: {num_heads})."
            )
        self.scaling = self.head_dim**-0.5

        self.k_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.v_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.q_proj = nn.Linear(embed_dim, embed_dim, bias=bias)
        self.out_proj = nn.Linear(embed_dim, embed_dim, bias=bias)

    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):
        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()

    def forward(
        self,
        hidden_states: torch.Tensor,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.Tensor] = None,
    ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
        """
        1. MitreModel uses MitreSdpaAttention, which is modified from M2M100SdpaAttention.
        Notably, neither of them supports 'output_attentions=True' or 'layer_head_mask is not None',
        meaning that attn_weights are not included in the output.
        Improving this feature is currently a low priority, and we leave this functionality for users to customize.
        2.We plan to enhance this code with Flash Attention v2 in the future.
        """
        bsz, tgt_len, _ = hidden_states.size()

        # get query proj
        query_states = self.q_proj(hidden_states)
        if past_key_value is not None:
            # reuse k, v, self_attention
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)
            key_states = torch.cat([past_key_value[0], key_states], dim=2)
            value_states = torch.cat([past_key_value[1], value_states], dim=2)
        else:
            # self_attention
            key_states = self._shape(self.k_proj(hidden_states), -1, bsz)
            value_states = self._shape(self.v_proj(hidden_states), -1, bsz)

        past_key_value = (key_states, value_states)

        query_states = self._shape(query_states, tgt_len, bsz)

        attn_output = torch.nn.functional.scaled_dot_product_attention(
            query_states,
            key_states,
            value_states,
            attn_mask=attention_mask,
            dropout_p=self.dropout if self.training else 0.0,
            is_causal=False,
        )

        if attn_output.size() != (bsz, self.num_heads, tgt_len, self.head_dim):
            raise ValueError(
                f"`attn_output` should be of size {(bsz, self.num_heads, tgt_len, self.head_dim)}, but is"
                f" {attn_output.size()}"
            )

        attn_output = attn_output.transpose(1, 2)

        # Use the `embed_dim` from the config (stored in the class) rather than `hidden_state` because `attn_output` can be
        # partitioned across GPUs when using tensor-parallelism.
        attn_output = attn_output.reshape(bsz, tgt_len, self.embed_dim)

        attn_output = self.out_proj(attn_output)

        return attn_output, None, past_key_value


# Modified from transformers.models.m2m_100.modeling_m2m100.M2M100DecoderLayer
class MitreDecoderLayer(nn.Module):
    def __init__(self, config: MitreConfig):
        super().__init__()
        self.embed_dim = config.d_model

        self.self_attn = MitreSdpaAttention(
            embed_dim=self.embed_dim,
            num_heads=config.decoder_attention_heads,
            dropout=config.attention_dropout,
            config=config,
        )
        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
        self.activation_dropout = config.activation_dropout

        self.self_attn_layer_norm = nn.LayerNorm(self.embed_dim)
        self.fc1 = nn.Linear(self.embed_dim, config.decoder_ffn_dim)
        self.fc2 = nn.Linear(config.decoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = nn.LayerNorm(self.embed_dim)

    def forward(
        self,
        hidden_states: torch.Tensor,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_value: Optional[Tuple[torch.Tensor]] = None,
        use_cache: Optional[bool] = True,
    ) -> torch.Tensor:
        """
        Args:
            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`
            attention_mask (`torch.FloatTensor`): attention mask of size
                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
            past_key_value (`Tuple(torch.FloatTensor)`): cached past key and value projection states
        """
        residual = hidden_states
        hidden_states = self.self_attn_layer_norm(hidden_states)

        # Self Attention
        # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
        self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
        # add present self-attn cache to positions 1,2 of present_key_value tuple
        hidden_states, _, present_key_value = self.self_attn(
            hidden_states=hidden_states,
            past_key_value=self_attn_past_key_value,
            attention_mask=attention_mask,
        )
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        # Fully Connected
        residual = hidden_states
        hidden_states = self.final_layer_norm(hidden_states)
        hidden_states = self.activation_fn(self.fc1(hidden_states))
        hidden_states = nn.functional.dropout(hidden_states, p=self.activation_dropout, training=self.training)
        hidden_states = self.fc2(hidden_states)
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        hidden_states = residual + hidden_states

        outputs = (hidden_states,)

        if use_cache:
            outputs += (present_key_value,)

        return outputs


class MitrePreTrainedModel(PreTrainedModel):
    config_class = MitreConfig
    base_model_prefix = "model"
    supports_gradient_checkpointing = True
    _no_split_modules = ["MitreDecoderLayer"]
    # we plan to implement codes for falsh attention v2
    _supports_flash_attn_2 = False
    _supports_sdpa = True

    def _init_weights(self, module):
        std = self.config.init_std
        if isinstance(module, nn.Linear):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.Embedding):
            module.weight.data.normal_(mean=0.0, std=std)
            if module.padding_idx is not None:
                module.weight.data[module.padding_idx].zero_()


class MitreDecoder(MitrePreTrainedModel):
    """
    Transformer decoder consisting of *config.decoder_layers* layers. Each layer is a [`MitreDecoderLayer`]

    Args:
        config: MitreConfig
        embed_tokens (nn.Embedding): output embedding
    """

    def __init__(self, config: MitreConfig):
        super().__init__(config)
        self.dropout = config.dropout
        self.padding_idx = config.pad_token_id
        self.max_target_positions = config.max_position_embeddings
        embed_scale = math.sqrt(config.d_model) if config.scale_embedding else 1.0

        self.embed_tokens = MitreScaledWordEmbedding(
            config.vocab_size, config.d_model, self.padding_idx, embed_scale=embed_scale
        )

        self.src_embed_positions = MitreSinusoidalPositionalEmbedding(
            config.max_position_embeddings,
            config.d_model,
            self.padding_idx,
        )
        self.register_embed_positions = MitreSinusoidalPositionalEmbedding(
            config.max_position_embeddings,
            config.d_model,
            self.padding_idx,
        )
        self.tgt_embed_positions = MitreSinusoidalPositionalEmbedding(
            config.max_position_embeddings,
            config.d_model,
            self.padding_idx,
        )
        self.layers = nn.ModuleList([MitreDecoderLayer(config) for _ in range(config.decoder_layers)])
        if config._attn_implementation != "sdpa":
            raise NotImplementedError("Other attention mechanism are not implemented yet.")

        # TODO implement flash atten v2 for MITRE
        # self._use_flash_attention_2 = config._attn_implementation == "flash_attention_2"
        self._use_sdpa = config._attn_implementation == "sdpa"
        self.layer_norm = nn.LayerNorm(config.d_model)

        self.gradient_checkpointing = False
        self._future_mask = torch.empty(0)
        # Initialize weights and apply final processing
        self.post_init()
    
    def create_registers(self, input_ids):
        '''
            create registers by duplicating the language tag respective to each sentence.
            length(registers) = length(real_tokens) = length(tokens) - length(pads)
        '''
        register_nums = (~input_ids.eq(self.padding_idx)).sum(dim=1)
        max_register_nums = register_nums.max().item()
        total_token_nums = input_ids.size(1) + max_register_nums
        batch_size = input_ids.size(0)
        registers = input_ids[range(batch_size), torch.argmax(input_ids, dim=-1)].unsqueeze(1).repeat(1, max_register_nums)
        return registers, register_nums, total_token_nums
    
    def get_token_indices(self, input_ids, total_token_nums, register_nums):
        '''
            return a token_indices for selecting source tokens from expanded_src_tokens
        '''
        token_indices = torch.arange(total_token_nums).expand(input_ids.size(0), -1).to(input_ids.device)
        token_indices = token_indices + register_nums.unsqueeze(1)
        return token_indices
    
    def get_batch_indices(self, input_ids, token_indices):
        '''
            return a batch_indices for selecting source tokens from expanded_src_tokens
        '''
        batch_indices = torch.arange(input_ids.shape[0]).unsqueeze(1).expand(-1, token_indices.size(1)).contiguous()
        return batch_indices

    def combine_src_and_registers(self, input_ids, registers):
        '''
            return a expanded_src_tokens for positional embedding.
        '''
        pads = torch.full_like(registers, self.padding_idx)
        expanded_src_tokens = torch.cat((pads, input_ids, registers), dim=1)
        return expanded_src_tokens
    
    def source_tokens_embedding_with_positions(self, expanded_src_tokens, total_token_nums, batch_indices, indices):
        '''
            return the embeds of source tokens
        '''
        inputs_embeds = self.embed_tokens(expanded_src_tokens)
        inputs_embeds_1 = inputs_embeds[:,:total_token_nums,:] + self.src_embed_positions(expanded_src_tokens[:,:total_token_nums])
        inputs_embeds_2 = inputs_embeds[:,total_token_nums:,:] + self.register_embed_positions(expanded_src_tokens[:,total_token_nums:])
        inputs_embeds = torch.cat((inputs_embeds_1, inputs_embeds_2), dim=1)
        inputs_embeds = inputs_embeds[batch_indices, indices]

        return inputs_embeds
    
    def fill_with_neg_inf(self, t):
        return t.float().fill_(float("-inf")).type_as(t)
    
    def check_contiguous(self, t: torch.Tensor):
        return t if t.is_contiguous() else t.contiguous()

    def build_future_mask(self, embeds, src_length, register_nums, past_key_values_length=0):
        b = register_nums.size(0)
        ns = src_length - register_nums
        if past_key_values_length == 0:
            # in training
            # 1. create mask by cache
            dim = embeds.size(1)
            if (
                self._future_mask.size(0) == 0
                or self._future_mask.size(0) < dim
                ):
                self._future_mask = torch.triu(self.fill_with_neg_inf(torch.zeros([dim, dim])), 1)
            if self._future_mask.device == embeds.device:
                mask = self._future_mask[:dim, :dim].clone()
            else:
                mask = self._future_mask[:dim, :dim].to(embeds, copy=True)

            # 2. bi-directional attention in source tokens and registers
            mask[ :src_length, :src_length] = 0.

            # 3. create batch mask
            batch_mask = mask.unsqueeze(0).expand(b, -1, -1).clone().contiguous()

            # 4. mask source tokens -> registers
            # 5. mask target -> source tokens
            batch_indices = torch.arange(b).to(batch_mask.device).view(-1, 1, 1).expand(b, dim, dim).contiguous()
            row_indices = torch.arange(dim).to(batch_mask.device).view(1, -1, 1).expand(b, dim, dim).contiguous()
            col_indices = torch.arange(dim).to(batch_mask.device).view(1, 1, -1).expand(b, dim, dim).contiguous()
            source_indices = (row_indices < ns.view(-1, 1, 1)) & (col_indices >= ns.view(-1, 1, 1)) & (col_indices < (ns + register_nums).view(-1, 1, 1)).contiguous()
            target_indices = (row_indices >= (ns + register_nums).view(-1, 1, 1)) & (col_indices < ns.view(-1, 1, 1)).contiguous()
            # 4
            batch_mask[batch_indices[source_indices], row_indices[source_indices], col_indices[source_indices]] = float('-inf')
            # 5
            batch_mask[batch_indices[target_indices], row_indices[target_indices], col_indices[target_indices]] = float('-inf')
            # shape: batch_size, head_num (1 for broadcasting), seq_len, seq_len
            batch_mask = batch_mask.unsqueeze(1)

        elif past_key_values_length > 0:
            # in generation
            # this block is only used in fairseq and is not used in huggingface,
            # because we reuse the mask by the cache.
            mask = torch.zeros(past_key_values_length + 1)
            mask = mask.to(embeds, copy=True)
            batch_mask = mask.unsqueeze(0).expand(b, -1).clone().contiguous()
            
            batch_indices = torch.arange(b).view(-1, 1).expand(b, past_key_values_length + 1).to(batch_mask.device)
            token_indices = torch.arange(past_key_values_length + 1).view(1, -1).expand(b, past_key_values_length + 1).to(batch_mask.device)
            target_to_source_mask = token_indices < ns.view(-1, 1)
            
            batch_mask[batch_indices[target_to_source_mask], token_indices[target_to_source_mask]] = float('-inf')
            batch_mask = batch_mask.unsqueeze(1)
        
        batch_mask = batch_mask.view(b, 1, batch_mask.shape[-2], batch_mask.shape[-1])
        return batch_mask

    def forward(
        self,
        input_ids: Optional[torch.Tensor] = None,
        decoder_input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[List[torch.FloatTensor]] = None,
        use_cache: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        registering_cache: dict = None,
    ):
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        # past_key_values_length
        past_key_values_length = past_key_values[0][0].shape[2] if past_key_values is not None else 0

        if past_key_values_length > 0:
            register_nums = registering_cache["register_nums"]
            src_length = registering_cache["src_length"]

        if input_ids is not None and past_key_values_length == 0:
            # ensure contiguous
            input_ids = self.check_contiguous(input_ids)
            decoder_input_ids = self.check_contiguous(decoder_input_ids)
            
            if attention_mask is None:
                # create registers from input_ids
                registers, register_nums, total_token_nums = self.create_registers(input_ids)
                # 'expanded_src_tokens' is combined by input_ids, registers, and pads. 
                expanded_src_tokens = self.combine_src_and_registers(input_ids, registers)
                token_indices = self.get_token_indices(input_ids, total_token_nums, register_nums)
                batch_indices = self.get_batch_indices(input_ids, token_indices)
                # source tokens (input_ids + registers)
                source_tokens = expanded_src_tokens[batch_indices, token_indices]

            else:
                # although we do not give the attention mask in training and the 1st step of generation,
                # we still leave this block here. 
                if registering_cache is None or \
                   not all(key in registering_cache for key in \
                   ("register_nums", "total_token_nums", "expanded_src_tokens",\
                    "batch_indices", "token_indices", "source_tokens")):
                    raise ValueError(
                        "If you generate registers by external codes, \
                        you must provide 'register_nums', 'total_token_nums', \
                        'expanded_src_tokens', 'batch_indices', 'token_indices' \
                        and 'source_tokens' in 'registering_cache' in the training."
                        )
                register_nums, total_token_nums = registering_cache["register_nums"], registering_cache["total_token_nums"]
                expanded_src_tokens = registering_cache["expanded_src_tokens"]
                batch_indices, token_indices = registering_cache["batch_indices"], registering_cache["token_indices"]
                source_tokens = registering_cache["source_tokens"]

            # ensure contiguous
            expanded_src_tokens = self.check_contiguous(expanded_src_tokens)
            source_tokens = self.check_contiguous(source_tokens)

            # get embeds with positions for source tokens (input_ids + registers)
            inputs_embeds = self.source_tokens_embedding_with_positions(expanded_src_tokens, total_token_nums, batch_indices, token_indices)

            # replace the inference trigger with langtok
            # namely, enc-tgt-dec-tgt strategy
            if decoder_input_ids[0][0].item() != source_tokens[0][-1].item():
                decoder_input_ids[:, 0] = source_tokens[:, -1]
            
            tokens = torch.cat([source_tokens, decoder_input_ids], dim=1)
            src_length = source_tokens.shape[1]
        
        decoder_inputs_embeds = self.embed_tokens(decoder_input_ids)
        decoder_inputs_embeds = decoder_inputs_embeds + self.tgt_embed_positions(decoder_input_ids, past_key_values_length, src_length=src_length)

        if past_key_values_length == 0:
            hidden_states = torch.cat([inputs_embeds, decoder_inputs_embeds], dim=1)
        else:
            hidden_states = decoder_inputs_embeds
        
        # ensure contiguous
        hidden_states = self.check_contiguous(hidden_states)

        # if attention_mask is NOT given, we build the attention mask from current hyperparams
        # if attention_mask is given, check the shape of attention mask 
        if attention_mask is None:
            attention_mask = self.build_future_mask(hidden_states, src_length, register_nums, past_key_values_length)
        else:
            bsz, src_len = hidden_states.shape[0], hidden_states.shape[1]
            tgt_len = hidden_states.shape[1] if past_key_values_length == 0 else past_key_values_length + 1
            if attention_mask.size() != (bsz, 1, src_len, tgt_len):
                raise ValueError(
                    f"Attention mask should be of size {(bsz, 1, src_len, tgt_len)}, but is {attention_mask.size()}"
                    )
        
        # ensure contiguous
        attention_mask = self.check_contiguous(attention_mask)

        # this is a param to turncate kv cache
        # in training, it's None, namely, unactivated.
        max_register_num = None
        # masking pads for attention_mask in the training or the 1st step of generation
        if past_key_values_length == 0:
            # if in generation, activate
            max_register_num = register_nums.max().item() if use_cache else None

            padding_mask = tokens.eq(self.padding_idx)
            if padding_mask.any():
                padding_mask = padding_mask.unsqueeze(1).unsqueeze(2)
                attention_mask = attention_mask.masked_fill(padding_mask == 1, float('-inf'))
        
        hidden_states = nn.functional.dropout(hidden_states, p=self.dropout, training=self.training)
        
        if self.gradient_checkpointing and self.training:
            if use_cache:
                logger.warning_once(
                    "`use_cache=True` is incompatible with gradient checkpointing. Setting" " `use_cache=False`..."
                )
                use_cache = False

        # decoder layers
        all_hidden_states = () if output_hidden_states else None
        next_decoder_cache = () if use_cache else None

        for idx, decoder_layer in enumerate(self.layers):
            if output_hidden_states:
                all_hidden_states += (hidden_states,)

            past_key_value = past_key_values[idx] if past_key_values is not None else None

            if self.gradient_checkpointing and self.training:
                layer_outputs = self._gradient_checkpointing_func(
                    decoder_layer.__call__,
                    hidden_states,
                    attention_mask,
                    past_key_value=None,
                    use_cache=use_cache,
                )
            else:
                layer_outputs = decoder_layer(
                    hidden_states,
                    attention_mask=attention_mask,
                    past_key_value=past_key_value,
                    use_cache=use_cache,
                )

                hidden_states = layer_outputs[0]

            if use_cache:
                if past_key_values_length > 0:
                    next_decoder_cache += (layer_outputs[1],)
                else:
                    cache_key, cache_value = layer_outputs[1]
                    clipped_rep = (
                        cache_key[:, :, src_length - max_register_num:, :],
                        cache_value[:, :, src_length - max_register_num:, :]
                    )
                    next_decoder_cache += (clipped_rep,)
        
        if past_key_values_length == 0:
            hidden_states = hidden_states[:,src_length:,:]

        hidden_states = self.layer_norm(hidden_states)

        # add hidden states from the last decoder layer
        if output_hidden_states:
            all_hidden_states += (hidden_states,)
        
        next_cache = next_decoder_cache if use_cache else None
        
        model_output = BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=next_cache,
            hidden_states=all_hidden_states,
        )

        # the registering cache used in generation
        # in the 1st step, we turncate the kv cache to save cost, so we have to change the src_length
        if use_cache:
            model_output.registering_cache = {
                "register_nums": register_nums,
                "src_length": src_length if past_key_values_length > 0 else max_register_num,
                "attention_mask": attention_mask if past_key_values_length > 0 else None
            }
        else:
            model_output.registering_cache = None
        
        return model_output
    

# Copied from transformers.models.m2m_100.modeling_m2m_100.M2M100ScaledWordEmbedding
class MitreScaledWordEmbedding(nn.Embedding):
    """
    This module overrides nn.Embeddings' forward by multiplying with embeddings scale.
    """
    def __init__(self, num_embeddings: int, embedding_dim: int, padding_idx: int, embed_scale: Optional[float] = 1.0):
        super().__init__(num_embeddings, embedding_dim, padding_idx)
        self.embed_scale = embed_scale

    def forward(self, input_ids: torch.Tensor):
        return super().forward(input_ids) * self.embed_scale


class MitreSinusoidalPositionalEmbedding(nn.Module):
    """This module produces sinusoidal positional embeddings of any length."""

    def __init__(self, num_positions: int, embedding_dim: int, padding_idx: Optional[int] = None):
        super().__init__()
        self.offset = 2
        self.embedding_dim = embedding_dim
        self.padding_idx = padding_idx
        self.make_weights(num_positions + self.offset, embedding_dim, padding_idx)

    def make_weights(self, num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
        emb_weights = self.get_embedding(num_embeddings, embedding_dim, padding_idx)
        if hasattr(self, "weights"):
            # in forward put the weights on the correct dtype and device of the param
            emb_weights = emb_weights.to(dtype=self.weights.dtype, device=self.weights.device)

        self.register_buffer("weights", emb_weights, persistent=False)

    @staticmethod
    def get_embedding(num_embeddings: int, embedding_dim: int, padding_idx: Optional[int] = None):
        """
        Build sinusoidal embeddings.

        This matches the implementation in tensor2tensor, but differs slightly from the description in Section 3.5 of
        "Attention Is All You Need".
        """
        half_dim = embedding_dim // 2
        emb = math.log(10000) / (half_dim - 1)
        emb = torch.exp(torch.arange(half_dim, dtype=torch.int64).float() * -emb)
        emb = torch.arange(num_embeddings, dtype=torch.int64).float().unsqueeze(1) * emb.unsqueeze(0)
        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=1).view(num_embeddings, -1)
        if embedding_dim % 2 == 1:
            # zero pad
            emb = torch.cat([emb, torch.zeros(num_embeddings, 1)], dim=1)
        if padding_idx is not None:
            emb[padding_idx, :] = 0

        return emb.to(torch.get_default_dtype())

    @torch.no_grad()
    def forward(
        self, input_ids: torch.Tensor = None, past_key_values_length: int = 0, src_length: int = 0
    ):
        bsz, seq_len = input_ids.size()
        # Create the position ids from the input token ids. Any padded tokens remain padded.
        position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx, past_key_values_length).to(
            input_ids.device
        )

        if past_key_values_length > 0 and src_length > 0:
            position_ids = torch.where(position_ids == 1, position_ids, position_ids - src_length)
        
        # expand embeddings if needed
        max_pos = self.padding_idx + 1 + seq_len + past_key_values_length
        
        if max_pos > self.weights.size(0):
            self.make_weights(max_pos + self.offset, self.embedding_dim, self.padding_idx)

        return self.weights.index_select(0, position_ids.view(-1)).view(bsz, seq_len, self.weights.shape[-1]).detach()

class MitreModel(MitrePreTrainedModel):
    _tied_weights_keys = ["decoder.embed_tokens.weight"]

    def __init__(self, config: MitreConfig):
        super().__init__(config)

        self.decoder = MitreDecoder(config)

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

    def get_input_embeddings(self):
        return self.decoder.embed_tokens

    def get_decoder(self):
        return self.decoder

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        decoder_input_ids: Optional[torch.Tensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        registering_cache: dict = None,
    ) -> Union[Tuple[torch.Tensor], Seq2SeqModelOutput]:
        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
        )
        use_cache = use_cache if use_cache is not None else self.config.use_cache

        decoder_outputs = self.decoder(
            input_ids=input_ids,
            decoder_input_ids=decoder_input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
            registering_cache=registering_cache
        )

        model_output = Seq2SeqModelOutput(
            last_hidden_state=decoder_outputs.last_hidden_state,
            past_key_values=decoder_outputs.past_key_values,
            decoder_hidden_states=decoder_outputs.hidden_states,
            decoder_attentions=decoder_outputs.attentions,
        )
        model_output.registering_cache = decoder_outputs.registering_cache
        return model_output

class MitreForConditionalGeneration(MitrePreTrainedModel, GenerationMixin):
    base_model_prefix = "model"
    _tied_weights_keys = ["decoder.embed_tokens.weight", "lm_head.weight"]

    def __init__(self, config: MitreConfig):
        super().__init__(config)
        self.model = MitreModel(config)
        self.lm_head = nn.Linear(config.d_model, self.model.decoder.embed_tokens.num_embeddings, bias=False)

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

    def get_decoder(self):
        return self.model.get_decoder()

    def get_output_embeddings(self):
        return self.lm_head

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

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        decoder_input_ids: Optional[torch.LongTensor] = None,
        attention_mask: Optional[torch.Tensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.FloatTensor]]] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        registering_cache: dict = None,
    ) -> Union[Tuple[torch.Tensor], Seq2SeqLMOutput]:
        
        outputs = self.model(
            input_ids=input_ids,
            decoder_input_ids=decoder_input_ids,
            attention_mask=attention_mask,
            past_key_values=past_key_values,
            use_cache=use_cache,
            output_hidden_states=output_hidden_states,
            registering_cache=registering_cache,
        )

        lm_logits = self.lm_head(outputs[0])
        
        if labels is not None:
            raise NotImplementedError("Please implement your loss function here.")

        model_output = Seq2SeqLMOutput(
            loss=None,
            logits=lm_logits,
            past_key_values=outputs.past_key_values,
            decoder_hidden_states=outputs.decoder_hidden_states,
            decoder_attentions=outputs.decoder_attentions,
        )
        model_output.registering_cache = outputs.registering_cache
        return model_output

    @staticmethod
    def _reorder_cache(past_key_values, beam_idx):
        reordered_past = ()
        for layer_past in past_key_values:
            reordered_past += (
                tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),
            )
        return reordered_past
    
    @staticmethod
    def _reorder_register_cache(t, beam_idx):
        """ a costumized reorder method """
        return t.index_select(dim=0, index=beam_idx.to(t.device))
    
    @staticmethod
    def _expand_inputs_for_generation(
        input_ids: Optional[torch.LongTensor] = None,
        beam_size: int = 1,
    ) -> torch.LongTensor:
        """
          Expands input_ids from [batch_size, len(tokens)] to [batch_size * expand_size, , len(tokens)]
          This is simplified from 'transformers.generation.utils.GenerationMixin._expand_inputs_for_generation'
        """
        if beam_size == 1:
            return input_ids

        return input_ids.repeat_interleave(beam_size, dim=0)
    
    def generate(self,
        input_ids: Optional[torch.Tensor] = None,
        generation_config: Optional[GenerationConfig] = None,
        **kwargs: Dict
        ):
        """
            Inference with beam search.
            This code is an improved version of transformers.generation.utils.GenerationMixin.generate.
            There are two main improvements:
            1. 'soft early_stop' in beam search.
                a) problem in the vanilla version.
                In multilingual translation models such as NLLB and M2M, the vanilla early stop in BeamSearchScorer
                (the official implementation by HuggingFace) marks ended sequences with pad(1). However, these ended
                sequences are still fed into the model, leading to significant memory waste. 
                b) our improvement.
                We implemented a "soft early stop" to address this issue. Instead of modifying BeamSearchScorer
                (to maintain code flexibility), we remove ended sequences from the input. Since this changes the
                shape of the output hidden states, we insert placeholders to maintain compatibility with
                BeamSearchScorer's state shapes.
                Based on our tests, this improvement reduces memory usage by half.
            2. mask reusing.
                a) problem:
                Registers require attention masks at each step.
                A sequence may consist of four parts: padding, source tokens, registers, and target tokens.
                During training, we mask all tokens before registers for target token generation. During generation,
                we cannot allow target tokens to "see" padding tokens, requiring masks at every step.
                This leads to computational inefficiency.
                b) our improvement.
                First, we turncate the source tokens and their representations to reduce cost.
                Second, for source tokens acting as placeholders, we modified the mask generation logic compared to
                our Fairseq implementation.
                Third, to avoid regenerating masks at each step, we cache the mask in 'registering_cache', where cached
                mask is managed like the key-value cache in beam search. Then, At every step, we add a column of zeros
                to maintain alignment.
        """
        if generation_config != None:
            assert type(generation_config) is GenerationConfig
            self.generation_config = generation_config
        self.generation_config.update(**kwargs)

        generation_config = self.generation_config

        batch_size = input_ids.shape[0]
        beam_size = generation_config.num_beams
        device = input_ids.device
        max_cache_length = generation_config.max_length
        eos_token_id = torch.Tensor([generation_config.eos_token_id])

        # initial the target tokens
        decoder_input_ids = torch.full(
            (batch_size, 1), 
            self.generation_config.decoder_start_token_id, 
            dtype=input_ids.dtype, 
            device=device
        )
        
        beam_scorer = BeamSearchScorer(
            batch_size=batch_size,
            num_beams=beam_size,
            device=device,
            length_penalty=self.generation_config.length_penalty,
            do_early_stopping=self.generation_config.early_stopping,
            num_beam_hyps_to_keep=self.generation_config.num_return_sequences,
            max_length=max_cache_length,
        )

        input_ids = self._expand_inputs_for_generation(input_ids, beam_size)
        decoder_input_ids = self._expand_inputs_for_generation(decoder_input_ids, beam_size)
        cur_len = decoder_input_ids.shape[1]

        this_peer_finished = False
        past_key_values = None
        registering_cache= None
        attention_mask = None
        # done_mask shows the ended sequences.
        # (~done_mask) shows the running sequences.
        done_mask = None
        
        # we follow the style of M2M and NLLB
        # so we simplify the initialization of thoes two processors.
        logits_processor = LogitsProcessorList()
        stopping_criteria = StoppingCriteriaList()

        beam_scores = torch.zeros((batch_size, beam_size), dtype=torch.float, device=input_ids.device)
        beam_scores[:, 1:] = -1e9
        beam_scores = beam_scores.view((batch_size * beam_size,))      
        while not this_peer_finished:

            if past_key_values is not None:
                decoder_input_ids_for_generation = decoder_input_ids[:, -1:]
                attention_mask = registering_cache["attention_mask"]
                # Get the mask when the first time using kv cache.
                # After it, we can simply repeat 0. (the last column of mask) to get the next mask.
                # As a result, we avoid generate the mask from scratch in kv cache and save memory.
                if attention_mask is not None:
                    attention_mask = torch.cat((attention_mask, attention_mask[..., -1:]), dim=-1)
            else:
                decoder_input_ids_for_generation = decoder_input_ids

            outputs = self(
                input_ids, 
                decoder_input_ids_for_generation, 
                attention_mask=attention_mask, 
                past_key_values=past_key_values, 
                use_cache=True, 
                registering_cache=registering_cache
            )
            del input_ids
            input_ids = None

            past_key_values = outputs.past_key_values
            registering_cache = outputs.registering_cache
            next_token_logits = outputs.logits[:, -1, :].clone().float()
            del outputs

            next_token_logits = next_token_logits.to(device)
            next_token_scores = nn.functional.log_softmax(
                next_token_logits, dim=-1
            )  # (batch_size * num_beams, vocab_size)

            next_token_scores_processed = logits_processor(decoder_input_ids, next_token_scores)

            # if any sequence is ended, we have to keep the shape of Scorer's states.
            # Details are described in the head of this function.
            if done_mask is not None:
                if done_mask.any():
                    # the placeholder of scores is '0.'
                    restored_tensor = torch.zeros(
                        (batch_size * beam_size, next_token_scores_processed.shape[1]),
                        dtype=next_token_scores_processed.dtype,
                        device=next_token_scores_processed.device
                    )
                    restored_tensor[~done_mask] = next_token_scores_processed
                    next_token_scores_processed = restored_tensor
                    # the placeholder of tokens is 'pad_token_id'
                    restored_tokens = torch.full(
                        (batch_size * beam_size, decoder_input_ids.shape[1]), 
                        self.generation_config.pad_token_id, 
                        dtype=decoder_input_ids.dtype, 
                        device=device
                        )
                    restored_tokens[~done_mask] = decoder_input_ids
                    decoder_input_ids = restored_tokens
            
            next_token_scores = next_token_scores_processed + beam_scores[:, None].expand_as(
                next_token_scores_processed
            )

            # reshape for beam search
            vocab_size = next_token_scores.shape[-1]
            next_token_scores = next_token_scores.view(batch_size, beam_size * vocab_size)

            # Beam token selection: pick 1 + eos_token_id.shape[0] next tokens for each beam so we have at least 1
            # non eos token per beam.
            n_eos_tokens = eos_token_id.shape[0] if eos_token_id is not None else 0
            n_tokens_to_keep = max(2, 1 + n_eos_tokens) * beam_size
            next_token_scores, next_tokens = torch.topk(
                    next_token_scores, n_tokens_to_keep, dim=1, largest=True, sorted=True
                )
            
            next_indices = torch.div(next_tokens, vocab_size, rounding_mode="floor")
            next_tokens = next_tokens % vocab_size

            beam_outputs = beam_scorer.process(
                decoder_input_ids,
                next_token_scores,
                next_tokens,
                next_indices,
                pad_token_id=generation_config.pad_token_id,
                eos_token_id=generation_config.eos_token_id,
                decoder_prompt_len=1,
            )
            beam_scores = beam_outputs["next_beam_scores"]
            beam_next_tokens = beam_outputs["next_beam_tokens"]
            beam_idx = beam_outputs["next_beam_indices"]

            # 'last_done_mask' is used for reordering cache
            # details are described in the next code block
            if done_mask is not None:
                last_done_mask = done_mask
            
            # get the newest status of sequences.
            # then, filter the beam_idx
            done_mask = beam_scorer._done.clone().view(-1)
            done_mask = self._expand_inputs_for_generation(done_mask, beam_size)
            beam_idx = beam_idx[~done_mask]

            decoder_input_ids = torch.cat([decoder_input_ids[beam_idx, :], beam_next_tokens[~done_mask].unsqueeze(-1)], dim=-1)
            
            # different from processing tokens, caches' order is decided by 'tokens', 'done_mask' and
            # 'beam_idx', simultaneously.
            if decoder_input_ids_for_generation.shape[0] < beam_next_tokens.shape[0]:
                # Take carefule! If the running sequences' num is small than the num of input sequences,
                # it means the Scorer decides to end it, but the cache still follows the last status.
                # Therefore, we should employ the last done mask rather than newest done mask.
                if (~done_mask).sum() < decoder_input_ids_for_generation.shape[0]:
                    count_mask = last_done_mask
                else:
                    count_mask = done_mask
                # For biasing the beam_idx
                # Example:
                # done_mask with beam size of 2: [f, f, t, t, f, f]
                # beam_idx: [0, 0, 2, 2, 4, 5]
                # reorder_idx: [0-0, 0-0, 4-2, 5-2]
                prefix_sum = torch.cat([
                    torch.zeros_like(count_mask[:1], dtype=torch.long),
                    torch.cumsum(count_mask.long(), dim=0)
                    ], dim=0)
                reorder_idx = beam_idx - prefix_sum[beam_idx]
                not_done = ~done_mask[beam_idx]
                reorder_idx = reorder_idx[not_done]     
            else:
                reorder_idx = beam_idx

            past_key_values = self._reorder_cache(past_key_values, reorder_idx)
            registering_cache["register_nums"] = self._reorder_register_cache(registering_cache["register_nums"], reorder_idx)
            if registering_cache["attention_mask"] is not None:
                registering_cache["attention_mask"] = self._reorder_register_cache(registering_cache["attention_mask"], reorder_idx)
            
            cur_len = cur_len + 1

            if beam_scorer.is_done:
                this_peer_finished = True

        sequence_outputs = beam_scorer.finalize(
            decoder_input_ids,
            beam_scores,
            next_tokens,
            next_indices,
            pad_token_id=generation_config.pad_token_id,
            eos_token_id=eos_token_id,
            max_length=stopping_criteria.max_length,
            decoder_prompt_len=1,
        )

        return sequence_outputs["sequences"]


MitreForConditionalGeneration.register_for_auto_class("AutoModel")