easydel.modules.mistral.__init__

class easydel.modules.mistral.__init__.MistralConfig(vocab_size: int = 32000, hidden_size: int = 4096, intermediate_size: int = 14336, head_dim: int = 128, num_hidden_layers: int = 32, num_attention_heads: int = 32, num_key_value_heads: int = 8, hidden_act='silu', max_position_embeddings=131072, initializer_range=0.02, rms_norm_eps=1e-06, use_cache=True, pad_token_id=None, bos_token_id: int = 1, eos_token_id: int = 2, tie_word_embeddings=False, rope_theta=10000.0, rope_scaling: Dict[str, Union[str, float]] = None, sliding_window=4096, gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, number_rep_kv: int = 1, attention_dropout: float = 0.0, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, bits: Optional[int] = None, attention_bias: bool = False, **kwargs)

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 32000) – Vocabulary size of the Mistral 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 4096) – Dimensionality of the encoder layers and the pooler layer.

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

• head_dim (int, defaults to 128) – Dimensionality of the head for attention.

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

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

• num_key_value_heads (int, optional, defaults to 8) – Number of key and value heads for each attention layer in the Transformer encoder.

• 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 4096 * 32) – 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.

• rms_norm_eps (float, optional, defaults to 1e-6) – The epsilon used by the rms 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) – The index of the padding token in the vocabulary.

• bos_token_id (int, optional, defaults to 1) – The id of the beginning-of-sequence token.

• eos_token_id (int, optional, defaults to 2) – The id of the end-of-sequence token.

• tie_word_embeddings (bool, optional, defaults to False) – 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.

• rope_scaling (tp.Dict[str, tp.Union[str, float]], optional) – The configuration for rope scaling.

• sliding_window (int, optional, defaults to 4096) – The sliding window size.

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

• number_rep_kv (int, optional, defaults to 1) – Number of repetitions for the key and value vectors.

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

• use_scan_mlp (bool, optional, defaults to False) – Whether to use the scan implementation for the MLP.

• scan_mlp_chunk_size (int, optional, defaults to 1024) – The chunk size to use when scanning the MLP.

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

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

attach_custom_arguments(gradient_checkpointing: EasyDeLGradientCheckPointers = EasyDeLGradientCheckPointers.NONE, use_scan_mlp: bool = False, scan_mlp_chunk_size: int = 1024, number_rep_kv: int = 1, bits: Optional[int] = None, attention_dropout: float = 0.0, rope_scaling: Dict[str, Union[str, float]] = None, attention_bias: bool = False, **kwargs)

The attach_custom_arguments function adds the following arguments to the model:

Parameters

• self – Bind the attributes and methods of a class to an instance of that class

• gradient_checkpointing – str: Determine whether to use gradient checkpointing

• use_scan_mlp – bool: Determine whether to use the scan_mlp function or notn

• scan_mlp_chunk_size – int: Chunk the input to the mlp

• number_rep_kv – int: Control the number of times that the key and value vectors are repeated

• bits – tp.Optional[int]: Specify the number of bits to use for quantization

• attention_dropout – float: Set the dropout rate for the attention layer

• attention_bias – bool: when ever to use attention_bias

• rope_scaling – tp.Dict[str, tp.Union[str, float]]: rope_scaling for rope

Return type

A tuple of the following

get_partition_rules(*args, **kwargs)

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

static get_weight_decay_exclusions()

property granted_freq_max_position_embedding: int

property granted_mask_max_position_embedding: int

model_type: str = 'mistral'

static rng_keys()

class easydel.modules.mistral.__init__.MistralForCausalLM(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Mistral model with a language modeling head for causal language modeling tasks.

This model is a transformer-based language model with sliding window attention applied to perform autoregressive language generation.

config

Configuration for the model.

Type

MistralConfig

dtype

Data type for computations.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.

class easydel.modules.mistral.__init__.MistralForSequenceClassification(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Mistral model for sequence classification tasks.

This class extends the base Mistral model by adding a linear classification head to perform sequence classification tasks such as sentiment analysis or text classification.

config

Configuration for the model.

Type

MistralConfig

dtype

Data type for computations.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.

class easydel.modules.mistral.__init__.MistralModel(*args: Any, **kwargs: Any)

Bases: EasyDeLBaseModule

Mistral model implementation.

This implements the Mistral language model architecture, utilizing transformer blocks with RMSNorm, sliding window attention, and rotary position embeddings.

config

Configuration for the model.

Type

MistralConfig

dtype

Data type for computations.

Type

jnp.dtype

param_dtype

Data type for parameters.

Type

jnp.dtype

precision

Precision setting for JAX operations.