refactai/codify_3b_multi

1import math2import warnings3from typing import Optional, Tuple, Union4 5import torch6import torch.utils.checkpoint7from torch import nn8from torch.nn import CrossEntropyLoss, LayerNorm9from torch.nn import functional as F10from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, \11    add_start_docstrings_to_model_forward12from transformers.modeling_outputs import (13    BaseModelOutputWithPast,14    CausalLMOutputWithPast,15)16from transformers.modeling_utils import PreTrainedModel17from transformers.utils import logging18 19from .configuration_codify import CodifyConfig20 21logger = logging.get_logger(__name__)22 23_CHECKPOINT_FOR_DOC = "smallcloudai/codify_medium_multi"24_CONFIG_FOR_DOC = "CodifyConfig"25_TOKENIZER_FOR_DOC = "CodifyTokenizerFast"26 27 28CODIFY_PRETRAINED_MODEL_ARCHIVE_LIST = [29    "smallcloudai/codify_medium_multi",30    "smallcloudai/codify_3b_multi"31]32 33def _make_causal_mask(34    input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int35) -> torch.BoolTensor:36    """37    Make causal mask used for self-attention.38    """39    batch_size, target_length = input_ids_shape40    mask = torch.empty((target_length, target_length + past_key_values_length), dtype=torch.bool, device=device)41    # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround42    seq_ids = torch.arange(target_length, device=device)43    mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]44 45    if past_key_values_length > 0:46        mask[:, :past_key_values_length] = False47 48    expanded_mask = mask[None, None, :, :].expand(batch_size, 1, target_length, target_length + past_key_values_length)49    return expanded_mask50 51 52def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:53    """54    Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.55    """56    batch_size, src_length = mask.shape57    tgt_length = tgt_length if tgt_length is not None else src_length58 59    expanded_mask = ~(mask[:, None, None, :].to(torch.bool))60    return expanded_mask.expand(batch_size, 1, tgt_length, src_length)61 62 63def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:64    """65    Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it66    relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value67    `softmax(l+a) = softmax(l)`. Based on68    https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L74269    TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly.70 71    Args:72    Returns tensor shaped (batch_size * num_heads, 1, max_seq_len)73        attention_mask (`torch.Tensor`):74            Token-wise attention mask, this should be of shape (batch_size, max_seq_len).75        num_heads (`int`, *required*):76            number of heads77        dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`):78            dtype of the output tensor79    """80    batch_size, seq_length = attention_mask.shape81    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))82    base = torch.tensor(83        2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float3284    )85    powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)86    slopes = torch.pow(base, powers)87 88    if closest_power_of_2 != num_heads:89        extra_base = torch.tensor(90            2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float3291        )92        num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)93        extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)94        slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)95 96    # Note: alibi will added to the attention bias that will be applied to the query, key product of attention97    # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)98    # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)99    # => the query_length dimension will then be broadcasted correctly100    # This is more or less identical to T5's relative position bias:101    # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527102    arange_tensor = ((attention_mask.cumsum(dim=-1)) * attention_mask)[:, None, :]103    alibi = slopes[..., None] * arange_tensor104    return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)105 106 107 108def codify_gelu_forward(x: torch.Tensor) -> torch.Tensor:109    """110    Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to111    make the model jitable.112 113    Args:114        x (`torch.tensor`, *required*):115            input hidden states116    """117    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))118 119 120def codify_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:121    """122    gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) +123    0.3989423 * x * torch.exp(-0.5 * x * x)124 125    Args:126        g (`torch.tensor`, *required*):127            gradient output tensor128        x (`torch.tensor`, *required*):129            input tensor130    """131    x = x[0]  # x is a tuple of 1 element, needs to unpack it first132    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))133    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243134    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)135    return ff * g136 137 138class GeLUFunction(torch.autograd.Function):139    @staticmethod140    def forward(ctx, input: torch.Tensor) -> torch.Tensor:141        ctx.save_for_backward(input)142        return codify_gelu_forward(input)143 144    @staticmethod145    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:146        input = ctx.saved_tensors147        tmp = codify_gelu_back(grad_output, input)148        return tmp149 150 151class CodifyGelu(nn.Module):152    def __init__(self):153        super().__init__()154 155    def forward(self, x: torch.Tensor) -> torch.Tensor:156        if self.training:157            return GeLUFunction.apply(x)158        else:159            return codify_gelu_forward(x)160 161 162class CodifyAttention(nn.Module):163    def __init__(self, config: CodifyConfig):164        super().__init__()165 166        self.hidden_size = config.hidden_size167        self.num_heads = config.num_attention_heads168        self.head_dim = self.hidden_size // self.num_heads169        self.split_size = self.hidden_size170 171        if self.head_dim * self.num_heads != self.hidden_size:172            raise ValueError(173                f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"174                f" {self.num_heads})."175            )176 177        # Layer-wise attention scaling178        # 8.0 = self.head_dim179        self.inv_norm_factor = 8.0 / self.head_dim180        self.beta = 1.0181 182        self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=True)183        self.dense = nn.Linear(self.hidden_size, self.hidden_size)184 185    def _split_heads(self, fused_qkv: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:186        """187        Split the last dimension into (num_heads, head_dim) without making any copies, results share same memory188        storage as `fused_qkv`189 190        Args:191            fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim]192 193        Returns:194            query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim]195            value: [batch_size, seq_length, num_heads, head_dim]196        """197        batch_size, seq_length, _ = fused_qkv.shape198        q, k, v = fused_qkv.chunk(3, dim=-1)199        return q.view(batch_size, seq_length, self.num_heads, self.head_dim),\200            k.view(batch_size, seq_length, self.num_heads, self.head_dim),\201            v.view(batch_size, seq_length, self.num_heads, self.head_dim)202 203    def _merge_heads(self, x: torch.Tensor) -> torch.Tensor:204        """205        Merge heads together over the last dimenstion206 207        Args:208            x: (`torch.tensor`, *required*): [batch_size * num_heads, seq_length, head_dim]209 210        Returns:211            torch.tensor: [batch_size, seq_length, num_heads * head_dim]212        """213        # What we want to achieve is:214        # batch_size * num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads * head_dim215        batch_size_and_num_heads, seq_length, _ = x.shape216        batch_size = batch_size_and_num_heads // self.num_heads217 218        # First view to decompose the batch size219        # batch_size * num_heads, seq_length, head_dim -> batch_size, num_heads, seq_length, head_dim220        x = x.view(batch_size, self.num_heads, seq_length, self.head_dim)221 222        # batch_size, num_heads, seq_length, head_dim -> batch_size, seq_length, num_heads, head_dim223        x = x.permute(0, 2, 1, 3)224 225        # batch_size, seq_length, num_heads, head_dim -> batch_size, seq_length, num_heads * head_dim226        return x.reshape(batch_size, seq_length, self.num_heads * self.head_dim)227 228    def forward(229        self,230        hidden_states: torch.Tensor,231        alibi: torch.Tensor,232        attention_mask: torch.Tensor,233        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,234        head_mask: Optional[torch.Tensor] = None,235        use_cache: bool = False,236        output_attentions: bool = False,237    ):238        fused_qkv = self.query_key_value(hidden_states)  # [batch_size, seq_length, 3 x hidden_size]239 240        # 3 x [batch_size, seq_length, num_heads, head_dim]241        (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)242 243        batch_size, q_length, _, _ = query_layer.shape244 245        query_layer = query_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)246        key_layer = key_layer.permute(0, 2, 3, 1).reshape(batch_size * self.num_heads, self.head_dim, q_length)247        value_layer = value_layer.transpose(1, 2).reshape(batch_size * self.num_heads, q_length, self.head_dim)248        if layer_past is not None:249            past_key, past_value = layer_past250            # concatenate along seq_length dimension:251            #  - key: [batch_size * self.num_heads, head_dim, kv_length]252            #  - value: [batch_size * self.num_heads, kv_length, head_dim]253            key_layer = torch.cat((past_key, key_layer), dim=2)254            value_layer = torch.cat((past_value, value_layer), dim=1)255 256        _, _, kv_length = key_layer.shape257 258        if use_cache is True:259            present = (key_layer, value_layer)260        else:261            present = None262 263        # [batch_size * num_heads, q_length, kv_length]264        # we use `torch.Tensor.baddbmm` instead of `torch.baddbmm` as the latter isn't supported by TorchScript v1.11265        matmul_result = alibi.baddbmm(266            batch1=query_layer,267            batch2=key_layer,268            beta=self.beta,269            alpha=self.inv_norm_factor,270        )271 272        # change view to [batch_size, num_heads, q_length, kv_length]273        attention_scores = matmul_result.view(batch_size, self.num_heads, q_length, kv_length)274 275        # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]276        input_dtype = attention_scores.dtype277        # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`278        if input_dtype == torch.float16:279            attention_scores = attention_scores.to(torch.float)280        attn_weights = torch.masked_fill(attention_scores, attention_mask, torch.finfo(attention_scores.dtype).min)281        attention_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(input_dtype)282 283        if head_mask is not None:284            attention_probs = attention_probs * head_mask285 286        # change view [batch_size x num_heads, q_length, kv_length]287        attention_probs_reshaped = attention_probs.view(batch_size * self.num_heads, q_length, kv_length)288 289        # matmul: [batch_size * num_heads, q_length, head_dim]290        context_layer = torch.bmm(attention_probs_reshaped, value_layer)291 292        # change view [batch_size, num_heads, q_length, head_dim]293        context_layer = self._merge_heads(context_layer)294 295        output_tensor = self.dense(context_layer)296        outputs = (output_tensor, present)297        if output_attentions:298            outputs += (attention_probs,)299 300        return outputs301 302 303class CodifyMLP(nn.Module):304    def __init__(self, config: CodifyConfig):305        super().__init__()306        hidden_size = config.hidden_size307        self.dense_h_to_4h = nn.Linear(hidden_size, config.mlp_mult * hidden_size)308        self.gelu_impl = CodifyGelu()309        self.dense_4h_to_h = nn.Linear(config.mlp_mult * hidden_size, hidden_size)310 311    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:312        hidden_states = self.gelu_impl(self.dense_h_to_4h(hidden_states))313        output = self.dense_4h_to_h(hidden_states)314        return output315 316 317class CodifyBlock(nn.Module):318    def __init__(self, config: CodifyConfig):319        super().__init__()320        hidden_size = config.hidden_size321 322        self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)323        self.num_heads = config.num_attention_heads324        self.self_attention = CodifyAttention(config)325        self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)326        self.mlp = CodifyMLP(config)327 328    def forward(329        self,330        hidden_states: torch.Tensor,331        alibi: torch.Tensor,332        attention_mask: torch.Tensor,333        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,334        head_mask: Optional[torch.Tensor] = None,335        use_cache: bool = False,336        output_attentions: bool = False,337    ):338        # hidden_states: [batch_size, seq_length, hidden_size]339 340        # Layer norm at the beginning of the transformer layer.341        layernorm_output = self.input_layernorm(hidden_states)342 343        # Self attention.344        attn_outputs = self.self_attention(345            layernorm_output,346            layer_past=layer_past,347            attention_mask=attention_mask,348            alibi=alibi,349            head_mask=head_mask,350            use_cache=use_cache,351            output_attentions=output_attentions,352        )353 354        attention_output = attn_outputs[0]355        outputs = attn_outputs[1:]356 357        attention_mix = attention_output + hidden_states358        layernorm_output = self.post_attention_layernorm(attention_mix)359 360        # MLP.361        output = self.mlp(layernorm_output)362        output = output + attention_output + hidden_states363 364        if use_cache:365            outputs = (output,) + outputs366        else:367            outputs = (output,) + outputs[1:]368 369        return outputs  # hidden_states, present, attentions370 371class CodifyPreTrainedModel(PreTrainedModel):372    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]373    """374    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained375    models.376    """377 378    config_class = CodifyConfig379    base_model_prefix = "transformer"380    supports_gradient_checkpointing = True381    _no_split_modules = ["CodifyBlock"]382 383    def __init__(self, *inputs, **kwargs):384        super().__init__(*inputs, **kwargs)385 386    def _init_weights(self, module: nn.Module):387        """Initialize the weights."""388        if isinstance(module, nn.Linear):389            # Slightly different from the TF version which uses truncated_normal for initialization390            # cf https://github.com/pytorch/pytorch/pull/5617391            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)392            if module.bias is not None:393                module.bias.data.zero_()394        elif isinstance(module, nn.Embedding):395            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)396            if module.padding_idx is not None:397                module.weight.data[module.padding_idx].zero_()398        elif isinstance(module, LayerNorm):399            module.bias.data.zero_()400            module.weight.data.fill_(1.0)401 402    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):403        if isinstance(module, CodifyModel):404            module.gradient_checkpointing = value405 406    @staticmethod407    def _convert_to_standard_cache(408        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int409    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:410        """411        Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size,412        num_heads, ...]))413        """414        batch_size_times_num_heads, head_dim, seq_length = past_key_value[0][0].shape415        num_heads = batch_size_times_num_heads // batch_size416        # key: [batch_size * num_heads, head_dim, seq_length] -> [batch_size, num_heads, head_dim, seq_length]417        # value: [batch_size * num_heads, seq_length, head_dim] -> [batch_size, num_heads, seq_length, head_dim]418        return tuple(419            (420                layer_past[0].view(batch_size, num_heads, head_dim, seq_length),421                layer_past[1].view(batch_size, num_heads, seq_length, head_dim),422            )423            for layer_past in past_key_value424        )425 426    @staticmethod427    def _convert_to_codify_cache(428        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]429    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:430        batch_size, num_heads, head_dim, seq_length = past_key_value[0][0].shape431        batch_size_times_num_heads = batch_size * num_heads432        # key:  [batch_size, num_heads, head_dim, seq_length] -> [batch_size * num_heads, head_dim, seq_length]433        # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size * num_heads, seq_length, head_dim]434        return tuple(435            (436                layer_past[0].view(batch_size_times_num_heads, head_dim, seq_length),437                layer_past[1].view(batch_size_times_num_heads, seq_length, head_dim),438            )439            for layer_past in past_key_value440        )441 442class CodifyModel(CodifyPreTrainedModel):443    def __init__(self, config: CodifyConfig):444        super().__init__(config)445 446        self.embed_dim = config.hidden_size447        self.num_heads = config.num_attention_heads448 449        # Embedding450        self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)451 452        # Transformer blocks453        self.h = nn.ModuleList([CodifyBlock(config) for _ in range(config.num_hidden_layers)])454 455        # Final Layer Norm456        self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)457 458        self.gradient_checkpointing = False459 460        # Initialize weights and apply final processing461        self.post_init()462 463    def get_input_embeddings(self):464        return self.word_embeddings465 466    def _prepare_attn_mask(467        self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int468    ) -> torch.BoolTensor:469        # create causal mask470        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]471        combined_attention_mask = None472        device = attention_mask.device473        _, src_length = input_shape474 475        if src_length > 1:476            combined_attention_mask = _make_causal_mask(477                input_shape, device=device, past_key_values_length=past_key_values_length478            )479 480        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]481        expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)482        combined_attention_mask = (483            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask484        )485 486        return combined_attention_mask487 488    def set_input_embeddings(self, new_embeddings: torch.Tensor):489        self.word_embeddings = new_embeddings490 491    @add_code_sample_docstrings(492        processor_class=_TOKENIZER_FOR_DOC,493        checkpoint=_CHECKPOINT_FOR_DOC,494        output_type=BaseModelOutputWithPast,495        config_class=_CONFIG_FOR_DOC,496    )497    def forward(498        self,499        input_ids: Optional[torch.LongTensor] = None,500        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,501        attention_mask: Optional[torch.Tensor] = None,502        head_mask: Optional[torch.LongTensor] = None,503        inputs_embeds: Optional[torch.LongTensor] = None,504        use_cache: Optional[bool] = None,505        output_attentions: Optional[bool] = None,506        output_hidden_states: Optional[bool] = None,507        return_dict: Optional[bool] = None,508        **deprecated_arguments509    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPast]:510        if deprecated_arguments.pop("position_ids", False) is not False:511            # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`512            warnings.warn(513                "`position_ids` have no functionality in Codify and will be removed in v5.0.0. You can safely ignore"514                " passing `position_ids`.",515                FutureWarning,516            )517        if len(deprecated_arguments) > 0:518            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")519 520        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions521        output_hidden_states = (522            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states523        )524        use_cache = use_cache if use_cache is not None else self.config.use_cache525        return_dict = return_dict if return_dict is not None else self.config.use_return_dict526 527        if input_ids is not None and inputs_embeds is not None:528            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")529        elif input_ids is not None:530            batch_size, seq_length = input_ids.shape531        elif inputs_embeds is not None:532            batch_size, seq_length, _ = inputs_embeds.shape533        else:534            raise ValueError("You have to specify either input_ids or inputs_embeds")535 536        if past_key_values is None:537            past_key_values = tuple([None] * len(self.h))538 539        # Prepare head mask if needed540        # 1.0 in head_mask indicate we keep the head541        # attention_probs has shape batch_size x num_heads x N x N542        # head_mask has shape n_layer x batch x num_heads x N x N543        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)544 545        if inputs_embeds is None:546            inputs_embeds = self.word_embeddings(input_ids)547 548        hidden_states = inputs_embeds549 550        presents = () if use_cache else None551        all_self_attentions = () if output_attentions else None552        all_hidden_states = () if output_hidden_states else None553 554        # Compute alibi tensor: check build_alibi_tensor documentation555        seq_length_with_past = seq_length556        past_key_values_length = 0557        if past_key_values[0] is not None:558            past_key_values_length = past_key_values[0][0].shape[2]559            seq_length_with_past = seq_length_with_past + past_key_values_length560        if attention_mask is None:561            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)562        else:563            attention_mask = attention_mask.to(hidden_states.device)564 565        alibi = build_alibi_tensor(attention_mask, self.num_heads, dtype=hidden_states.dtype)566 567        causal_mask = self._prepare_attn_mask(568            attention_mask,569            input_shape=(batch_size, seq_length),570            past_key_values_length=past_key_values_length,571        )572 573        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):574 575            if output_hidden_states:576                all_hidden_states = all_hidden_states + (hidden_states,)577 578            if self.gradient_checkpointing and self.training:579 580                if use_cache:581                    logger.warning(582                        "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."583                    )584                    use_cache = False585 586                def create_custom_forward(module):587                    def custom_forward(*inputs):588                        # None for past_key_value589                        return module(*inputs, use_cache=use_cache, output_attentions=output_attentions)590 591                    return custom_forward592 593                outputs = torch.utils.checkpoint.checkpoint(594                    create_custom_forward(block),595                    hidden_states,596                    alibi,597                    causal_mask,598                    head_mask[i],599                )600            else:601                outputs = block(602                    hidden_states,603                    layer_past=layer_past,604                    attention_mask=causal_mask,605                    head_mask=head_mask[i],606                    use_cache=use_cache,607                    output_attentions=output_attentions,608                    alibi=alibi,609                )610 611            hidden_states = outputs[0]612            if use_cache is True:613                presents = presents + (outputs[1],)614 615            if output_attentions:616                all_self_attentions = all_self_attentions + (outputs[2 if use_cache else 1],)617 618        # Add last hidden state619        hidden_states = self.ln_f(hidden_states)620 621        if output_hidden_states:622            all_hidden_states = all_hidden_states + (hidden_states,)623 624        if not return_dict:625            return tuple(v for v in [hidden_states, presents, all_hidden_states, all_self_attentions] if v is not None)626 627        return BaseModelOutputWithPast(628            last_hidden_state=hidden_states,629            past_key_values=presents,630            hidden_states=all_hidden_states,631            attentions=all_self_attentions,632        )633 634 635class CodifyForCausalLM(CodifyPreTrainedModel):636    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]637 638    def __init__(self, config: CodifyConfig):639        super().__init__(config)640        self.transformer = CodifyModel(config)641        self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)642 643        # Initialize weights and apply final processing644        self.post_init()645 646    def get_output_embeddings(self):647        return self.lm_head648 649    def set_output_embeddings(self, new_embeddings: torch.Tensor):650        self.lm_head = new_embeddings651 652    def prepare_inputs_for_generation(653        self,654        input_ids: torch.LongTensor,655        past: Optional[torch.Tensor] = None,656        attention_mask: Optional[torch.Tensor] = None,657        **kwargs658    ) -> dict:659        # only last token for input_ids if past is not None660        if past:661            input_ids = input_ids[:, -1].unsqueeze(-1)662 663            if past[0][0].shape[0] == input_ids.shape[0]:664                past = self._convert_to_codify_cache(past)665 666        return {667            "input_ids": input_ids,668            "past_key_values": past,669            "use_cache": kwargs.get("use_cache"),670            "attention_mask": attention_mask,671        }672 673    @add_code_sample_docstrings(674        processor_class=_TOKENIZER_FOR_DOC,675        checkpoint=_CHECKPOINT_FOR_DOC,676        output_type=CausalLMOutputWithPast,677        config_class=_CONFIG_FOR_DOC,678    )679    def forward(680        self,681        input_ids: Optional[torch.LongTensor] = None,682        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,683        attention_mask: Optional[torch.Tensor] = None,684        head_mask: Optional[torch.Tensor] = None,685        inputs_embeds: Optional[torch.Tensor] = None,686        labels: Optional[torch.Tensor] = None,687        use_cache: Optional[bool] = None,688        output_attentions: Optional[bool] = None,689        output_hidden_states: Optional[bool] = None,690        return_dict: Optional[bool] = None,691        **deprecated_arguments692    ) -> Union[Tuple[torch.Tensor], CausalLMOutputWithPast]:693        r"""694        labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):695            Labels for language modeling. Note that the labels **are shifted** inside the model, i.e. you can set696            `labels = input_ids` Indices are selected in `[-100, 0, ..., config.vocab_size]` All labels set to `-100`697            are ignored (masked), the loss is only computed for labels in `[0, ..., config.vocab_size]`698        """699        if deprecated_arguments.pop("position_ids", False) is not False:700            # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`701            warnings.warn(702                "`position_ids` have no functionality in Codify and will be removed in v5.0.0. You can safely ignore"703                " passing `position_ids`.",704                FutureWarning,705            )706        if len(deprecated_arguments) > 0:707            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")708 709        return_dict = return_dict if return_dict is not None else self.config.use_return_dict710 711        transformer_outputs = self.transformer(712            input_ids,713            past_key_values=past_key_values,714            attention_mask=attention_mask,715            head_mask=head_mask,716            inputs_embeds=inputs_embeds,717            use_cache=use_cache,718            output_attentions=output_attentions,719            output_hidden_states=output_hidden_states,720            return_dict=return_dict,721        )722        hidden_states = transformer_outputs[0]723 724        lm_logits = self.lm_head(hidden_states / 2.0)725 726        loss = None727        if labels is not None:728            # Shift so that tokens < n predict n729            shift_logits = lm_logits[..., :-1, :].contiguous()730            shift_labels = labels[..., 1:].contiguous()731            batch_size, seq_length, vocab_size = shift_logits.shape732            # Flatten the tokens733            loss_fct = CrossEntropyLoss()734            loss = loss_fct(735                shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)736            )737 738        if not return_dict:739            output = (lm_logits,) + transformer_outputs[1:]740            return ((loss,) + output) if loss is not None else output741 742        return CausalLMOutputWithPast(743            loss=loss,744            logits=lm_logits,745            past_key_values=transformer_outputs.past_key_values,746            hidden_states=transformer_outputs.hidden_states,747            attentions=transformer_outputs.attentions,748        )749 750    def _reorder_cache(751        self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor752    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:753        """754        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or755        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct756        beam_idx at every generation step.757 758        Output shares the same memory storage as `past`.759        """760        standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx))761 762        # Get a copy of `beam_idx` on all the devices where we need those indices.763        device_to_beam_idx = {764            past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past765        }766        reordered_past = tuple(767            (768                layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),769                layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),770            )771            for layer_past in standardized_past772        )773        return self._convert_to_codify_cache(reordered_past)774