base_models.py

from typing import Optional, Tuple, Union, List

import copy
import pickle
import numbers
from functools import partial

from transformers import (
    AutoConfig, GPT2Model, GPT2LMHeadModel
)
from transformers.modeling_outputs import (
    BaseModelOutputWithPastAndCrossAttentions,
)

from transformers.pytorch_utils import Conv1D
from transformers.models.gpt2.modeling_gpt2 import GPT2Attention, GPT2MLP
from transformers.modeling_outputs import CausalLMOutputWithCrossAttentions

import torch
from torch import nn
from torch.nn.parameter import Parameter
from torch.nn import init
from torch import Tensor, Size
import torch.nn.functional as F
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss

from utils import get_position_embeds, sinusoidal_positional_embedding



_shape_t = Union[int, List[int], Size]


def get_correct_hidden_state_attn(att_uses_str, hidden_states, hidden_states_tok, hidden_states_pos):
    if att_uses_str == "mixed":
        return hidden_states
    elif att_uses_str == "tok_emb":
        return hidden_states_tok
    elif att_uses_str == "pos_emb":
        return hidden_states_pos

class CustomLayerNorm(nn.Module):
    """ 
        This is same as nn.LayerNorm. Why we need this? Because GPT2PreTrainedModel class's _init_weights method has:
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)

        And this will break when elementwise_affine is False. So we need to create a custom class for which this elif block will not break.
    """
    __constants__ = ['normalized_shape', 'eps', 'elementwise_affine']
    normalized_shape: Tuple[int, ...]
    eps: float
    elementwise_affine: bool

    def __init__(self, normalized_shape: _shape_t, eps: float = 1e-5, elementwise_affine: bool = True,
                 device=None, dtype=None) -> None:
        factory_kwargs = {'device': device, 'dtype': dtype}
        super(CustomLayerNorm, self).__init__()
        if isinstance(normalized_shape, numbers.Integral):
            # mypy error: incompatible types in assignment
            normalized_shape = (normalized_shape,)  # type: ignore[assignment]
        self.normalized_shape = tuple(normalized_shape)  # type: ignore[arg-type]
        self.eps = eps
        self.elementwise_affine = elementwise_affine
        if self.elementwise_affine:
            self.weight = Parameter(torch.empty(self.normalized_shape, **factory_kwargs))
            self.bias = Parameter(torch.empty(self.normalized_shape, **factory_kwargs))
        else:
            self.register_parameter('weight', None)
            self.register_parameter('bias', None)

        self.reset_parameters()

    def reset_parameters(self) -> None:
        if self.elementwise_affine:
            init.ones_(self.weight)
            init.zeros_(self.bias)

    def forward(self, input: Tensor) -> Tensor:
        return F.layer_norm(
            input, self.normalized_shape, self.weight, self.bias, self.eps)

    def extra_repr(self) -> str:
        return '{normalized_shape}, eps={eps}, ' \
            'elementwise_affine={elementwise_affine}'.format(**self.__dict__)

