start point

README.md

@@ -1,3 +1 @@
----
-license: apache-2.0
----
+https://github.com/CoinCheung/gdGPT

config.json

@@ -0,0 +1,40 @@
+{
+  "_name_or_path": "coincheung/cc-bloomz-7b",
+  "apply_residual_connection_post_layernorm": false,
+  "architectures": [
+    "CCBloomForCausalLM"
+  ],
+  "attention_dropout": 0.0,
+  "attention_softmax_in_fp32": true,
+  "auto_map": {
+    "AutoConfig": "configuration_cc_bloom.CCBloomConfig",
+    "AutoModel": "modeling_cc_bloom.CCBloomForCausalLM",
+    "AutoModelForCausalLM": "modeling_cc_bloom.CCBloomForCausalLM"
+  },
+  "bias_dropout_fusion": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "hidden_dropout": 0.0,
+  "hidden_size": 4096,
+  "initializer_range": 0.02,
+  "layer_norm_epsilon": 1e-05,
+  "masked_softmax_fusion": true,
+  "model_type": "ccbloom",
+  "n_head": 32,
+  "n_inner": null,
+  "n_layer": 30,
+  "pad_token_id": 0,
+  "pretraining_tp": 1,
+  "rope_scaling": null,
+  "seq_length": 2048,
+  "skip_bias_add": true,
+  "skip_bias_add_qkv": false,
+  "slow_but_exact": false,
+  "tie_word_embeddings": false,
+  "torch_dtype": "float32",
+  "transformers_version": "4.31.0",
+  "unk_token_id": 0,
+  "use_cache": true,
+  "use_flash_attn": false,
+  "vocab_size": 64000
+}

configuration_cc_bloom.py

@@ -0,0 +1,127 @@
+# coding=utf-8
+# Copyright 2022 the Big Science Workshop and HuggingFace Inc. team.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Bloom configuration"""
+from collections import OrderedDict
+from typing import TYPE_CHECKING, Any, List, Mapping, Optional
+
+from packaging import version
+
+
+if TYPE_CHECKING:
+    from transformers import PreTrainedTokenizer, TensorType
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import is_torch_available, logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class CCBloomConfig(PretrainedConfig):
+    """
+    This is the configuration class to store the configuration of a [`BloomModel`]. It is used to instantiate a Bloom
+    model according to the specified arguments, defining the model architecture. Instantiating a configuration with the
+    defaults will yield a similar configuration to the Bloom architecture
+    [bigscience/bloom](https://huggingface.co/bigscience/bloom).
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 250880):
+            Vocabulary size of the Bloom model. Defines the maximum number of different tokens that can be represented
+            by the `inputs_ids` passed when calling [`BloomModel`]. Check [this
+            discussion](https://huggingface.co/bigscience/bloom/discussions/120#633d28389addb8530b406c2a) on how the
+            `vocab_size` has been defined.
+        hidden_size (`int`, *optional*, defaults to 64):
+            Dimensionality of the embeddings and hidden states.
+        n_layer (`int`, *optional*, defaults to 2):
+            Number of hidden layers in the Transformer encoder.
+        n_head (`int`, *optional*, defaults to 8):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
+            The epsilon to use in the layer normalization layers.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        apply_residual_connection_post_layernorm (`bool`, *optional*, defaults to `False`):
+            If enabled, use the layer norm of the hidden states as the residual in the transformer blocks
+        hidden_dropout (`float`, *optional*, defaults to 0.1):
+            Dropout rate of the dropout function on the bias dropout.
+        attention_dropout (`float`, *optional*, defaults to 0.1):
+            Dropout rate applied to the attention probs
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models).
+        pretraining_tp (`int`, *optional*, defaults to `1`):
+            Experimental feature. Tensor parallelism rank used during pretraining with Megatron. Please refer to [this
+            document](https://huggingface.co/docs/transformers/parallelism) to understand more about it. This value is
+            necessary to ensure exact reproducibility of the pretraining results. Please refer to [this
+            issue](https://github.com/pytorch/pytorch/issues/76232). Note also that this is enabled only when
+            `slow_but_exact=True`.
+        slow_but_exact (`bool`, *optional*, defaults to `False`):
+            Experimental feature. Whether to use slow but exact implementation of the attention mechanism. While
+            merging the TP rank tensors, due to slicing operations the results may be slightly different between the
+            model trained on Megatron and our model. Please refer to [this
+            issue](https://github.com/pytorch/pytorch/issues/76232). A solution to obtain more accurate results is to
+            enable this feature. Enabling this will hurt the computational time of the inference. Will be probably
+            resolved in the future once the main model has been fine-tuned with TP_rank=1.
+    """
+
+    model_type = "ccbloom"
+    keys_to_ignore_at_inference = ["past_key_values"]
+    attribute_map = {
+        "num_hidden_layers": "n_layer",
+        "num_attention_heads": "n_head",
+    }
+
+    def __init__(
+        self,
+        vocab_size=250880,
+        hidden_size=64,
+        n_layer=2,
+        n_head=8,
+        layer_norm_epsilon=1e-5,
+        initializer_range=0.02,
+        use_cache=True,
+        bos_token_id=1,
+        eos_token_id=2,
+        apply_residual_connection_post_layernorm=False,
+        hidden_dropout=0.0,
+        attention_dropout=0.0,
+        pretraining_tp=1,  # TP rank used when training with megatron
+        slow_but_exact=False,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        # Backward compatibility with n_embed kwarg
+        n_embed = kwargs.pop("n_embed", None)
+        self.hidden_size = hidden_size if n_embed is None else n_embed
+        self.n_layer = n_layer
+        self.n_head = n_head
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.initializer_range = initializer_range
+        self.use_cache = use_cache
+        self.pretraining_tp = pretraining_tp
+        self.apply_residual_connection_post_layernorm = apply_residual_connection_post_layernorm
+        self.hidden_dropout = hidden_dropout
+        self.attention_dropout = attention_dropout
+
+        self.bos_token_id = bos_token_id
+        self.eos_token_id = eos_token_id
+        self.slow_but_exact = slow_but_exact
+
+        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
+
+

generation_config.json

@@ -0,0 +1,7 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": 2,
+  "pad_token_id": 3,
+  "transformers_version": "4.31.0"
+}

modeling_cc_bloom.py

