codingutf8 2022 the big science workshop and huggingface inc team 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 orglicenseslicense2 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 typechecking any list mapping optional from packaging import version if typechecking from import pretrainedtokenizer tensortype from configurationutils import pretrainedconfig from onnx import onnxconfigwithpast patchingspec from utils import istorchavailable logging logger logging getloggername bloompretrainedconfigarchivemap bigsciencebloom https huggingface cobigsciencebloomresolvemainconfig json bigsciencebloom560m https huggingface cobigsciencebloom560mblobmainconfig json bigsciencebloom1b1 https huggingface cobigsciencebloom1b1blobmainconfig json bigsciencebloom1b7 https huggingface cobigsciencebloom1b7blobmainconfig json bigsciencebloom3b https huggingface cobigsciencebloom3bblobmainconfig json bigsciencebloom7b1 https huggingface cobigsciencebloom7b1blobmainconfig json class bloomconfigpretrainedconfig backward compatibility with nembed kwarg todo how to do that better bloom stores values on dynamic axis 2 for more details see https github comhuggingfacetransformerspull18344 we need to order the input in the way they appears in the forward need to add the pastkeys not using the same length for pastkeyvalues coding utf 8 2022 the big science workshop and huggingface inc team 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 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 example python from transformers import bloomconfig bloommodel initializing a bloom configuration configuration bloomconfig initializing a model with random weights from the configuration model bloommodel configuration accessing the model configuration configuration model config tp rank used when training with megatron backward compatibility with n_embed kwarg todo how to do that better bloom stores values on dynamic axis 2 for more details see https github com huggingface transformers pull 18344 we need to order the input in the way they appears in the forward need to add the past_keys not using the same length for past_key_values
from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional

from packaging import version


if TYPE_CHECKING:
    from ... import PreTrainedTokenizer, TensorType

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfigWithPast, PatchingSpec
from ...utils import is_torch_available, logging


logger = logging.get_logger(__name__)

BLOOM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    "bigscience/bloom": "https://huggingface.co/bigscience/bloom/resolve/main/config.json",
    "bigscience/bloom-560m": "https://huggingface.co/bigscience/bloom-560m/blob/main/config.json",
    "bigscience/bloom-1b1": "https://huggingface.co/bigscience/bloom-1b1/blob/main/config.json",
    "bigscience/bloom-1b7": "https://huggingface.co/bigscience/bloom-1b7/blob/main/config.json",
    "bigscience/bloom-3b": "https://huggingface.co/bigscience/bloom-3b/blob/main/config.json",
    "bigscience/bloom-7b1": "https://huggingface.co/bigscience/bloom-7b1/blob/main/config.json",
}


class BloomConfig(PretrainedConfig):
    

    model_type = "bloom"
    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,  
        slow_but_exact=False,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        
        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)


class BloomOnnxConfig(OnnxConfigWithPast):
    torch_onnx_minimum_version = version.parse("1.12")

    def __init__(
        self,
        config: PretrainedConfig,
        task: str = "default",
        patching_specs: List[PatchingSpec] = None,
        use_past: bool = False,
    ):
        super().__init__(config, task=task, patching_specs=patching_specs, use_past=use_past)
        if not getattr(self._config, "pad_token_id", None):
            
            self._config.pad_token_id = 0

    @property
    def inputs(self) -> Mapping[str, Mapping[int, str]]:
        common_inputs = OrderedDict({"input_ids": {0: "batch", 1: "sequence"}})
        if self.use_past:
            
            self.fill_with_past_key_values_(common_inputs, direction="inputs", inverted_values_shape=True)
            common_inputs["attention_mask"] = {0: "batch", 1: "past_sequence + sequence"}
        else:
            common_inputs["attention_mask"] = {0: "batch", 1: "sequence"}

        return common_inputs

    @property
    def num_layers(self) -> int:
        return self._config.n_layer

    @property
    def num_attention_heads(self) -> int:
        return self._config.n_head

    @property
    def atol_for_validation(self) -> float:
        return 1e-3

    def generate_dummy_inputs(
        self,
        tokenizer: "PreTrainedTokenizer",
        batch_size: int = -1,
        seq_length: int = -1,
        is_pair: bool = False,
        framework: Optional["TensorType"] = None,
    ) -> Mapping[str, Any]:
        common_inputs = super(OnnxConfigWithPast, self).generate_dummy_inputs(
            tokenizer, batch_size=batch_size, seq_length=seq_length, is_pair=is_pair, framework=framework
        )

        
        ordered_inputs = OrderedDict({"input_ids": common_inputs["input_ids"]})

        
        if self.use_past:
            if not is_torch_available():
                raise ValueError("Cannot generate dummy past_keys inputs without PyTorch installed.")
            else:
                import torch

                batch, seqlen = common_inputs["input_ids"].shape
                
                past_key_values_length = seqlen + 2
                head_dim = self._config.hidden_size // self.num_attention_heads
                past_key_shape = (
                    batch * self.num_attention_heads,
                    head_dim,
                    past_key_values_length,
                )
                past_value_shape = (
                    batch * self.num_attention_heads,
                    past_key_values_length,
                    head_dim,
                )
                ordered_inputs["past_key_values"] = [
                    (torch.zeros(past_key_shape), torch.zeros(past_value_shape)) for _ in range(self.num_layers)
                ]

        ordered_inputs["attention_mask"] = common_inputs["attention_mask"]
        if self.use_past:
            mask_dtype = ordered_inputs["attention_mask"].dtype
            ordered_inputs["attention_mask"] = torch.cat(
                [ordered_inputs["attention_mask"], torch.ones(batch, past_key_values_length, dtype=mask_dtype)], dim=1
            )

        return ordered_inputs

    @property
    def default_onnx_opset(self) -> int:
        return 13