easydel.modules.cohere.cohere_configuration#

class easydel.modules.cohere.cohere_configuration.CohereConfig(vocab_size=256000, hidden_size=8192, intermediate_size=22528, logit_scale=0.0625, num_hidden_layers=40, num_attention_heads=64, num_key_value_heads=None, hidden_act='silu', max_position_embeddings=8192, initializer_range=0.02, layer_norm_eps=1e-05, use_cache=True, pad_token_id=0, bos_token_id=5, eos_token_id=255001, tie_word_embeddings=True, rope_theta=10000.0, attention_bias=False, attention_dropout=0.0, use_qk_norm: bool = False, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, **kwargs)[source]#

Bases: EasyDeLBaseConfig

Configuration objects inherit from [EasyDeLBaseConfig] and can be used to control the model outputs. Read the documentation from [EasyDeLBaseConfig] for more information.

Parameters
  • vocab_size (int, optional, defaults to 256000) – Vocabulary size of the Cohere model. Defines the number of different tokens that can be represented by the inputs_ids passed to the forward method.

  • hidden_size (int, optional, defaults to 8192) – Dimensionality of the encoder layers and the pooler layer.

  • intermediate_size (int, optional, defaults to 22528) – Dimensionality of the “intermediate” (i.e., feed-forward) layer in the Transformer encoder.

  • logit_scale (float, optional, defaults to 0.0625) – A logit scale value used in the attention layer.

  • num_hidden_layers (int, optional, defaults to 40) – Number of hidden layers in the Transformer encoder.

  • num_attention_heads (int, optional, defaults to 64) – Number of attention heads for each attention layer in the Transformer encoder.

  • num_key_value_heads (int, optional) – Number of key and value heads for each attention layer in the Transformer encoder. Will default to num_attention_heads if not set.

  • hidden_act (str or function, optional, defaults to “silu”) – The non-linear activation function (function or string) to use in the encoder and pooler. If string, “gelu”, “relu”, “swish” and “gelu_new” are supported.

  • max_position_embeddings (int, optional, defaults to 8192) – The maximum sequence length that this model might ever be used with. Typically set this to something large just in case (e.g., 2048 or 4096).

  • initializer_range (float, optional, defaults to 0.02) – The standard deviation of the truncated_normal_initializer for initializing all weight matrices.

  • layer_norm_eps (float, optional, defaults to 1e-5) – The epsilon used by the layer normalization layers.

  • use_cache (bool, optional, defaults to True) – Whether or not the model should return the last key/values attentions (not used by all models). Only relevant if config.is_decoder=True.

  • pad_token_id (int, optional, defaults to 0) – The index of the padding token in the vocabulary.

  • bos_token_id (int, optional, defaults to 5) – The index of the beginning of sequence token in the vocabulary.

  • eos_token_id (int, optional, defaults to 255001) – The index of the end of sequence token in the vocabulary.

  • tie_word_embeddings (bool, optional, defaults to True) – Whether to tie the weights of the input embeddings and the output embeddings.

  • rope_theta (float, optional, defaults to 10000.0) – The theta value to use for rotary position embeddings.

  • attention_bias (bool, optional, defaults to False) – Whether to use attention bias.

  • attention_dropout (float, optional, defaults to 0.0) – The dropout ratio for the attention probabilities.

  • use_qk_norm (bool, optional, defaults to False) – Whether to use query and key normalization.

  • gradient_checkpointing (str, optional, defaults to “nothing_saveable”) – The gradient checkpointing configuration.

  • bits (int, optional) – The number of bits to quantize the model to.

attach_custom_arguments(gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, bits: Optional[int] = None, **kwargs)[source]#

The attach_custom_arguments function adds the following arguments to the Transformer class:

Parameters
  • self – Refer to the current object

  • gradient_checkpointing (EasyDeLGradientCheckPointers) – Control the amount of memory used by jax

  • bits (Optional[int]) – Determine the number of bits used in the quantization

  • **kwargs – Additional keyword arguments.

get_partition_rules(*args, **kwargs)[source]#

Get the partition rules for the model. :returns: The partition rules. :rtype: tp.Tuple[tp.Tuple[str, PartitionSpec]]

static get_weight_decay_exclusions()[source]#

Returns a tuple of parameter names for which weight decay should be excluded.

property granted_freq_max_position_embedding: int#

Returns the maximum position embedding size for frequency-based position embeddings, falling back to max_position_embeddings.

property granted_mask_max_position_embedding: int#

Returns the maximum position embedding size for mask-based position embeddings, falling back to max_position_embeddings.

model_type: str = 'cohere'#
static rng_keys()[source]#

Returns the names of the random number generator keys used by the model.