codingutf8 2022 google and huggingface inc 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 switch transformers model configuration from configurationutils import pretrainedconfig from utils import logging logger logging getloggername switchtransformerspretrainedconfigarchivemap googleswitchbase8 https huggingface cogoogleswitchbase8blobmainconfig json class switchtransformersconfigpretrainedconfig r this is the configuration class to store the configuration of a switchtransformersmodel it is used to instantiate a switchtransformers model according to the specified arguments defining the model architecture instantiating a configuration with the defaults will yield a similar configuration to that of the switchtransformers googleswitchbase8https huggingface cogoogleswitchbase8 architecture configuration objects inherit from pretrainedconfig and can be used to control the model outputs read the documentation from pretrainedconfig for more information arguments vocabsize int optional defaults to 32128 vocabulary size of the switchtransformers model defines the number of different tokens that can be represented by the inputsids passed when calling switchtransformersmodel dmodel int optional defaults to 512 size of the encoder layers and the pooler layer dkv int optional defaults to 64 size of the key query value projections per attention head dkv has to be equal to dmodel numheads dff int optional defaults to 2048 size of the intermediate feed forward layer in each switchtransformersblock expertcapacity int optional defaults to 64 number of tokens that can be stored in each expert if set to 1 the model will behave like a regular transformer numlayers int optional defaults to 12 number of dense hidden layers in the transformer encoder layer numsparseencoderlayers int optional defaults to 6 number of sparse moe dense hidden layers in the transformer encoder layer numdecoderlayers int optional defaults to 12 number of hidden layers in the transformer decoder will use the same value as numlayers if not set numsparsedecoderlayers int optional defaults to 12 number of sparse moe dense hidden layers in the transformer decoder layer numheads int optional defaults to 8 number of attention heads for each attention layer in the transformer encoder numexperts int optional defaults to 8 number of experts for each switchtransformer layer routertype str optional defaults to tokensmasked router type choose between tokensmasked tokensscatter and expertsmasked routerbias bool optional defaults to true whether to add a bias to the router routerjitternoise float optional defaults to 0 1 amount of noise to add to the router routerdtype str optional default to float32 the dtype used for the routers it is preferable to keep the dtype to float32 as specified in the selective precision discussion in the paperhttps arxiv orgabs2101 03961 routerignorepaddingtokens bool optional defaults to false whether to ignore padding tokens when routing relativeattentionnumbuckets int optional defaults to 32 the number of buckets to use for each attention layer relativeattentionmaxdistance int optional defaults to 128 the maximum distance of the longer sequences for the bucket separation dropoutrate float optional defaults to 0 1 the ratio for all dropout layers layernormeps float optional defaults to 1e6 the epsilon used by the layer normalization layers routerzlosscoef float optional defaults to 0 001 the z loss factor for the total loss routerauxlosscoef float optional defaults to 0 001 the aux loss factor for the total loss initializerfactor float optional defaults to 1 a factor for initializing all weight matrices should be kept to 1 used internally for initialization testing feedforwardproj string optional defaults to relu type of feed forward layer to be used should be one of relu or gatedgelu switchtransformersv1 1 uses the gatedgelu feed forward projection original switchtransformers uses relu addrouterprobs bool optional defaults to false whether to output router probabilities to compute router auxiliary loss usecache bool optional defaults to true whether or not the model should return the last keyvalues attentions not used by all models this tells us each how many encoder layer we ll have to set a sparse layer this tells us each how many encoder layer we ll have to set a sparse layer coding utf 8 2022 google and huggingface inc 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 switch transformers model configuration this is the configuration class to store the configuration of a switchtransformersmodel it is used to instantiate a switchtransformers model according to the specified arguments defining the model architecture instantiating a configuration with the defaults will yield a similar configuration to that of the switchtransformers google switch base 8 https huggingface co google switch base 8 architecture configuration objects inherit from pretrainedconfig and can be used to control the model outputs read the documentation from pretrainedconfig for more information arguments vocab_size int optional defaults to 32128 vocabulary size of the switchtransformers model defines the number of different tokens that can be represented by the inputs_ids passed when calling switchtransformersmodel d_model int optional defaults to 512 size of the encoder layers and the pooler layer d_kv int optional defaults to 64 size of the key query value projections per attention head d_kv has to be equal to d_model num_heads d_ff int optional defaults to 2048 size of the intermediate feed forward layer in each switchtransformersblock expert_capacity int optional defaults to 64 number of tokens that can be stored in each expert if set to 1 the model will behave like a regular transformer num_layers int optional defaults to 12 number of dense hidden layers in the transformer encoder layer num_sparse_encoder_layers int optional defaults to 6 number of sparse moe dense hidden layers in the transformer encoder layer num_decoder_layers int optional defaults to 12 number of hidden layers in the transformer decoder will use the same value as num_layers if not set num_sparse_decoder_layers int optional defaults to 12 number of sparse moe dense hidden layers in the transformer decoder layer num_heads int optional defaults to 8 number of attention heads for each attention layer in the transformer encoder num_experts int optional defaults to 8 number of experts for each switchtransformer layer router_type str optional defaults to tokens_masked router type choose between tokens_masked tokens_scatter and experts_masked router_bias bool optional defaults to true whether to add a bias to the router router_jitter_noise float optional defaults to 0 1 amount of noise to add to the router router_dtype str optional default to float32 the dtype used for the routers it is preferable to keep the dtype to float32 as specified in the selective precision discussion in the paper https arxiv org abs 2101 03961 router_ignore_padding_tokens bool optional defaults to false whether to ignore padding tokens when routing relative_attention_num_buckets int optional defaults to 32 the number of buckets to use for each attention layer relative_attention_max_distance int optional defaults to 128 the maximum distance of the longer sequences for the bucket separation dropout_rate float optional defaults to 0 1 the ratio for all dropout layers layer_norm_eps float optional defaults to 1e 6 the epsilon used by the layer normalization layers router_z_loss_coef float optional defaults to 0 001 the z loss factor for the total loss router_aux_loss_coef float optional defaults to 0 001 the aux loss factor for the total loss initializer_factor float optional defaults to 1 a factor for initializing all weight matrices should be kept to 1 used internally for initialization testing feed_forward_proj string optional defaults to relu type of feed forward layer to be used should be one of relu or gated gelu switchtransformersv1 1 uses the gated gelu feed forward projection original switchtransformers uses relu add_router_probs bool optional defaults to false whether to output router probabilities to compute router auxiliary loss use_cache bool optional defaults to true whether or not the model should return the last key values attentions not used by all models default symmetry this tells us each how many encoder layer we ll have to set a sparse layer hack this will create 0 sparse layers this tells us each how many encoder layer we ll have to set a sparse layer hack this will create 0 sparse layers
from ...configuration_utils import PretrainedConfig
from ...utils import logging


