easydel.modules.pixtral.__init__

easydel.modules.pixtral.__init__#

class easydel.modules.pixtral.__init__.PixtralVisionConfig(hidden_size: int = 1024, intermediate_size: int = 4096, num_hidden_layers: int = 24, num_attention_heads: int = 16, num_channels: int = 3, image_size: int = 1024, patch_size: int = 16, hidden_act: str = 'gelu', attention_dropout: float = 0.0, rope_theta: float = 10000.0, initializer_range: int = 0.02, **kwargs)[source]#

Bases: EasyDeLBaseConfig

This is the configuration class to store the configuration of a [PixtralVisionModel]. It is used to instantiate an Pixtral vision encoder according to the specified arguments, defining the model architecture. Instantiating a configuration with the defaults will yield a similar configuration to the vision encoder used by Pixtral-12B.

e.g. [pixtral-hf/pixtral-9b](https://huggingface.co/pixtral-hf/pixtral-9b)

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

Parameters

  • hidden_size (int, optional, defaults to 1024) – Dimension of the hidden representations.

  • intermediate_size (int, optional, defaults to 4096) – Dimension of the MLP representations.

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

  • num_attention_heads (int, optional, defaults to 16) – Number of attention heads in the Transformer encoder.

  • num_channels (int, optional, defaults to 3) – Number of input channels in the input images.

  • image_size (int, optional, defaults to 1024) – Max dimension of the input images.

  • patch_size (int, optional, defaults to 16) – Size of the image patches.

  • hidden_act (str, optional, defaults to “gelu”) – Activation function used in the hidden layers.

  • attention_dropout (float, optional, defaults to 0.0) – Dropout probability for the attention layers.

  • rope_theta (float, optional, defaults to 10000.0) – The base period of the RoPE embeddings.

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

Example:

```python >>> from transformers import PixtralVisionModel, PixtralVisionConfig

>>> # Initializing a Pixtral-12B style configuration
>>> config = PixtralVisionConfig()

>>> # Initializing a model (with randomly initialized weights) from the configuration
>>> model = PixtralVisionModel(configuration)

>>> # Accessing the model configuration
>>> configuration = model.config
```
get_partition_rules(*args, **kwargs)[source]#

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

model_type: str = 'pixtral'#
class easydel.modules.pixtral.__init__.PixtralVisionModel(*args: Any, **kwargs: Any)[source]#

Bases: EasyDeLBaseModule

The Pixtral Vision Model transformer.

This class implements the complete Pixtral vision model, including patch embedding via convolution and the main transformer stack.

config#

Configuration object for the model.

Type

PixtralVisionConfig

dtype#

Data type for computations.

Type

jnp.dtype

param_dtype#

Data type for parameters.

Type

jnp.dtype

precision#

Precision setting for JAX operations.

Type

jax.lax.PrecisionLike

rngs#

Random number generators.

Type

nn.Rngs

patch_conv#

Convolutional layer for patch embedding.

Type

nn.Conv

transformer#

The main transformer stack.

Type

PixtralTransformer

ln_pre#

Layer normalization applied before the transformer blocks.

Type

RMSNorm

property frequencies#

Cached property to compute and retrieve RoPE frequencies.

Contents