@@ -0,0 +1,1552 @@
+# coding=utf-8
+# Copyright 2022 HuggingFace Inc. team and BigScience workshop.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BLOOM model."""
+
+import os
+import os.path as osp
+import math
+import warnings
+from typing import Optional, Tuple, Union
+
+import torch
+import torch.utils.checkpoint
+from torch import nn
+from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, LayerNorm, MSELoss
+from torch.nn import functional as F
+
+from transformers.file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
+from transformers.modeling_outputs import (
+    BaseModelOutputWithPastAndCrossAttentions,
+    CausalLMOutputWithCrossAttentions,
+    QuestionAnsweringModelOutput,
+    SequenceClassifierOutputWithPast,
+    TokenClassifierOutput,
+)
+from transformers.modeling_utils import PreTrainedModel
+from transformers.utils import logging
+
+try:
+    from flash_attn import flash_attn_func
+except ImportError:
+    flash_attn_func = None
+
+from transformers.models.bloom.configuration_bloom import BloomConfig as CCBloomConfig
+
+#  from transformers.models.llama.modeling_llama import LlamaDynamicNTKScalingRotaryEmbedding, LlamaLinearScalingRotaryEmbedding, LlamaRotaryEmbedding
+
+
+logger = logging.get_logger(__name__)
+
+_CHECKPOINT_FOR_DOC = "bigscience/bloom-560m"
+_CONFIG_FOR_DOC = "CCBloomConfig"
+
+BLOOM_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "bigscience/bigscience-small-testing",
+    "bigscience/bloom-560m",
+    "bigscience/bloom-1b1",
+    "bigscience/bloom-1b7",
+    "bigscience/bloom-3b",
+    "bigscience/bloom-7b1",
+    "bigscience/bloom",
+]
+
+
+def _make_causal_mask(
+    input_ids_shape: torch.Size, device: torch.device, past_key_values_length: int
+) -> torch.BoolTensor:
+    """
+    Make causal mask used for self-attention.
+    """
+    batch_size, target_length = input_ids_shape
+    mask = torch.empty((target_length, target_length + past_key_values_length), dtype=torch.bool, device=device)
+    # ONNX doesn't support `torch.Tensor.triu` properly, thus we use this workaround
+    seq_ids = torch.arange(target_length, device=device)
+    mask[:, past_key_values_length:] = seq_ids[:, None] < seq_ids[None, :]
+
+    if past_key_values_length > 0:
+        mask[:, :past_key_values_length] = False
+
+    expanded_mask = mask[None, None, :, :].expand(batch_size, 1, target_length, target_length + past_key_values_length)
+    return expanded_mask
+
+
+def _expand_mask(mask: torch.Tensor, tgt_length: int) -> torch.BoolTensor:
+    """
+    Expands attention_mask from `[batch_size, src_length]` to `[batch_size, 1, tgt_length, src_length]`.
+    """
+    batch_size, src_length = mask.shape
+    tgt_length = tgt_length if tgt_length is not None else src_length
+
+    expanded_mask = ~(mask[:, None, None, :].to(torch.bool))
+    return expanded_mask.expand(batch_size, 1, tgt_length, src_length)
+
+
+def build_alibi_tensor(attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
+    """
+    Link to paper: https://arxiv.org/abs/2108.12409 Alibi tensor is not causal as the original paper mentions, it
+    relies on a translation invariance of softmax for quick implementation: with l being a tensor, and a fixed value
+    `softmax(l+a) = softmax(l)`. Based on
+    https://github.com/ofirpress/attention_with_linear_biases/blob/a35aaca144e0eb6b789dfcb46784c4b8e31b7983/fairseq/models/transformer.py#L742
+    TODO @thomasw21 this doesn't work as nicely due to the masking strategy, and so masking varies slightly.
+
+    Args:
+    Returns tensor shaped (batch_size * num_heads, 1, max_seq_len)
+        attention_mask (`torch.Tensor`):
+            Token-wise attention mask, this should be of shape (batch_size, max_seq_len).
+        num_heads (`int`, *required*):
+            number of heads
+        dtype (`torch.dtype`, *optional*, default=`torch.bfloat16`):
+            dtype of the output tensor
+    """
+    batch_size, seq_length = attention_mask.shape
+    closest_power_of_2 = 2 ** math.floor(math.log2(num_heads))
+    base = torch.tensor(
+        2 ** (-(2 ** -(math.log2(closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+    )
+    powers = torch.arange(1, 1 + closest_power_of_2, device=attention_mask.device, dtype=torch.int32)
+    slopes = torch.pow(base, powers)
+
+    if closest_power_of_2 != num_heads:
+        extra_base = torch.tensor(
+            2 ** (-(2 ** -(math.log2(2 * closest_power_of_2) - 3))), device=attention_mask.device, dtype=torch.float32
+        )
+        num_remaining_heads = min(closest_power_of_2, num_heads - closest_power_of_2)
+        extra_powers = torch.arange(1, 1 + 2 * num_remaining_heads, 2, device=attention_mask.device, dtype=torch.int32)
+        slopes = torch.cat([slopes, torch.pow(extra_base, extra_powers)], dim=0)
+
+    # Note: alibi will added to the attention bias that will be applied to the query, key product of attention
+    # => therefore alibi will have to be of shape (batch_size, num_heads, query_length, key_length)
+    # => here we set (batch_size=1, num_heads=num_heads, query_length=1, key_length=max_length)
+    # => the query_length dimension will then be broadcasted correctly
+    # This is more or less identical to T5's relative position bias:
+    # https://github.com/huggingface/transformers/blob/f681437203baa7671de3174b0fa583c349d9d5e1/src/transformers/models/t5/modeling_t5.py#L527
+    arange_tensor = ((attention_mask.cumsum(dim=-1) - 1) * attention_mask)[:, None, :]
+    alibi = slopes[..., None] * arange_tensor
+    return alibi.reshape(batch_size * num_heads, 1, seq_length).to(dtype)
+
+
+def dropout_add(x: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
+    """
+    Dropout add function
+
+    Args:
+        x (`torch.tensor`, *required*):
+            input tensor
+        residual (`torch.tensor`, *required*):
+            esidual tensor
+        prob (`float`, *required*):
+            dropout probability
+        training (`bool`, *required*):
+            training mode
+    """
+    out = F.dropout(x, p=prob, training=training)
+    out = residual + out
+    return out
+
+
+def bloom_gelu_forward(x: torch.Tensor) -> torch.Tensor:
+    """
+    Custom bias GELU function. Adapted from Megatron-DeepSpeed code. Here we use a simple implementation (inference) to
+    make the model jitable.
+
+    Args:
+        x (`torch.tensor`, *required*):
+            input hidden states
+    """
+    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+
+def bloom_gelu_back(g: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    """
+    gradient of tanh approximation of gelu gradient of actual gelu is: 0.5 * (1. + torch.erf(x * 0.70710678)) +
+    0.3989423 * x * torch.exp(-0.5 * x * x)
+
+    Args:
+        g (`torch.tensor`, *required*):
+            gradient output tensor
+        x (`torch.tensor`, *required*):
+            input tensor
+    """
+    x = x[0]  # x is a tuple of 1 element, needs to unpack it first
+    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
+    return ff * g
+
+
+class GeLUFunction(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, input: torch.Tensor) -> torch.Tensor:
+        ctx.save_for_backward(input)
+        return bloom_gelu_forward(input)
+
+    @staticmethod
+    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:
+        input = ctx.saved_tensors
+        tmp = bloom_gelu_back(grad_output, input)
+        return tmp
+
+
+class BloomGelu(nn.Module):
+    """
+    BloomBiasGelu wrapper function that make use of the simple function on inference mode to make the model
+    torchscriptable and use the autograd function in training mode to get the accurate results of the gradients Partly
+    copied from Megatron-DeepSpeed code and adapted for our needs
+
+    See here why autograd functions are not torchscriptable: https://github.com/pytorch/pytorch/issues/22329
+    """
+
+    def __init__(self):
+        super().__init__()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.training:
+            return GeLUFunction.apply(x)
+        else:
+            return bloom_gelu_forward(x)
+
+
+class LlamaRotaryEmbedding(nn.Module):
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):
+        super().__init__()
+
+        self.dim = dim
+        self.max_position_embeddings = max_position_embeddings
+        self.base = base
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        # Build here to make `torch.jit.trace` work.
+        #  self._set_cos_sin_cache(
+        #      seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.float32
+        #  )
+        self.max_seq_len_cached = max_position_embeddings
+        self.cos_cached = None
+        self.sin_cached = None
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :].to(dtype)
+        self.sin_cached = emb.sin()[None, None, :, :].to(dtype)
+
+    def forward(self, x, seq_len=None):
+        # x: [bs, num_attention_heads, seq_len, head_size]
+        if self.cos_cached is None:
+            self._set_cos_sin_cache(seq_len=self.max_seq_len_cached, device=x.device, dtype=torch.float32)
+        if seq_len > self.max_seq_len_cached:
+            print('reset cos/sin cache')
+            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=torch.float32)
+
+        self.cos_cached = self.cos_cached.to(x.device)
+        self.sin_cached = self.sin_cached.to(x.device)
+        return (
+            self.cos_cached[:, :, :seq_len, ...],
+            self.sin_cached[:, :, :seq_len, ...],
+        )
+
+
+class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):
+    """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+        t = t / self.scaling_factor
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :].to(dtype)
+        self.sin_cached = emb.sin()[None, None, :, :].to(dtype)
+
+
+class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):
+    """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""
+
+    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):
+        self.scaling_factor = scaling_factor
+        super().__init__(dim, max_position_embeddings, base, device)
+
+    def _set_cos_sin_cache(self, seq_len, device, dtype):
+        self.max_seq_len_cached = seq_len
+
+        if seq_len > self.max_position_embeddings:
+            base = self.base * (
+                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)
+            ) ** (self.dim / (self.dim - 2))
+            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))
+            self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)
+
+        freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+        # Different from paper, but it uses a different permutation in order to obtain the same calculation
+        emb = torch.cat((freqs, freqs), dim=-1)
+        self.cos_cached = emb.cos()[None, None, :, :].to(dtype)
+        self.sin_cached = emb.sin()[None, None, :, :].to(dtype)
+
+
+def rotate_half(x):
+    """Rotates half the hidden dims of the input."""
+    x1 = x[..., : x.shape[-1] // 2]
+    x2 = x[..., x.shape[-1] // 2 :]
+    return torch.cat((-x2, x1), dim=-1)
+
+
+def apply_rotary_pos_emb(q, k, cos, sin, position_ids):
+    '''
+    q: [bs, seq_len, num_heads, head_dim]
+    k: [bs, seq_len, num_heads, head_dim]
+
+    cos: [1, 1, seq_len, dim]
+    sin: [1, 1, seq_len, dim]
+
+    position_ids: [bs, seq_len]
+    '''
+    # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.
+    cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]
+    sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]
+
+    cos = cos[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
+    sin = sin[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]
+
+    # q_embed and k_embed: [bs, seq_len, num_heads, dim]
+    q_embed = (q.float() * cos) + (rotate_half(q.float()) * sin)
+    k_embed = (k.float() * cos) + (rotate_half(k.float()) * sin)
+    return q_embed.to(q.dtype), k_embed.to(k.dtype)
+
+
+class BloomAttention(nn.Module):
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__()
+
+        self.config = config
+        self.pretraining_tp = config.pretraining_tp
+        self.slow_but_exact = config.slow_but_exact
+
+        self.hidden_size = config.hidden_size
+        self.num_heads = config.n_head
+        self.head_dim = self.hidden_size // self.num_heads
+        self.split_size = self.hidden_size
+        self.hidden_dropout = config.hidden_dropout
+        self.p_attn_dropout = config.attention_dropout
+        self.max_position_embeddings = config.seq_length
+
+        if self.head_dim * self.num_heads != self.hidden_size:
+            raise ValueError(
+                f"`hidden_size` must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and `num_heads`:"
+                f" {self.num_heads})."
+            )
+
+        # Layer-wise attention scaling
+        self.inv_norm_factor = 1.0 / math.sqrt(self.head_dim)
+        self.beta = 1.0
+
+        self.query_key_value = nn.Linear(self.hidden_size, 3 * self.hidden_size, bias=True)
+        self.dense = nn.Linear(self.hidden_size, self.hidden_size)
+        self.attention_dropout = nn.Dropout(config.attention_dropout)
+
+        self._init_rope()
+
+
+    def _init_rope(self):
+        if self.config.rope_scaling is None:
+            self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)
+        else:
+            scaling_type = self.config.rope_scaling["type"]
+            scaling_factor = self.config.rope_scaling["factor"]
+            if scaling_type == "linear":
+                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(
+                    self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor
+                )
+            elif scaling_type == "dynamic":
+                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(
+                    self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor
+                )
+            else:
+                raise ValueError(f"Unknown RoPE scaling type {scaling_type}")
+
+
+    def _split_heads(self, fused_qkv: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+        """
+        Split the last dimension into (num_heads, head_dim) without making any copies, results share same memory
+        storage as `fused_qkv`
+
+        Args:
+            fused_qkv (`torch.tensor`, *required*): [batch_size, seq_length, num_heads * 3 * head_dim]
+
+        Returns:
+            query: [batch_size, seq_length, num_heads, head_dim] key: [batch_size, seq_length, num_heads, head_dim]
+            value: [batch_size, seq_length, num_heads, head_dim]
+        """
+        batch_size, seq_length, three_times_hidden_size = fused_qkv.shape
+        fused_qkv = fused_qkv.view(batch_size, seq_length, self.num_heads, 3, self.head_dim)
+        return fused_qkv[..., 0, :], fused_qkv[..., 1, :], fused_qkv[..., 2, :]
+
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor,
+        #  alibi: torch.Tensor,
+        attention_mask: torch.Tensor,
+        causal_mask: torch.Tensor,
+        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+    ):
+        fused_qkv = self.query_key_value(hidden_states)  # [batch_size, seq_length, 3 x hidden_size]
+
+        # 3 x [batch_size, seq_length, num_heads, head_dim]
+        (query_layer, key_layer, value_layer) = self._split_heads(fused_qkv)
+        batch_size, q_length, num_heads, head_dim = query_layer.shape
+
+        kv_length = key_layer.shape[1]
+        position_ids = attention_mask.cumsum(dim=-1) - 1
+        if layer_past is not None:
+            kv_length += layer_past[0].shape[1]
+            position_ids = position_ids[:, -1:]
+
+        ## add rope
+        cos, sin = self.rotary_emb(value_layer, seq_len=kv_length)
+        query_layer, key_layer = apply_rotary_pos_emb(query_layer, key_layer, cos, sin, position_ids)
+        # still [bs, seq_len, num_heads, head_dim]
+
+        if layer_past is not None:
+            past_key, past_value = layer_past
+            # concatenate along seq_length dimension:
+            #  - key: [batch_size, kv_length, num_heads, head_dim]
+            #  - value: [batch_size, kv_length, num_heads, head_dim]
+            key_layer = torch.cat((past_key, key_layer), dim=1)
+            value_layer = torch.cat((past_value, value_layer), dim=1)
+
+        if use_cache is True:
+            present = (key_layer, value_layer)
+        else:
+            present = None
+
+        context_layer = self.compute_qkv_attn(query_layer, key_layer,
+                value_layer, causal_mask, head_mask=head_mask,
+                causal=layer_past is None)
+
+        # aggregate results across tp ranks. See here: https://github.com/pytorch/pytorch/issues/76232
+        if self.pretraining_tp > 1 and self.slow_but_exact:
+            slices = self.hidden_size / self.pretraining_tp
+            output_tensor = torch.zeros_like(context_layer)
+            for i in range(self.pretraining_tp):
+                output_tensor = output_tensor + F.linear(
+                    context_layer[:, :, int(i * slices) : int((i + 1) * slices)],
+                    self.dense.weight[:, int(i * slices) : int((i + 1) * slices)],
+                )
+        else:
+            output_tensor = self.dense(context_layer)
+
+        output_tensor = dropout_add(output_tensor, residual, self.hidden_dropout, self.training)
+
+        outputs = (output_tensor, present)
+        if output_attentions:
+            outputs += (attention_probs,)
+
+        return outputs
+
+    def compute_qkv_attn(self, query_layer, key_layer, value_layer, causal_mask, head_mask=None, causal=False):
+        batch_size, q_length, num_heads, head_dim = query_layer.shape
+        kv_length = key_layer.size(1)
+
+        query_layer = query_layer.transpose(1, 2)
+        key_layer = key_layer.permute(0, 2, 3, 1)
+        attn_weights = torch.matmul(query_layer, key_layer).mul(self.inv_norm_factor)
+        # attn_weights is [batch_size, num_heads, q_length, kv_length]
+
+        # cast attention scores to fp32, compute scaled softmax and cast back to initial dtype - [batch_size, num_heads, q_length, kv_length]
+        # `float16` has a minimum value of -65504.0, whereas `bfloat16` and `float32` have a minimum value of `-3.4e+38`
+        input_dtype = attn_weights.dtype
+        if input_dtype == torch.float16:
+            attn_weights = attn_weights.to(torch.float)
+
+        attn_weights = torch.masked_fill(attn_weights, causal_mask,
+                torch.finfo(attn_weights.dtype).min)
+        attention_probs = F.softmax(attn_weights, dim=-1, dtype=torch.float32).to(input_dtype)
+        # [batch_size, num_heads, q_length, kv_length]
+
+        attention_probs = self.attention_dropout(attention_probs)
+
+        if head_mask is not None:
+            attention_probs = attention_probs * head_mask
+
+        # change v view [bs, num_heads, kv_length, head_dim]
+        value_layer = value_layer.transpose(1, 2)
+
+        # matmul: [batch_size, num_heads, q_length, head_dim]
+        num = batch_size * num_heads
+        attention_probs = attention_probs.reshape(num, q_length, kv_length)
+        value_layer = value_layer.reshape(num, kv_length, head_dim)
+        context_layer = torch.bmm(attention_probs, value_layer)
+        context_layer = context_layer.reshape(batch_size, num_heads, q_length, head_dim)
+        context_layer = torch.einsum('bhld->blhd', context_layer)
+        context_layer = context_layer.flatten(2)
+        return context_layer
+
+
+class BloomAttentionFlashAttn(BloomAttention):
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__(config)
+
+    def compute_qkv_attn(self, query_layer, key_layer, value_layer, causal_mask, head_mask=None, causal=False):
+        batch_size, q_length, num_heads, head_dim = query_layer.shape
+
+        # flash attention requires qkv to be [bs, seq_len, num_heads, head_dim]
+        context_layer = flash_attn_func(query_layer, key_layer,
+                value_layer, dropout_p=self.p_attn_dropout,
+                softmax_scale=self.inv_norm_factor, causal=causal)
+        # output is: [batch_size, q_length, num_heads, head_dim]
+
+        context_layer = context_layer.flatten(2)
+        return context_layer
+
+
+class BloomAttentionTorchFast(BloomAttention):
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__(config)
+
+    def compute_qkv_attn(self, query_layer, key_layer, value_layer, causal_mask, head_mask=None, causal=False):
+        batch_size, q_length, num_heads, head_dim = query_layer.shape
+        kv_length = key_layer.size(1)
+
+        ## requires qkv to have shape: [batch_size, num_heads, length, head_dim]
+        #  query_layer = torch.einsum('blhd->bhld', query_layer)
+        #  key_layer   = torch.einsum('blhd->bhld', key_layer)
+        #  value_layer = torch.einsum('blhd->bhld', value_layer)
+        query_layer = query_layer.transpose(1, 2)
+        key_layer   = key_layer.transpose(1, 2)
+        value_layer = value_layer.transpose(1, 2)
+        attn_mask = torch.zeros(batch_size, 1, q_length, kv_length,
+                device=query_layer.device, dtype=query_layer.dtype)
+        attn_mask.masked_fill_(causal_mask, torch.finfo(attn_mask.dtype).min)
+        with torch.backends.cuda.sdp_kernel(enable_flash=True, enable_math=True, enable_mem_efficient=True):
+            context_layer = F.scaled_dot_product_attention(query_layer, key_layer,
+                    value_layer, dropout_p=self.p_attn_dropout, attn_mask=attn_mask,
+                    scale=self.inv_norm_factor)
+        context_layer = context_layer.transpose(1, 2)
+        #  context_layer = torch.einsum('blhd->bhld', context_layer)
+
+        context_layer = context_layer.flatten(2)
+        return context_layer
+
+
+
+class BloomMLP(nn.Module):
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__()
+        hidden_size = config.hidden_size
+
+        self.pretraining_tp = config.pretraining_tp
+        self.slow_but_exact = config.slow_but_exact
+        self.dense_h_to_4h = nn.Linear(hidden_size, 4 * hidden_size)
+        self.gelu_impl = BloomGelu()
+        self.dense_4h_to_h = nn.Linear(4 * hidden_size, hidden_size)
+        self.hidden_dropout = config.hidden_dropout
+
+    def forward(self, hidden_states: torch.Tensor, residual: torch.Tensor) -> torch.Tensor:
+        hidden_states = self.gelu_impl(self.dense_h_to_4h(hidden_states))
+
+        if self.pretraining_tp > 1 and self.slow_but_exact:
+            intermediate_output = torch.zeros_like(residual)
+            slices = self.dense_4h_to_h.weight.shape[-1] / self.pretraining_tp
+            for i in range(self.pretraining_tp):
+                intermediate_output = intermediate_output + F.linear(
+                    hidden_states[:, :, int(i * slices) : int((i + 1) * slices)],
+                    self.dense_4h_to_h.weight[:, int(i * slices) : int((i + 1) * slices)],
+                )
+        else:
+            intermediate_output = self.dense_4h_to_h(hidden_states)
+
+        output = dropout_add(intermediate_output, residual, self.hidden_dropout, self.training)
+
+        return output
+
+
+class BloomBlock(nn.Module):
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__()
+        self.config = config
+        hidden_size = config.hidden_size
+
+        self.input_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+        self.num_heads = config.n_head
+
+        if config.use_flash_attn:
+            self.self_attention = BloomAttentionFlashAttn(config)
+        else:
+            self.self_attention = BloomAttention(config)
+            #  self.self_attention = BloomAttentionTorchFast(config)
+
+        self.post_attention_layernorm = LayerNorm(hidden_size, eps=config.layer_norm_epsilon)
+
+        self.mlp = BloomMLP(config)
+
+        self.apply_residual_connection_post_layernorm = config.apply_residual_connection_post_layernorm
+        self.hidden_dropout = config.hidden_dropout
+
+    #  @torch.compile
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: torch.Tensor,
+        causal_mask: torch.Tensor,
+        layer_past: Optional[Tuple[torch.Tensor, torch.Tensor]] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        use_cache: bool = False,
+        output_attentions: bool = False,
+    ):
+        # hidden_states: [batch_size, seq_length, hidden_size]
+
+        # Layer norm at the beginning of the transformer layer.
+        layernorm_output = self.input_layernorm(hidden_states)
+
+        # Layer norm post the self attention.
+        if self.apply_residual_connection_post_layernorm:
+            residual = layernorm_output
+        else:
+            residual = hidden_states
+
+        # Self attention.
+        attn_outputs = self.self_attention(
+            layernorm_output,
+            residual,
+            layer_past=layer_past,
+            attention_mask=attention_mask,
+            causal_mask=causal_mask,
+            #  alibi=alibi,
+            head_mask=head_mask,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+        )
+
+        attention_output = attn_outputs[0]
+
+        outputs = attn_outputs[1:]
+
+        layernorm_output = self.post_attention_layernorm(attention_output)
+
+        # Get residual
+        if self.apply_residual_connection_post_layernorm:
+            residual = layernorm_output
+        else:
+            residual = attention_output
+
+        # MLP.
+        output = self.mlp(layernorm_output, residual)
+
+        if use_cache:
+            outputs = (output,) + outputs
+        else:
+            outputs = (output,) + outputs[1:]
+
+        return outputs  # hidden_states, present, attentions
+
+
+class NormHead(nn.Module):
+
+    def __init__(self, hidden_size, vocab_size, bias=False):
+        super().__init__()
+        self.weight = nn.Parameter(torch.empty((vocab_size, hidden_size)))
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        self.first_flag = True
+
+    def forward(self, hidden_states):
+        if self.training:
+            norm_weight = nn.functional.normalize(self.weight)
+        elif self.first_flag:
+            self.first_flag = False
+            self.weight = nn.Parameter(nn.functional.normalize(self.weight))
+            norm_weight = self.weight
+        else:
+            norm_weight = self.weight
+        return F.linear(hidden_states, norm_weight)
+
+
+class SmoothEmbedding(nn.Embedding):
+
+    def __init__(self, vocab_size, embed_dim):
+        super().__init__(vocab_size, embed_dim)
+
+    def forward(self, input_ids):
+        word_emb = super().forward(input_ids)
+        mean_emb = self.weight.mean(dim=0, keepdim=True)
+        emb = word_emb * 0.9 + mean_emb * 0.1
+        return emb
+
+
+class BloomPreTrainedModel(PreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]
+    """
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
+    """
+
+    config_class = CCBloomConfig
+    base_model_prefix = "transformer"
+    supports_gradient_checkpointing = True
+    _no_split_modules = ["BloomBlock"]
+    _skip_keys_device_placement = "past_key_values"
+
+    def __init__(self, *inputs, **kwargs):
+        super().__init__(*inputs, **kwargs)
+
+    def _init_weights(self, module: nn.Module):
+        """Initialize the weights."""
+        if isinstance(module, nn.Linear):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.bias is not None:
+                module.bias.data.zero_()
+        elif isinstance(module, nn.Embedding):
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+            if module.padding_idx is not None:
+                module.weight.data[module.padding_idx].zero_()
+        elif isinstance(module, LayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+
+    def _set_gradient_checkpointing(self, module: nn.Module, value: bool = False):
+        if isinstance(module, BloomModel):
+            module.gradient_checkpointing = value
+
+    @staticmethod
+    def _convert_to_standard_cache(
+        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]], batch_size: int
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Standardizes the format of the cache so as to match most implementations, i.e. to tuple(tuple([batch_size,
+        num_heads, ...]))
+        """
+        #  seq_length, num_heads, head_dim = past_key_value[0][0].shape
+        #  batch_size_times_num_heads, seq_length, head_dim = past_key_value[0][0].shape
+        #  num_heads = batch_size_times_num_heads // batch_size
+        # key: [batch_size, seq_length, num_heads, head_dim] -> [batch_size, num_heads, head_dim, seq_length]
+        # value: [batch_size, seq_length, num_heads, head_dim] -> [batch_size, num_heads, seq_length, head_dim]
+        return tuple(
+            (
+                layer_past[0].permute(0, 2, 3, 1),
+                layer_past[1].permute(0, 2, 1, 3),
+            )
+            for layer_past in past_key_value
+        )
+
+    @staticmethod
+    def _convert_to_bloom_cache(
+        past_key_value: Tuple[Tuple[torch.Tensor, torch.Tensor]]
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor]]:
+        """
+        Converts the cache to the format expected by Bloom, i.e. to tuple(tuple([batch_size * num_heads, ...]))
+        """
+        # key:  [batch_size, num_heads, head_dim, seq_length] -> [batch_size, seq_length, num_heads, head_dim]
+        # value: [batch_size, num_heads, seq_length, head_dim] -> [batch_size, seq_length, num_heads, head_dim]
+        return tuple(
+            (
+                torch.einsum('bhdl->blhd', layer_past[0]),
+                torch.einsum('bhld->blhd', layer_past[1])
+            )
+            for layer_past in past_key_value
+        )
+
+
+BLOOM_START_DOCSTRING = r"""
+
+    This model inherits from [`PreTrainedModel`]. Check the superclass documentation for the generic methods the
+    library implements for all its model (such as downloading or saving, resizing the input embeddings etc.)
+
+    This model is also a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage
+    and behavior.
+
+    Parameters:
+        config ([`CCBloomConfig`]): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights.
+"""
+
+BLOOM_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (`torch.LongTensor` of shape `(batch_size, input_ids_length)`):
+            `input_ids_length` = `sequence_length` if `past_key_values` is `None` else `past_key_values[0][0].shape[2]`
+            (`sequence_length` of input past key value states). Indices of input sequence tokens in the vocabulary.
+
+            If `past_key_values` is used, only `input_ids` that do not have their past calculated should be passed as
+            `input_ids`.
+
+            Indices can be obtained using [`AutoTokenizer`]. See [`PreTrainedTokenizer.encode`] and
+            [`PreTrainedTokenizer.__call__`] for details.
+
+            [What are input IDs?](../glossary#input-ids)
+        past_key_values (`Tuple[Tuple[torch.Tensor]]` of length `config.n_layers`):
+            Contains precomputed hidden-states (key and values in the attention blocks) as computed by the model (see
+            `past_key_values` output below). Can be used to speed up sequential decoding. The `input_ids` which have
+            their past given to this model should not be passed as `input_ids` as they have already been computed.
+
+            Each element of `past_key_values` is a tuple (past_key, past_value):
+            - past_key: [batch_size * num_heads, head_dim, kv_length]
+            - past_value: [batch_size * num_heads, kv_length, head_dim]
+        attention_mask (`torch.FloatTensor` of shape `(batch_size, sequence_length)`, *optional*):
+            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            [What are attention masks?](../glossary#attention-mask)
+        head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        inputs_embeds (`torch.FloatTensor` of shape `(batch_size, sequence_length, hidden_size)`, *optional*):
+            Optionally, instead of passing `input_ids` you can choose to directly pass an embedded representation. This
+            is useful if you want more control over how to convert `input_ids` indices into associated vectors than the
+            model's internal embedding lookup matrix.
+
+            If `past_key_values` is used, optionally only the last `inputs_embeds` have to be input (see
+            `past_key_values`).
+        use_cache (`bool`, *optional*):
+            If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding (see
+            `past_key_values`).
+        output_attentions (`bool`, *optional*):
+            Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned
+            tensors for more detail.
+        output_hidden_states (`bool`, *optional*):
+            Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for
+            more detail.
+        return_dict (`bool`, *optional*):
+            Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare Bloom Model transformer outputting raw hidden-states without any specific head on top.",
+    BLOOM_START_DOCSTRING,
+)
+class BloomModel(BloomPreTrainedModel):
+    _auto_class = 'ccbloom'
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__(config)
+
+        self.embed_dim = config.hidden_size
+        self.num_heads = config.n_head
+
+        # Embedding + LN Embedding
+        self.word_embeddings = SmoothEmbedding(config.vocab_size, self.embed_dim)
+        #  self.word_embeddings = nn.Embedding(config.vocab_size, self.embed_dim)
+        self.word_embeddings_layernorm = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+
+        # Transformer blocks
+        self.h = nn.ModuleList([BloomBlock(config) for _ in range(config.num_hidden_layers)])
+
+        # Final Layer Norm
+        self.ln_f = LayerNorm(self.embed_dim, eps=config.layer_norm_epsilon)
+
+        self.gradient_checkpointing = False
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def build_alibi_tensor(self, attention_mask: torch.Tensor, num_heads: int, dtype: torch.dtype) -> torch.Tensor:
+        return build_alibi_tensor(attention_mask, num_heads, dtype)
+
+    def get_input_embeddings(self):
+        return self.word_embeddings
+
+    def _prepare_attn_mask(
+        self, attention_mask: torch.Tensor, input_shape: Tuple[int, int], past_key_values_length: int
+    ) -> torch.BoolTensor:
+        # create causal mask
+        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
+        combined_attention_mask = None
+        device = attention_mask.device
+        _, src_length = input_shape
+
+        if src_length > 1:
+            combined_attention_mask = _make_causal_mask(
+                input_shape, device=device, past_key_values_length=past_key_values_length
+            )
+
+        # [batch_size, seq_length] -> [batch_size, 1, tgt_length, src_length]
+        expanded_attn_mask = _expand_mask(attention_mask, tgt_length=src_length)
+        combined_attention_mask = (
+            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask | combined_attention_mask
+        )
+
+        return combined_attention_mask
+
+    def set_input_embeddings(self, new_embeddings: torch.Tensor):
+        self.word_embeddings = new_embeddings
+
+    @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=BaseModelOutputWithPastAndCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.LongTensor] = None,
+        inputs_embeds: 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,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor, ...], BaseModelOutputWithPastAndCrossAttentions]:
+        if deprecated_arguments.pop("position_ids", False) is not False:
+            # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
+            warnings.warn(
+                "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
+                " passing `position_ids`.",
+                FutureWarning,
+            )
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+
+        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
+
+        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:
+            batch_size, seq_length = input_ids.shape
+        elif inputs_embeds is not None:
+            batch_size, seq_length, _ = inputs_embeds.shape
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        if past_key_values is None:
+            past_key_values = tuple([None] * len(self.h))
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape batch_size x num_heads x N x N
+        # head_mask has shape n_layer x batch x num_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.word_embeddings(input_ids)
+
+        hidden_states = self.word_embeddings_layernorm(inputs_embeds)
+
+        presents = () if use_cache else None
+        all_self_attentions = () if output_attentions else None
+        all_hidden_states = () if output_hidden_states else None
+
+        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
+
+        # Compute alibi tensor: check build_alibi_tensor documentation
+        seq_length_with_past = seq_length
+        past_key_values_length = 0
+        if past_key_values[0] is not None:
+            past_key_values_length = past_key_values[0][0].shape[2]
+            seq_length_with_past = seq_length_with_past + past_key_values_length
+        if attention_mask is None:
+            attention_mask = torch.ones((batch_size, seq_length_with_past), device=hidden_states.device)
+        else:
+            attention_mask = attention_mask.to(hidden_states.device)
+
+        ## lower triangular mask
+        causal_mask = None
+        if self.config.use_flash_attn == False:
+        #  if True:
+            causal_mask = self._prepare_attn_mask(
+                attention_mask,
+                input_shape=(batch_size, seq_length),
+                past_key_values_length=past_key_values_length,
+            )
+
+        for i, (block, layer_past) in enumerate(zip(self.h, past_key_values)):
+            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=use_cache, output_attentions=output_attentions)
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    attention_mask,
+                    causal_mask,
+                    layer_past,
+                    head_mask[i],
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=attention_mask,
+                    causal_mask=causal_mask,
+                    head_mask=head_mask[i],
+                    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],)
+
+        # Add last hidden state
+        hidden_states = self.ln_f(hidden_states)
+
+        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] 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,
+        )
+
+
+@add_start_docstrings(
+    """
+    The Bloom Model transformer with a language modeling head on top (linear layer with weights tied to the input
+    embeddings).
+    """,
+    BLOOM_START_DOCSTRING,
+)
+class CCBloomForCausalLM(BloomPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__(config)
+        self.transformer = BloomModel(config)
+        #  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+        self.lm_head = NormHead(config.hidden_size, config.vocab_size)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    def get_output_embeddings(self):
+        return self.lm_head
+
+    def set_output_embeddings(self, new_embeddings: torch.Tensor):
+        self.lm_head = new_embeddings
+
+    def prepare_inputs_for_generation(
+        self,
+        input_ids: torch.LongTensor,
+        past_key_values: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> dict:
+        # only last token for input_ids if past is not None
+        if past_key_values:
+            input_ids = input_ids[:, -1].unsqueeze(-1)
+
+            # the cache may be in the stardard format (e.g. in contrastive search), convert to bloom's format if needed
+            #  if past_key_values[0][0].shape[0] == input_ids.shape[0]:
+            #      past_key_values = self._convert_to_bloom_cache(past_key_values)
+
+        # if `inputs_embeds` are passed, we only want to use them in the 1st generation step
+        if inputs_embeds is not None and past_key_values is None:
+            model_inputs = {"inputs_embeds": inputs_embeds}
+        else:
+            model_inputs = {"input_ids": input_ids}
+
+        model_inputs.update(
+            {
+                "past_key_values": past_key_values,
+                "use_cache": kwargs.get("use_cache"),
+                "attention_mask": attention_mask,
+            }
+        )
+        return model_inputs
+
+    @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=CausalLMOutputWithCrossAttentions,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor], 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]`
+        """
+        if deprecated_arguments.pop("position_ids", False) is not False:
+            # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
+            warnings.warn(
+                "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
+                " passing `position_ids`.",
+                FutureWarning,
+            )
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+
+        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,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        hidden_states = transformer_outputs[0]
+
+        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()
+            batch_size, seq_length, vocab_size = shift_logits.shape
+            # Flatten the tokens
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(
+                shift_logits.view(batch_size * seq_length, vocab_size), shift_labels.view(batch_size * seq_length)
+            )
+
+        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,
+        )
+
+    def _reorder_cache(
+        self, past: Tuple[Tuple[torch.Tensor, torch.Tensor], ...], beam_idx: torch.LongTensor
+    ) -> Tuple[Tuple[torch.Tensor, torch.Tensor], ...]:
+        """
+        This function is used to re-order the `past_key_values` cache if [`~PreTrainedModel.beam_search`] or
+        [`~PreTrainedModel.beam_sample`] is called. This is required to match `past_key_values` with the correct
+        beam_idx at every generation step.
+
+        Output shares the same memory storage as `past`.
+        """
+        standardized_past = self._convert_to_standard_cache(past, batch_size=len(beam_idx))
+
+        # Get a copy of `beam_idx` on all the devices where we need those indices.
+        device_to_beam_idx = {
+            past_state.device: beam_idx.to(past_state.device) for layer_past in past for past_state in layer_past
+        }
+        reordered_past = tuple(
+            (
+                layer_past[0].index_select(0, device_to_beam_idx[layer_past[0].device]),
+                layer_past[1].index_select(0, device_to_beam_idx[layer_past[0].device]),
+            )
+            for layer_past in standardized_past
+        )
+        return self._convert_to_bloom_cache(reordered_past)
+
+
+@add_start_docstrings(
+    """
+    The Bloom Model transformer with a sequence classification head on top (linear layer).
+
+    [`BloomForSequenceClassification`] uses the last token in order to do the classification, as other causal models
+    (e.g. GPT-1) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If
+    no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the
+    padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in
+    each row of the batch).
+    """,
+    BLOOM_START_DOCSTRING,
+)
+class BloomForSequenceClassification(BloomPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.transformer = BloomModel(config)
+        self.score = nn.Linear(config.hidden_size, config.num_labels, bias=False)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=SequenceClassifierOutputWithPast,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor], SequenceClassifierOutputWithPast]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        if deprecated_arguments.pop("position_ids", False) is not False:
+            # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
+            warnings.warn(
+                "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
+                " passing `position_ids`.",
+                FutureWarning,
+            )
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+
+        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,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = transformer_outputs[0]
+        logits = self.score(hidden_states)
+
+        if input_ids is not None:
+            batch_size = input_ids.shape[0]
+        else:
+            batch_size = inputs_embeds.shape[0]
+
+        if self.config.pad_token_id is None and batch_size != 1:
+            raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if input_ids is not None:
+                sequence_lengths = (torch.ne(input_ids, self.config.pad_token_id).sum(-1) - 1).to(logits.device)
+            else:
+                sequence_lengths = -1
+                logger.warning(
+                    f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                    "unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+                )
+
+        pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]
+
+        loss = None
+        if labels is not None:
+            if self.config.problem_type is None:
+                if self.num_labels == 1:
+                    self.config.problem_type = "regression"
+                elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):
+                    self.config.problem_type = "single_label_classification"
+                else:
+                    self.config.problem_type = "multi_label_classification"
+
+            if self.config.problem_type == "regression":
+                loss_fct = MSELoss()
+                if self.num_labels == 1:
+                    loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())
+                else:
+                    loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "single_label_classification":
+                loss_fct = CrossEntropyLoss()
+                loss = loss_fct(pooled_logits, labels)
+            elif self.config.problem_type == "multi_label_classification":
+                loss_fct = BCEWithLogitsLoss()
+                loss = loss_fct(pooled_logits, labels)
+        if not return_dict:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return SequenceClassifierOutputWithPast(
+            loss=loss,
+            logits=pooled_logits,
+            past_key_values=transformer_outputs.past_key_values,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    Bloom Model with a token classification head on top (a linear layer on top of the hidden-states output) e.g. for
+    Named-Entity-Recognition (NER) tasks.
+    """,
+    BLOOM_START_DOCSTRING,
+)
+class BloomForTokenClassification(BloomPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]
+
+    def __init__(self, config: CCBloomConfig):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.transformer = BloomModel(config)
+        if hasattr(config, "classifier_dropout") and config.classifier_dropout is not None:
+            classifier_dropout = config.classifier_dropout
+        elif hasattr(config, "hidden_dropout") and config.hidden_dropout is not None:
+            classifier_dropout = config.hidden_dropout
+        else:
+            classifier_dropout = 0.1
+        self.dropout = nn.Dropout(classifier_dropout)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        checkpoint=_CHECKPOINT_FOR_DOC,
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        past_key_values: Optional[Tuple[Tuple[torch.Tensor, torch.Tensor], ...]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        head_mask: Optional[torch.Tensor] = None,
+        inputs_embeds: Optional[torch.Tensor] = None,
+        labels: Optional[torch.Tensor] = None,
+        use_cache: Optional[bool] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+        **deprecated_arguments,
+    ) -> Union[Tuple[torch.Tensor], TokenClassifierOutput]:
+        r"""
+        labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,
+            config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If
+            `config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        """
+        if deprecated_arguments.pop("position_ids", False) is not False:
+            # `position_ids` could have been `torch.Tensor` or `None` so defaulting pop to `False` allows to detect if users were passing explicitly `None`
+            warnings.warn(
+                "`position_ids` have no functionality in BLOOM and will be removed in v5.0.0. You can safely ignore"
+                " passing `position_ids`.",
+                FutureWarning,
+            )
+        if len(deprecated_arguments) > 0:
+            raise ValueError(f"Got unexpected arguments: {deprecated_arguments}")
+
+        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,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        hidden_states = transformer_outputs[0]
+        hidden_states = self.dropout(hidden_states)
+        logits = self.classifier(hidden_states)
+
+        loss = None
+        if labels is not None:
+            # move labels to correct device to enable model parallelism
+            labels = labels.to(logits.device)
+            batch_size, seq_length = labels.shape
+            loss_fct = CrossEntropyLoss()
+            loss = loss_fct(
+                logits.view(batch_size * seq_length, self.num_labels), labels.view(batch_size * seq_length)
+            )
+
+        if not return_dict:
+            output = (logits,) + transformer_outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """
+    The BLOOM Model transformer with a span classification head on top for extractive question-answering tasks like
+    SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`).
+    """,
+    BLOOM_START_DOCSTRING,
+)
+class BloomForQuestionAnswering(BloomPreTrainedModel):
+    _keys_to_ignore_on_load_missing = [r"h.*.self_attention.scale_mask_softmax.causal_mask", r"lm_head.weight"]
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.transformer = BloomModel(config)
+        self.qa_outputs = nn.Linear(config.hidden_size, 2)
+
+        # Initialize weights and apply final processing
+        self.post_init()
+
+    @add_start_docstrings_to_model_forward(BLOOM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        position_ids: Optional[torch.LongTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        start_positions: Optional[torch.LongTensor] = None,
+        end_positions: Optional[torch.LongTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, QuestionAnsweringModelOutput]:
+        r"""
+        start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
+            are not taken into account for computing the loss.
+        end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`). Position outside of the sequence
+            are not taken into account for computing the loss.
+        """
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.transformer(
+            input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1).contiguous()
+        end_logits = end_logits.squeeze(-1).contiguous()
+
+        total_loss = None
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions = start_positions.clamp(0, ignored_index)
+            end_positions = end_positions.clamp(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+
+        if not return_dict:
+            output = (start_logits, end_logits) + outputs[2:]
+            return ((total_loss,) + output) if total_loss is not None else output
+
+        return QuestionAnsweringModelOutput(
+            loss=total_loss,
+            start_logits=start_logits,
+            end_logits=end_logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+

pytorch_model-00001-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:df5ba779f2b1d0ffcc19fda8d8a958a430743a14eae2c0fc3fa198dd5bc23693
+size 927086401

pytorch_model-00002-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ed09b6a4a5ae5b1917ee05f5f7484581a98b92b4f5ac2d1167734a4b140e8a50
+size 939796331

pytorch_model-00003-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7d027fb4930f1f7d22682211889de321a164d6092bc5d30d8e7b8c15cf1cdd7
+size 939778659

pytorch_model-00004-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:29058b10814361ae89c78eabced2d7e9a1abfca3c2fe0c77f17e29ac48b3533b
+size 939771079

pytorch_model-00005-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fffda31283253b879a348fbb094110873a0a32d01983f1368f691df827b1732d
+size 939796331

pytorch_model-00006-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a0a732472370f3025055b988379d9c289c0d1ffc3d28c4106e1e92ba265c9482
+size 939778659

pytorch_model-00007-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ff1b12fe6e14db5aa2680298df1fe76dac02247c7590c12f5e1a0d02c7e59105
+size 939771079

pytorch_model-00008-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2b7fdd228fdd86bc2ddc61d6671d1ea187749f5a7ec0012f2938a9219ed0a6aa
+size 939796331

pytorch_model-00009-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:64fc0567bb455d1763bf8262876c785db3a0eab04a51115feae4bccf84b10bfc
+size 939778659

pytorch_model-00010-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:488679b37dad5e9416e87129689944b89c796809b7a2e5ab93e3beba8e8c58f9
+size 939771079

pytorch_model-00011-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:44a3ad467d6b8c8b43d48fb6ce8a9855d671c8739376fd6fd6be2253b320d806
+size 939796331

pytorch_model-00012-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7be7478ffbc5e6f6075dc2eea80a29a34335f9158036506fdebd37ae67fd5946
+size 939778659

pytorch_model-00013-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c261788ad6d5003622fac486b098b3e55ea46c686a7d3d28c226dd686c437a8
+size 939771079

pytorch_model-00014-of-00014.bin

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:6f37f159283ca191a4af43fcc60c0909c6f09c8334073c45b4cd587c4ca986c3
+size 927052345

pytorch_model.bin.index.json

@@ -0,0 +1,373 @@
+{
+  "metadata": {
+    "total_size": 13131399168
+  },
+  "weight_map": {
+    "lm_head.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.0.input_layernorm.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.input_layernorm.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.mlp.dense_4h_to_h.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.mlp.dense_4h_to_h.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.mlp.dense_h_to_4h.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.mlp.dense_h_to_4h.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.post_attention_layernorm.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.post_attention_layernorm.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.self_attention.dense.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.self_attention.dense.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.self_attention.query_key_value.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.0.self_attention.query_key_value.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.1.input_layernorm.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.1.input_layernorm.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.h.1.mlp.dense_4h_to_h.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.mlp.dense_h_to_4h.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.post_attention_layernorm.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.post_attention_layernorm.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.self_attention.dense.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.self_attention.dense.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.self_attention.query_key_value.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.1.self_attention.query_key_value.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.10.input_layernorm.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.input_layernorm.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.mlp.dense_4h_to_h.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.10.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.10.mlp.dense_h_to_4h.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.10.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.10.post_attention_layernorm.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.post_attention_layernorm.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.self_attention.dense.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.self_attention.dense.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.self_attention.query_key_value.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.10.self_attention.query_key_value.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.11.input_layernorm.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.input_layernorm.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.mlp.dense_4h_to_h.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.mlp.dense_4h_to_h.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.mlp.dense_h_to_4h.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.post_attention_layernorm.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.post_attention_layernorm.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.self_attention.dense.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.self_attention.dense.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.self_attention.query_key_value.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.11.self_attention.query_key_value.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.input_layernorm.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.input_layernorm.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.mlp.dense_4h_to_h.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.12.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.12.mlp.dense_h_to_4h.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.mlp.dense_h_to_4h.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.post_attention_layernorm.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.post_attention_layernorm.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.self_attention.dense.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.self_attention.dense.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.self_attention.query_key_value.bias": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.12.self_attention.query_key_value.weight": "pytorch_model-00006-of-00014.bin",
+    "transformer.h.13.input_layernorm.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.input_layernorm.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.mlp.dense_4h_to_h.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.mlp.dense_h_to_4h.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.post_attention_layernorm.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.post_attention_layernorm.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.self_attention.dense.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.self_attention.dense.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.self_attention.query_key_value.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.13.self_attention.query_key_value.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.input_layernorm.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.input_layernorm.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.mlp.dense_4h_to_h.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.mlp.dense_4h_to_h.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.mlp.dense_h_to_4h.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.mlp.dense_h_to_4h.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.post_attention_layernorm.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.post_attention_layernorm.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.self_attention.dense.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.self_attention.dense.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.self_attention.query_key_value.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.14.self_attention.query_key_value.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.15.input_layernorm.bias": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.15.input_layernorm.weight": "pytorch_model-00007-of-00014.bin",
+    "transformer.h.15.mlp.dense_4h_to_h.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.mlp.dense_h_to_4h.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.post_attention_layernorm.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.post_attention_layernorm.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.self_attention.dense.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.self_attention.dense.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.self_attention.query_key_value.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.15.self_attention.query_key_value.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.input_layernorm.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.input_layernorm.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.mlp.dense_4h_to_h.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.mlp.dense_4h_to_h.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.mlp.dense_h_to_4h.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.mlp.dense_h_to_4h.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.post_attention_layernorm.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.post_attention_layernorm.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.self_attention.dense.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.self_attention.dense.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.self_attention.query_key_value.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.16.self_attention.query_key_value.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.input_layernorm.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.input_layernorm.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.mlp.dense_4h_to_h.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.17.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.17.mlp.dense_h_to_4h.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.17.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.17.post_attention_layernorm.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.post_attention_layernorm.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.self_attention.dense.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.self_attention.dense.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.self_attention.query_key_value.bias": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.17.self_attention.query_key_value.weight": "pytorch_model-00008-of-00014.bin",
+    "transformer.h.18.input_layernorm.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.input_layernorm.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.mlp.dense_4h_to_h.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.mlp.dense_4h_to_h.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.mlp.dense_h_to_4h.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.post_attention_layernorm.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.post_attention_layernorm.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.self_attention.dense.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.self_attention.dense.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.self_attention.query_key_value.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.18.self_attention.query_key_value.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.input_layernorm.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.input_layernorm.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.mlp.dense_4h_to_h.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.19.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.19.mlp.dense_h_to_4h.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.mlp.dense_h_to_4h.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.post_attention_layernorm.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.post_attention_layernorm.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.self_attention.dense.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.self_attention.dense.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.self_attention.query_key_value.bias": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.19.self_attention.query_key_value.weight": "pytorch_model-00009-of-00014.bin",
+    "transformer.h.2.input_layernorm.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.input_layernorm.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.mlp.dense_4h_to_h.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.mlp.dense_4h_to_h.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.mlp.dense_h_to_4h.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.mlp.dense_h_to_4h.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.post_attention_layernorm.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.post_attention_layernorm.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.self_attention.dense.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.self_attention.dense.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.self_attention.query_key_value.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.2.self_attention.query_key_value.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.20.input_layernorm.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.input_layernorm.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.mlp.dense_4h_to_h.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.mlp.dense_h_to_4h.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.mlp.dense_h_to_4h.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.post_attention_layernorm.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.post_attention_layernorm.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.self_attention.dense.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.self_attention.dense.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.self_attention.query_key_value.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.20.self_attention.query_key_value.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.input_layernorm.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.input_layernorm.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.mlp.dense_4h_to_h.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.mlp.dense_4h_to_h.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.mlp.dense_h_to_4h.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.mlp.dense_h_to_4h.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.post_attention_layernorm.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.post_attention_layernorm.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.self_attention.dense.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.self_attention.dense.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.self_attention.query_key_value.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.21.self_attention.query_key_value.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.22.input_layernorm.bias": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.22.input_layernorm.weight": "pytorch_model-00010-of-00014.bin",
+    "transformer.h.22.mlp.dense_4h_to_h.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.mlp.dense_4h_to_h.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.mlp.dense_h_to_4h.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.mlp.dense_h_to_4h.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.post_attention_layernorm.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.post_attention_layernorm.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.self_attention.dense.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.self_attention.dense.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.self_attention.query_key_value.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.22.self_attention.query_key_value.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.input_layernorm.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.input_layernorm.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.mlp.dense_4h_to_h.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.mlp.dense_4h_to_h.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.mlp.dense_h_to_4h.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.mlp.dense_h_to_4h.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.post_attention_layernorm.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.post_attention_layernorm.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.self_attention.dense.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.self_attention.dense.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.self_attention.query_key_value.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.23.self_attention.query_key_value.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.input_layernorm.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.input_layernorm.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.mlp.dense_4h_to_h.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.24.mlp.dense_4h_to_h.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.24.mlp.dense_h_to_4h.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.24.mlp.dense_h_to_4h.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.24.post_attention_layernorm.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.post_attention_layernorm.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.self_attention.dense.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.self_attention.dense.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.self_attention.query_key_value.bias": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.24.self_attention.query_key_value.weight": "pytorch_model-00011-of-00014.bin",
+    "transformer.h.25.input_layernorm.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.input_layernorm.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.mlp.dense_4h_to_h.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.mlp.dense_4h_to_h.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.mlp.dense_h_to_4h.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.mlp.dense_h_to_4h.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.post_attention_layernorm.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.post_attention_layernorm.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.self_attention.dense.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.self_attention.dense.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.self_attention.query_key_value.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.25.self_attention.query_key_value.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.input_layernorm.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.input_layernorm.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.mlp.dense_4h_to_h.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.26.mlp.dense_4h_to_h.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.26.mlp.dense_h_to_4h.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.mlp.dense_h_to_4h.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.post_attention_layernorm.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.post_attention_layernorm.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.self_attention.dense.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.self_attention.dense.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.self_attention.query_key_value.bias": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.26.self_attention.query_key_value.weight": "pytorch_model-00012-of-00014.bin",
+    "transformer.h.27.input_layernorm.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.input_layernorm.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.mlp.dense_4h_to_h.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.mlp.dense_4h_to_h.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.mlp.dense_h_to_4h.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.mlp.dense_h_to_4h.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.post_attention_layernorm.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.post_attention_layernorm.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.self_attention.dense.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.self_attention.dense.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.self_attention.query_key_value.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.27.self_attention.query_key_value.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.input_layernorm.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.input_layernorm.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.mlp.dense_4h_to_h.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.mlp.dense_4h_to_h.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.mlp.dense_h_to_4h.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.mlp.dense_h_to_4h.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.post_attention_layernorm.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.post_attention_layernorm.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.self_attention.dense.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.self_attention.dense.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.self_attention.query_key_value.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.28.self_attention.query_key_value.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.29.input_layernorm.bias": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.29.input_layernorm.weight": "pytorch_model-00013-of-00014.bin",
+    "transformer.h.29.mlp.dense_4h_to_h.bias": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.mlp.dense_4h_to_h.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.mlp.dense_h_to_4h.bias": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.mlp.dense_h_to_4h.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.post_attention_layernorm.bias": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.post_attention_layernorm.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.self_attention.dense.bias": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.self_attention.dense.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.self_attention.query_key_value.bias": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.29.self_attention.query_key_value.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.h.3.input_layernorm.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.input_layernorm.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.mlp.dense_4h_to_h.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.3.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.3.mlp.dense_h_to_4h.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.3.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.3.post_attention_layernorm.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.post_attention_layernorm.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.self_attention.dense.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.self_attention.dense.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.self_attention.query_key_value.bias": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.3.self_attention.query_key_value.weight": "pytorch_model-00002-of-00014.bin",
+    "transformer.h.4.input_layernorm.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.input_layernorm.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.mlp.dense_4h_to_h.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.mlp.dense_4h_to_h.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.mlp.dense_h_to_4h.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.post_attention_layernorm.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.post_attention_layernorm.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.self_attention.dense.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.self_attention.dense.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.self_attention.query_key_value.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.4.self_attention.query_key_value.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.input_layernorm.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.input_layernorm.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.mlp.dense_4h_to_h.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.5.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.5.mlp.dense_h_to_4h.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.mlp.dense_h_to_4h.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.post_attention_layernorm.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.post_attention_layernorm.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.self_attention.dense.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.self_attention.dense.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.self_attention.query_key_value.bias": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.5.self_attention.query_key_value.weight": "pytorch_model-00003-of-00014.bin",
+    "transformer.h.6.input_layernorm.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.input_layernorm.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.mlp.dense_4h_to_h.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.mlp.dense_h_to_4h.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.post_attention_layernorm.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.post_attention_layernorm.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.self_attention.dense.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.self_attention.dense.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.self_attention.query_key_value.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.6.self_attention.query_key_value.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.input_layernorm.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.input_layernorm.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.mlp.dense_4h_to_h.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.mlp.dense_4h_to_h.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.mlp.dense_h_to_4h.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.mlp.dense_h_to_4h.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.post_attention_layernorm.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.post_attention_layernorm.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.self_attention.dense.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.self_attention.dense.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.self_attention.query_key_value.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.7.self_attention.query_key_value.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.8.input_layernorm.bias": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.8.input_layernorm.weight": "pytorch_model-00004-of-00014.bin",
+    "transformer.h.8.mlp.dense_4h_to_h.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.mlp.dense_h_to_4h.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.post_attention_layernorm.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.post_attention_layernorm.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.self_attention.dense.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.self_attention.dense.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.self_attention.query_key_value.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.8.self_attention.query_key_value.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.input_layernorm.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.input_layernorm.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.mlp.dense_4h_to_h.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.mlp.dense_4h_to_h.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.mlp.dense_h_to_4h.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.mlp.dense_h_to_4h.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.post_attention_layernorm.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.post_attention_layernorm.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.self_attention.dense.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.self_attention.dense.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.self_attention.query_key_value.bias": "pytorch_model-00005-of-00014.bin",
+    "transformer.h.9.self_attention.query_key_value.weight": "pytorch_model-00005-of-00014.bin",
+    "transformer.ln_f.bias": "pytorch_model-00014-of-00014.bin",
+    "transformer.ln_f.weight": "pytorch_model-00014-of-00014.bin",
+    "transformer.word_embeddings.weight": "pytorch_model-00001-of-00014.bin",
+    "transformer.word_embeddings_layernorm.bias": "pytorch_model-00001-of-00014.bin",
+    "transformer.word_embeddings_layernorm.weight": "pytorch_model-00001-of-00014.bin"
+  }
+}

special_tokens_map.json

@@ -0,0 +1,30 @@
+{
+  "bos_token": {
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "eos_token": {
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "pad_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "unk_token": {
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}

tokenizer.model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4be54af290d93c113bcbf421115ae9eed9d6340408f564898f1e966dc738ef01
+size 1136699

tokenizer_config.json

@@ -0,0 +1,42 @@
+{
+  "add_bos_token": false,
+  "add_eos_token": false,
+  "bos_token": {
+    "__type": "AddedToken",
+    "content": "<s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "clean_up_tokenization_spaces": false,
+  "eos_token": {
+    "__type": "AddedToken",
+    "content": "</s>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "legacy": true,
+  "model_max_length": 1000000000000000019884624838656,
+  "pad_token": {
+    "__type": "AddedToken",
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  },
+  "padding_side": "left",
+  "sp_model_kwargs": {},
+  "tokenizer_class": "LlamaTokenizer",
+  "unk_token": {
+    "__type": "AddedToken",
+    "content": "<unk>",
+    "lstrip": false,
+    "normalized": true,
+    "rstrip": false,
+    "single_word": false
+  }
+}