class GPT2AttentionCustom(GPT2Attention):
    def __init__(self, config, custom_args, is_cross_attention=False, layer_idx=None):
        super().__init__(config, is_cross_attention=False, layer_idx=None)
        self.custom_args = custom_args
        max_positions = config.max_position_embeddings
        self.register_buffer(
            "bias",
            torch.tril(torch.ones((max_positions, max_positions), dtype=torch.bool)).view(
                1, 1, max_positions, max_positions
            ),
            persistent=False,
        )
        self.register_buffer("masked_bias", torch.tensor(-1e4), persistent=False)

        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = self.embed_dim // self.num_heads
        self.split_size = self.embed_dim
        if self.head_dim * self.num_heads != self.embed_dim:
            raise ValueError(
                f"`embed_dim` must be divisible by num_heads (got `embed_dim`: {self.embed_dim} and `num_heads`:"
                f" {self.num_heads})."
            )

        self.scale_attn_weights = config.scale_attn_weights
        self.is_cross_attention = is_cross_attention

        # Layer-wise attention scaling, reordering, and upcasting
        self.scale_attn_by_inverse_layer_idx = config.scale_attn_by_inverse_layer_idx
        self.layer_idx = layer_idx
        self.reorder_and_upcast_attn = config.reorder_and_upcast_attn

        if self.is_cross_attention:
            self.c_attn = Conv1D(2 * self.embed_dim, self.embed_dim)
            self.q_attn = Conv1D(self.embed_dim, self.embed_dim)
        else:
            att_uses_list = list(set(list(self.custom_args["att_uses"].values())))
            if self.layer_idx != 0 or (len(att_uses_list) == 1 and att_uses_list[0] == "mixed"):
                self.kqv_mixed = True
                self.c_attn = Conv1D(3 * self.embed_dim, self.embed_dim)
            else:
                self.kqv_mixed = False
                self.c_attn_q = Conv1D(self.embed_dim, self.embed_dim)
                self.c_attn_k = Conv1D(self.embed_dim, self.embed_dim)
                self.c_attn_v = Conv1D(self.embed_dim, self.embed_dim)
        self.c_proj = Conv1D(self.embed_dim, self.embed_dim)

        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)

        self.pruned_heads = set()

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        hidden_states_tok: Optional[Tuple[torch.FloatTensor]] = None,
        hidden_states_pos: Optional[Tuple[torch.FloatTensor]] = None,
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Tuple[Union[torch.Tensor, Tuple[torch.Tensor]], ...]:
        if encoder_hidden_states is not None:
            if not hasattr(self, "q_attn"):
                raise ValueError(
                    "If class is used as cross attention, the weights `q_attn` have to be defined. "
                    "Please make sure to instantiate class with `GPT2Attention(..., is_cross_attention=True)`."
                )

            query = self.q_attn(hidden_states)
            key, value = self.c_attn(encoder_hidden_states).split(self.split_size, dim=2)
            attention_mask = encoder_attention_mask
        else:
            if self.layer_idx > 0 or self.kqv_mixed:
                query, key, value = self.c_attn(hidden_states).split(self.split_size, dim=2)
            else:
                k_uses = get_correct_hidden_state_attn(self.custom_args["att_uses"]["k"], hidden_states, hidden_states_tok, hidden_states_pos)
                q_uses = get_correct_hidden_state_attn(self.custom_args["att_uses"]["q"], hidden_states, hidden_states_tok, hidden_states_pos)
                v_uses = get_correct_hidden_state_attn(self.custom_args["att_uses"]["v"], hidden_states, hidden_states_tok, hidden_states_pos)
                query = self.c_attn_q(q_uses)
                key = self.c_attn_k(k_uses)
                value = self.c_attn_v(v_uses)

        query = self._split_heads(query, self.num_heads, self.head_dim)
        key = self._split_heads(key, self.num_heads, self.head_dim)
        value = self._split_heads(value, self.num_heads, self.head_dim)

        if layer_past is not None:
            past_key, past_value = layer_past
            key = torch.cat((past_key, key), dim=-2)
            value = torch.cat((past_value, value), dim=-2)

        if use_cache is True:
            present = (key, value)
        else:
            present = None

        if self.reorder_and_upcast_attn:
            attn_output, attn_weights = self._upcast_and_reordered_attn(query, key, value, attention_mask, head_mask)
        else:
            attn_output, attn_weights = self._attn(query, key, value, attention_mask, head_mask)

        attn_output = self._merge_heads(attn_output, self.num_heads, self.head_dim)
        attn_output = self.c_proj(attn_output)
        attn_output = self.resid_dropout(attn_output)

        outputs = (attn_output, present)
        if output_attentions:
            outputs += (attn_weights,)

        return outputs  # a, present, (attentions)