logger = logging.get_logger(__name__)

SWITCH_TRANSFORMERS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
    "google/switch-base-8": "https://huggingface.co/google/switch-base-8/blob/main/config.json",
}


class SwitchTransformersConfig(PretrainedConfig):
    r

    model_type = "switch_transformers"
    keys_to_ignore_at_inference = ["past_key_values"]
    attribute_map = {"hidden_size": "d_model", "num_attention_heads": "num_heads", "num_hidden_layers": "num_layers"}

    def __init__(
        self,
        vocab_size=32128,
        d_model=768,
        d_kv=64,
        d_ff=2048,
        expert_capacity=64,
        num_layers=12,
        num_sparse_encoder_layers=3,
        num_decoder_layers=12,
        num_sparse_decoder_layers=3,
        num_heads=12,
        num_experts=8,
        router_bias=False,
        router_jitter_noise=0.01,
        router_dtype="float32",
        router_ignore_padding_tokens=False,
        relative_attention_num_buckets=32,
        relative_attention_max_distance=128,
        dropout_rate=0.1,
        layer_norm_epsilon=1e-6,
        router_z_loss_coef=0.001,
        router_aux_loss_coef=0.001,
        initializer_factor=1.0,
        dense_act_fn="relu",
        is_encoder_decoder=True,
        add_router_probs=False,
        use_cache=True,
        pad_token_id=0,
        eos_token_id=1,
        **kwargs,
    ):
        self.vocab_size = vocab_size
        self.d_model = d_model
        self.d_kv = d_kv
        self.d_ff = d_ff

        self.num_sparse_encoder_layers = num_sparse_encoder_layers

        self.num_layers = num_layers
        self.num_decoder_layers = (
            num_decoder_layers if num_decoder_layers is not None else self.num_layers
        )  
        self.num_sparse_decoder_layers = num_sparse_decoder_layers

        
        if self.num_sparse_encoder_layers > 0:
            self.encoder_sparse_step = self.num_layers // self.num_sparse_encoder_layers
        else:
            self.encoder_sparse_step = self.num_layers  

        
        if self.num_sparse_decoder_layers > 0:
            self.decoder_sparse_step = self.num_decoder_layers // self.num_sparse_decoder_layers
        else:
            self.decoder_sparse_step = self.num_decoder_layers  

        self.num_heads = num_heads
        self.num_experts = num_experts
        self.expert_capacity = expert_capacity
        self.router_bias = router_bias
        self.router_jitter_noise = router_jitter_noise
        if router_dtype not in ["float32", "float16", "bfloat16"]:
            raise ValueError(f"`router_dtype` must be one of 'float32', 'float16' or 'bfloat16', got {router_dtype}")
        self.router_dtype = router_dtype

        self.router_ignore_padding_tokens = router_ignore_padding_tokens
        self.relative_attention_num_buckets = relative_attention_num_buckets
        self.relative_attention_max_distance = relative_attention_max_distance

        self.dropout_rate = dropout_rate
        self.layer_norm_epsilon = layer_norm_epsilon
        self.initializer_factor = initializer_factor
        self.use_cache = use_cache
        self.add_router_probs = add_router_probs

        self.router_z_loss_coef = router_z_loss_coef
        self.router_aux_loss_coef = router_aux_loss_coef
        self.dense_act_fn = dense_act_fn

        super().__init__(
            pad_token_id=pad_token_id,
            eos_token_id=eos_token_id,
            is_encoder_decoder=is_encoder_decoder,
            **kwargs,
        )