modalities.nn package

Subpackages

• modalities.nn.model_initialization package
  • Submodules
  • modalities.nn.model_initialization.composed_initialization module
    • ComposedInitializationRoutines
      • ComposedInitializationRoutines.get_composed_model_initializer()
      • ComposedInitializationRoutines.get_model_initializer_wrapper()
    • ComposedModelInitializationConfig
      • ComposedModelInitializationConfig.hidden_dim
      • ComposedModelInitializationConfig.mean
      • ComposedModelInitializationConfig.model_config
      • ComposedModelInitializationConfig.model_type
      • ComposedModelInitializationConfig.num_layers
      • ComposedModelInitializationConfig.std
      • ComposedModelInitializationConfig.weight_init_type
    • ModelInitializerWrapper
      • ModelInitializerWrapper.initialize_in_place()
    • ModelInitializerWrapperConfig
      • ModelInitializerWrapperConfig.model_config
      • ModelInitializerWrapperConfig.model_initializers
  • modalities.nn.model_initialization.initialization_if module
    • ModelInitializationIF
      • ModelInitializationIF.initialize_in_place()
  • modalities.nn.model_initialization.initialization_routines module
    • InitializationRoutines
      • InitializationRoutines.get_plain_initialization()
      • InitializationRoutines.get_scaled_embed_initialization()
      • InitializationRoutines.get_scaled_initialization()
    • NamedParameterwiseNormalInitialization
      • NamedParameterwiseNormalInitialization.initialize_in_place()
    • PlainInitializationConfig
      • PlainInitializationConfig.check_std_and_hidden_dim()
      • PlainInitializationConfig.hidden_dim
      • PlainInitializationConfig.mean
      • PlainInitializationConfig.model_config
      • PlainInitializationConfig.parameter_name_regexes
      • PlainInitializationConfig.std
    • ScaledEmbedInitializationConfig
      • ScaledEmbedInitializationConfig.mean
      • ScaledEmbedInitializationConfig.model_config
      • ScaledEmbedInitializationConfig.parameter_name_regexes
    • ScaledInitializationConfig
      • ScaledInitializationConfig.mean
      • ScaledInitializationConfig.model_config
      • ScaledInitializationConfig.num_layers
      • ScaledInitializationConfig.parameter_name_regexes
      • ScaledInitializationConfig.std
  • modalities.nn.model_initialization.parameter_name_filters module
    • RegexFilter
      • RegexFilter.biases
      • RegexFilter.model_config
      • RegexFilter.weights
    • SupportWeightInitModels
      • SupportWeightInitModels.COCA
      • SupportWeightInitModels.GPT2
    • WeightInitTypes
      • WeightInitTypes.PLAIN
      • WeightInitTypes.SCALED
      • WeightInitTypes.SCALED_EMBED
  • Module contents

Submodules

modalities.nn.attention module

class modalities.nn.attention.AttentionConfig(**data)

Bases: BaseModel

Create a new model by parsing and validating input data from keyword arguments.

Raises [ValidationError][pydantic_core.ValidationError] if the input data cannot be validated to form a valid model.

self is explicitly positional-only to allow self as a field name.

Parameters:

attention_engine_type (AttentionEngineType)

attention_engine_type: AttentionEngineType

model_config: ClassVar[ConfigDict] = {}

Configuration for the model, should be a dictionary conforming to [ConfigDict][pydantic.config.ConfigDict].

class modalities.nn.attention.AttentionEngineType(value)

Bases: str, Enum

DEFAULT_ATTENTION = 'default_attention'

PYTORCH_FLASH_ATTENTION = 'pytorch_flash_attention'

class modalities.nn.attention.AttentionType(value)

Bases: str, Enum

CAUSAL_SELF_ATTENTION = 'causal_self_attention'

CROSS_ATTENTION = 'cross_attention'

NON_CAUSAL_SELF_ATTENTION = 'non_causal_self_attention'

class modalities.nn.attention.MultiHeadAttention(attention_config=None, attention_type=AttentionType.CAUSAL_SELF_ATTENTION, n_embd=768, n_head=8, bias=True, dropout=0.0, block_size=1024)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• attention_config (AttentionConfig)

• attention_type (AttentionType)

• n_embd (int)

• n_head (int)

• bias (bool)

• dropout (float)

• block_size (int)

forward(x, context=None)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

Parameters:

• x (Tensor)

• context (Tensor | None)

modalities.nn.mlp module

class modalities.nn.mlp.MLP(in_features, hidden_features=None, out_features=None, bias=True, dropout=0.0, act_fn=<class 'torch.nn.modules.activation.GELU'>)

Bases: Module

Initialize internal Module state, shared by both nn.Module and ScriptModule.

Parameters:

• in_features (int)

• hidden_features (int | None)

• out_features (int | None)

• bias (bool)

• dropout (float)

• act_fn (Callable[[], Module])

forward(x)

Define the computation performed at every call.

Should be overridden by all subclasses.

Note

Although the recipe for forward pass needs to be defined within this function, one should call the Module instance afterwards instead of this since the former takes care of running the registered hooks while the latter silently ignores them.

Return type:

Tensor

Parameters:

x (Tensor)

Module contents