class GPT2BlockCustom(nn.Module):
    def __init__(self, config, custom_args, layer_idx=None):
        super().__init__()
        self.layer_idx = layer_idx

        if custom_args["diff_embs"]:
            config.hidden_size = custom_args["model_hid_dim"]
        hidden_size = config.hidden_size
        inner_dim = config.n_inner if config.n_inner is not None else 4 * hidden_size
        self.custom_args = custom_args
        customLayerNormClass = nn.LayerNorm

        if custom_args["layer_norm_type"] == "no_gamma_beta":
            customLayerNormClass = partial(CustomLayerNorm, elementwise_affine=False)

        if "attn" in self.custom_args["layer_norm"]:
            self.ln_1 = customLayerNormClass(hidden_size, eps=config.layer_norm_epsilon)
        self.attn = GPT2AttentionCustom(config, custom_args, layer_idx=layer_idx)

        if config.add_cross_attention:
            self.crossattention = GPT2AttentionCustom(config, custom_args, is_cross_attention=True, layer_idx=layer_idx)
            if "attn" in self.custom_args["layer_norm"]:
                self.ln_cross_attn = customLayerNormClass(hidden_size, eps=config.layer_norm_epsilon)

        if self.custom_args["is_mlp"]:
            if "mlp" in self.custom_args["layer_norm"]:
                self.ln_2 = customLayerNormClass(hidden_size, eps=config.layer_norm_epsilon)
            self.mlp = GPT2MLP(inner_dim, config)

    def forward(
        self,
        hidden_states: Optional[Tuple[torch.FloatTensor]],
        hidden_states_tok: Optional[Tuple[torch.FloatTensor]] = None,
        hidden_states_pos: Optional[Tuple[torch.FloatTensor]] = None,
        layer_past: Optional[Tuple[torch.Tensor]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = False,
        output_attentions: Optional[bool] = False,
    ) -> Union[Tuple[torch.Tensor], Optional[Tuple[torch.Tensor, Tuple[torch.FloatTensor, ...]]]]:

        if self.layer_idx != 0:
            # hidden_states_tok and hidden_states_pos are only defined for the first layer
            # Make sure no other layer is trying to access them
            hidden_states_pos = None
            hidden_states_tok = None

        residual = hidden_states
        if "attn" in self.custom_args["layer_norm"]:
            hidden_states = self.ln_1(hidden_states)
            if hidden_states_tok is not None:
                hidden_states_tok = self.ln_1(hidden_states_tok)
            if hidden_states_pos is not None:
                hidden_states_pos = self.ln_1(hidden_states_pos)
        attn_outputs = self.attn(
            hidden_states,
            hidden_states_tok=hidden_states_tok,
            hidden_states_pos=hidden_states_pos,
            layer_past=layer_past,
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
        )
        attn_output = attn_outputs[0]  # output_attn: a, present, (attentions)
        outputs = attn_outputs[1:]
        # residual connection
        if self.custom_args["is_resid_conn"]:
            hidden_states = attn_output + residual
        else:
            hidden_states = attn_output

        if encoder_hidden_states is not None:
            # add one self-attention block for cross-attention
            if not hasattr(self, "crossattention"):
                raise ValueError(
                    f"If `encoder_hidden_states` are passed, {self} has to be instantiated with "
                    "cross-attention layers by setting `config.add_cross_attention=True`"
                )
            residual = hidden_states
            if "attn" in self.custom_args["layer_norm"]:
                hidden_states = self.ln_cross_attn(hidden_states)
            cross_attn_outputs = self.crossattention(
                hidden_states,
                attention_mask=attention_mask,
                head_mask=head_mask,
                encoder_hidden_states=encoder_hidden_states,
                encoder_attention_mask=encoder_attention_mask,
                output_attentions=output_attentions,
            )
            attn_output = cross_attn_outputs[0]
            # residual connection
            if self.custom_args["is_resid_conn"]:
                hidden_states = residual + attn_output
            else:
                hidden_states = attn_output

            outputs = outputs + cross_attn_outputs[2:]  # add cross attentions if we output attention weights

        residual = hidden_states
        
        if self.custom_args["is_mlp"]:
            if "mlp" in self.custom_args["layer_norm"]:
                hidden_states = self.ln_2(hidden_states)
            feed_forward_hidden_states = self.mlp(hidden_states)
            # residual connection
            if self.custom_args["is_resid_conn"]:
                hidden_states = residual + feed_forward_hidden_states
            else:
                hidden_states = feed_forward_hidden_states

        if use_cache:
            outputs = (hidden_states,) + outputs
        else:
            outputs = (hidden_states,) + outputs[1:]

        return outputs  # hidden_states, present, (attentions, cross_attentions)


class GPT2ModelCustom(GPT2Model):
    _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.attn\.bias", r"h\.\d+\.attn\.masked_bias"]
    _keys_to_ignore_on_load_missing = [r"attn.masked_bias", r"h\.\d+\.attn\.masked_bias", r"h\.\d+\.attn\.bias"]

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

        self.custom_args = custom_args
        customLayerNormClass = nn.LayerNorm 
        if custom_args["layer_norm_type"] == "no_gamma_beta":
            customLayerNormClass = partial(CustomLayerNorm, elementwise_affine=False)

        gpt2_blocks_list = []
        mlp_layers_list = self.custom_args["mlp_layers"]
        block_custom_args = copy.deepcopy(self.custom_args)
        block_custom_args['is_mlp'] = True
        for i in range(config.num_hidden_layers):
            curr_gpt2_block = None
            if i in mlp_layers_list:
                curr_gpt2_block = GPT2BlockCustom(config, block_custom_args, layer_idx=i)
            else:
                block_custom_args_copy = copy.deepcopy(block_custom_args)
                block_custom_args_copy["is_mlp"] = False
                curr_gpt2_block = GPT2BlockCustom(config, block_custom_args_copy, layer_idx=i)
            gpt2_blocks_list.append(curr_gpt2_block)
        self.h = nn.ModuleList(gpt2_blocks_list)
        self.ln_f = customLayerNormClass(self.embed_dim, eps=config.layer_norm_epsilon)
        if custom_args["diff_embs"]:
            self.ln_f = customLayerNormClass(custom_args["model_hid_dim"], eps=config.layer_norm_epsilon)

        if custom_args["diff_embs"]:
            self.scale_emb_linear = nn.Linear(self.embed_dim, custom_args["model_hid_dim"], bias=True)

        if not self.custom_args["is_pos_encode"] or self.custom_args['pos_enc_type'] != "learnable":
            del self.wpe
        if "final" not in self.custom_args["layer_norm"]:
            del self.ln_f

    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
    ) -> Union[Tuple, BaseModelOutputWithPastAndCrossAttentions]:
        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
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        hidden_states_pos = None
        hidden_states_tok = 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:
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
            batch_size = input_ids.shape[0]
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
            batch_size = inputs_embeds.shape[0]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device

        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])
        if position_ids is not None:
            position_ids = position_ids.view(-1, input_shape[-1])

        if past_key_values is None:
            past_length = 0
            past_key_values = tuple([None] * len(self.h))
        else:
            past_length = past_key_values[0][0].size(-2)
        if position_ids is None:
            position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1]) # [[0, 1, 2, .... max_len-1]] -- tensor of shape 1 x max_len

        # GPT2Attention mask.
        if attention_mask is not None:
            if batch_size <= 0:
                raise ValueError("batch_size has to be defined and > 0")
            attention_mask = attention_mask.view(batch_size, -1)
            
            # We create a 3D attention mask from a 2D tensor mask.
            # Sizes are [batch_size, 1, 1, to_seq_length]
            # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
            # this attention mask is more simple than the triangular masking of causal attention
            # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
            attention_mask = attention_mask[:, None, None, :]

            # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
            # masked positions, this operation will create a tensor which is 0.0 for
            # positions we want to attend and the dtype's smallest value for masked positions.
            # Since we are adding it to the raw scores before the softmax, this is
            # effectively the same as removing these entirely.
            attention_mask = attention_mask.to(dtype=self.dtype)  # fp16 compatibility
            attention_mask = (1.0 - attention_mask) * torch.finfo(self.dtype).min

        # If a 2D or 3D attention mask is provided for the cross-attention
        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
        if self.config.add_cross_attention and encoder_hidden_states is not None:
            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
            if encoder_attention_mask is None:
                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
            encoder_attention_mask = self.invert_attention_mask(encoder_attention_mask)
        else:
            encoder_attention_mask = None

        # 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
        # head_mask has shape n_layer x batch x n_heads x N x N
        head_mask = self.get_head_mask(head_mask, self.config.n_layer)

        if inputs_embeds is None:
            inputs_embeds = self.wte(input_ids) # batch, max_len, emb

        if "tok_emb" in self.custom_args["kqv_needs"]:
            hidden_states_tok = inputs_embeds
            hidden_states_tok = self.drop(hidden_states_tok)
        
        if self.custom_args["only_pos_at_pos_index"] is not None: # we want to pass only the position encoding as the input for these specified position indices, i.e. we want to zero out the token embedding at these positions
            index_to_zero_token_emb = self.custom_args["only_pos_at_pos_index"]
            inputs_embeds[:, index_to_zero_token_emb, :] = 0.

        if self.custom_args['is_pos_encode']:
            if self.custom_args['pos_enc_type'] == "sinusoidal":
                position_embeds = get_position_embeds(enc_type="sinusoidal", max_len=input_shape[-1], pos_emb_dim=self.embed_dim).to(device)
                hidden_states = inputs_embeds + position_embeds
            elif self.custom_args['pos_enc_type'] == "learnable":
                position_embeds = self.wpe(position_ids) # tensor of shape 1 x max_len x pos_emb_dim
                hidden_states = inputs_embeds + position_embeds

            if "pos_emb" in self.custom_args["kqv_needs"]:
                hidden_states_pos = position_embeds
                hidden_states_pos = self.drop(hidden_states_pos)
            if self.custom_args["diff_embs"]:
                hidden_states = self.scale_emb_linear(hidden_states)
        else:
            hidden_states = inputs_embeds

        if token_type_ids is not None:
            token_type_embeds = self.wte(token_type_ids)
            hidden_states = hidden_states + token_type_embeds

        hidden_states = self.drop(hidden_states)

        output_shape = input_shape + (hidden_states.size(-1),)

        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

        presents = () if use_cache else None
        all_self_attentions = () if output_attentions else None
        all_cross_attentions = () if output_attentions and self.config.add_cross_attention else None
        all_hidden_states = () if output_hidden_states else None
        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
            # Model parallel
            if self.model_parallel:
                torch.cuda.set_device(hidden_states.device)
                # Ensure layer_past is on same device as hidden_states (might not be correct)
                if layer_past is not None:
                    layer_past = tuple(past_state.to(hidden_states.device) for past_state in layer_past)
                # Ensure that attention_mask is always on the same device as hidden_states
                if attention_mask is not None:
                    attention_mask = attention_mask.to(hidden_states.device)
                if isinstance(head_mask, torch.Tensor):
                    head_mask = head_mask.to(hidden_states.device)
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

            if self.gradient_checkpointing and self.training:

                def create_custom_forward(module):
                    def custom_forward(*inputs):
                        # None for past_key_value
                        return module(*inputs, use_cache, output_attentions)

                    return custom_forward

                outputs = torch.utils.checkpoint.checkpoint(
                    create_custom_forward(block),
                    hidden_states,
                    None,
                    attention_mask,
                    head_mask[i],
                    encoder_hidden_states,
                    encoder_attention_mask,
                )
            else:
                outputs = block(
                    hidden_states,
                    hidden_states_tok=hidden_states_tok,
                    hidden_states_pos=hidden_states_pos,
                    layer_past=layer_past,
                    attention_mask=attention_mask,
                    head_mask=head_mask[i],
                    encoder_hidden_states=encoder_hidden_states,
                    encoder_attention_mask=encoder_attention_mask,
                    use_cache=use_cache,
                    output_attentions=output_attentions,
                )

            hidden_states = outputs[0]
            if use_cache is True:
                presents = presents + (outputs[1],)

            if output_attentions:
                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)
                if self.config.add_cross_attention:
                    all_cross_attentions = all_cross_attentions + (outputs[3 if use_cache else 2],)

            # Model Parallel: If it's the last layer for that device, put things on the next device
            if self.model_parallel:
                for k, v in self.device_map.items():
                    if i == v[-1] and "cuda:" + str(k) != self.last_device:
                        hidden_states = hidden_states.to("cuda:" + str(k + 1))

        if "final" in self.custom_args["layer_norm"]:
            hidden_states = self.ln_f(hidden_states)

        hidden_states = hidden_states.view(output_shape)
        # Add last hidden state
        if output_hidden_states:
            all_hidden_states = all_hidden_states + (hidden_states,)

        if not return_dict:
            return tuple(
                v
                for v in [hidden_states, presents, all_hidden_states, all_self_attentions, all_cross_attentions]
                if v is not None
            )

        return BaseModelOutputWithPastAndCrossAttentions(
            last_hidden_state=hidden_states,
            past_key_values=presents,
            hidden_states=all_hidden_states,
            attentions=all_self_attentions,
            cross_attentions=all_cross_attentions,
        )


class GPT2LMHeadModelCustom(GPT2LMHeadModel):
    _keys_to_ignore_on_load_missing = [r"lm_head.weight"]
    _keys_to_ignore_on_load_unexpected = [r"h\.\d+\.attn\.masked_bias", r"h\.\d+\.attn\.bias"]

    def __init__(self, config, custom_args):
        super().__init__(config)
        self.transformer = GPT2ModelCustom(config, custom_args)
        if custom_args["diff_embs"]:
            self.lm_head = nn.Linear(custom_args["model_hid_dim"], config.vocab_size, bias=False)
        
        self.post_init()    # add this line to tie weights
    
    def forward(
        self,
        input_ids: Optional[torch.LongTensor] = None,
        past_key_values: Optional[Tuple[Tuple[torch.Tensor]]] = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        token_type_ids: Optional[torch.LongTensor] = None,
        position_ids: Optional[torch.LongTensor] = None,
        head_mask: Optional[torch.FloatTensor] = None,
        inputs_embeds: Optional[torch.FloatTensor] = None,
        encoder_hidden_states: Optional[torch.Tensor] = None,
        encoder_attention_mask: Optional[torch.FloatTensor] = None,
        labels: Optional[torch.LongTensor] = None,
        use_cache: Optional[bool] = None,
        output_attentions: Optional[bool] = None,
        output_hidden_states: Optional[bool] = None,
        return_dict: Optional[bool] = None,
        loss_reduction="mean", # options: mean, none; mean -- loss is averaged over all tokens across all samples (a scalar is returned), none -- loss is computed for each token and not averaged (batch x seq tensor is returned)
        **kwargs
    ) -> Union[Tuple, CausalLMOutputWithCrossAttentions]:
        r"""
        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):
            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set
            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`
            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`
        """
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict

        transformer_outputs = self.transformer(
            input_ids,
            past_key_values=past_key_values,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
            use_cache=use_cache,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
            return_dict=return_dict,
        )
        hidden_states = transformer_outputs[0]

        # Set device for model parallelism
        if self.model_parallel:
            torch.cuda.set_device(self.transformer.first_device)
            hidden_states = hidden_states.to(self.lm_head.weight.device)

        lm_logits = self.lm_head(hidden_states)

        loss = None
        if labels is not None:
            # move labels to correct device to enable model parallelism
            labels = labels.to(lm_logits.device)
            # Shift so that tokens < n predict n
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = labels[..., 1:].contiguous()
            if 'loss_mask' in kwargs: # this runs for addition, greater than, etc. tasks where we pass loss mask to the model
                loss_mask = kwargs['loss_mask']
                shift_loss_mask = loss_mask[..., 1:].contiguous()
                loss_fct = CrossEntropyLoss(reduction="none")
                loss_x = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
                # Resize and average loss per sample
                non_reduced_loss_per_sample = loss_x.view(shift_logits.size(0), shift_logits.size(1)) # batch x seq
                
                if loss_reduction == "mean":
                    loss = (non_reduced_loss_per_sample * shift_loss_mask).sum()/shift_loss_mask.sum()
                elif loss_reduction == "none":
                    loss = non_reduced_loss_per_sample

            else: # this runs for main world models train.py where we do not pass loss mask to the model
                # Flatten the tokens
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))

        if not return_dict:
            output = (lm_logits,) + transformer_outputs[1:]
            return ((loss,) + output) if loss is not None else output

        return CausalLMOutputWithCrossAttentions(
            loss=loss,
            logits=lm_logits,
            past_key_values=transformer_outputs.past_key_values,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
            cross_attentions=transformer_outputs.cross_attentions,
        )