easydel.infra.factory

class easydel.infra.factory.ConfigType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration defining types of configurations that can be registered.

MODULE_CONFIG

Represents standard module configuration classes.

MODULE_CONFIG = 'module-config'

class easydel.infra.factory.ModuleRegistration(module: type[easydel.infra.base_module.EasyDeLBaseModule], config: type[easydel.infra.base_config.EasyDeLBaseConfig], embedding_layer_names: Optional[List[str]] = None, layernorm_names: Optional[List[str]] = None)

Bases: object

A container class holding information about a registered EasyDeL module.

This class stores the module class itself, its corresponding configuration class, and optional metadata like the names of embedding and LayerNorm layers, which can be useful for parameter transformation or analysis.

module

The class of the registered EasyDeL module.

Type

type[EasyDeLBaseModule]

config

The configuration class associated with the module.

Type

type[EasyDeLBaseConfig]

embedding_layer_names

A list of names identifying embedding layers within the module structure. Defaults to None.

Type

tp.Optional[tp.List[str]]

layernorm_names

A list of names identifying Layer Normalization layers within the module structure. Defaults to None.

Type

tp.Optional[tp.List[str]]

config: type[easydel.infra.base_config.EasyDeLBaseConfig]

embedding_layer_names: Optional[List[str]] = None

classmethod from_dict(data: Dict[str, Any]) → T

Deserializes a dictionary into a PyTree object.

classmethod from_json(json_str: str) → T

Deserializes a JSON string into a PyTree object.

layernorm_names: Optional[List[str]] = None

module: type[easydel.infra.base_module.EasyDeLBaseModule]

replace(**kwargs)

Creates a new instance with specified fields replaced.

to_dict() → Dict[str, Any]

Serializes the PyTree object to a dictionary.

to_json(**kwargs) → str

Serializes the PyTree object to a JSON string.

class easydel.infra.factory.Registry

Bases: object

A central registry for managing EasyDeL configurations and modules.

This class provides decorators (register_config, register_module) to easily add new configurations and module implementations. It organizes registrations by configuration type and task type, allowing for retrieval based on identifiers.

property config_registry

Provides access to the underlying configuration registry dictionary.

get_config(config_type: str, config_field: ConfigType = ConfigType.MODULE_CONFIG) → Type

Retrieves a registered configuration class by its type identifier.

Parameters

• config_type (str) – The identifier of the configuration class (e.g., “llama”).

• config_field (ConfigType) – The category of the configuration. Defaults to ConfigType.MODULE_CONFIG.

Returns

The registered configuration class.

Return type

tp.Type

Raises

KeyError – If the config_type is not found in the specified config_field registry.

get_module_registration(task_type: Union[TaskType, Literal['causal-language-model', 'sequence-classification', 'vision-language-model', 'audio-classification', 'base-module', 'sequence-to-sequence']], model_type: str) → ModuleRegistration

Retrieves the ModuleRegistration information for a given task and model type.

Parameters

• task_type (TaskType | Literal) – The task type (enum or string literal).

• model_type (str) – The identifier of the model type (e.g., “llama”).

Returns

The registration information containing the module class, config class,

and optional metadata.

Return type

ModuleRegistration

Raises

AssertionError – If the task_type or model_type is not found in the registry.

register_config(config_type: str, config_field: ConfigType = ConfigType.MODULE_CONFIG) → callable

Decorator factory to register a configuration class.

Parameters

• config_type (str) – A unique string identifier for this configuration class (e.g., “llama”).

• config_field (ConfigType) – The category under which to register the config. Defaults to ConfigType.MODULE_CONFIG.

Returns

A decorator that takes the configuration class, registers it,

and enhances its string representation.

Return type

callable

register_module(task_type: TaskType, config: EasyDeLBaseConfig, model_type: str, embedding_layer_names: Optional[List[str]] = None, layernorm_names: Optional[List[str]] = None) → callable

Decorator factory to register an EasyDeL module class for a specific task.

Parameters

• task_type (TaskType) – The task the module is designed for (e.g., TaskType.CAUSAL_LM).

• config (EasyDeLBaseConfig) – The configuration class associated with this module.

• model_type (str) – A unique string identifier for this model implementation (e.g., “llama”).

• embedding_layer_names (tp.Optional[tp.List[str]]) – Optional list of embedding layer names. Defaults to None.

• layernorm_names (tp.Optional[tp.List[str]]) – Optional list of LayerNorm layer names. Defaults to None.

Returns

A decorator that takes the module class, registers it with its metadata,

and sets internal _model_task and _model_type attributes on the class.

Return type

callable

property task_registry

Provides access to the underlying task registry dictionary.

class easydel.infra.factory.TaskType(value, names=None, *, module=None, qualname=None, type=None, start=1, boundary=None)

Bases: str, Enum

Enumeration defining different model task types supported by the registry.

CAUSAL_LM

Causal Language Modeling (e.g., GPT-style models).

VISION_LM

Vision Language Modeling (models combining vision and text).

IMAGE_TEXT_TO_TEXT

Models that take image and text input to produce text output.

BASE_MODULE

Basic, potentially abstract, modules.

BASE_VISION

Basic vision modules.

SEQUENCE_TO_SEQUENCE

Sequence-to-sequence tasks (e.g., translation, summarization).

SPEECH_SEQUENCE_TO_SEQUENCE

Speech-to-text or other speech sequence tasks.

ZERO_SHOT_IMAGE_CLASSIFICATION

Image classification without task-specific training.

SEQUENCE_CLASSIFICATION

Classifying entire sequences (e.g., sentiment analysis).

AUDIO_CLASSIFICATION

Classifying audio data.

IMAGE_CLASSIFICATION

Classifying images. AUTO_BIND: Whenever to automatically decide what todo.

AUDIO_CLASSIFICATION = 'audio-classification'

AUTO_BIND = 'auto-bind'

BASE_MODULE = 'base-module'

BASE_VISION = 'vision-module'

CAUSAL_LM = 'causal-language-model'

IMAGE_CLASSIFICATION = 'image-classification'

IMAGE_TEXT_TO_TEXT = 'image-text-to-text'

SEQUENCE_CLASSIFICATION = 'sequence-classification'

SEQUENCE_TO_SEQUENCE = 'sequence-to-sequence'

SPEECH_SEQUENCE_TO_SEQUENCE = 'speech-sequence-to-sequence'

VISION_LM = 'vision-language-model'

ZERO_SHOT_IMAGE_CLASSIFICATION = 'zero-shot-image-classification'

easydel.infra.factory.register_config(config_type: str, config_field: ConfigType = ConfigType.MODULE_CONFIG) → callable

Decorator factory to register a configuration class.

Parameters

• config_type (str) – A unique string identifier for this configuration class (e.g., “llama”).

• config_field (ConfigType) – The category under which to register the config. Defaults to ConfigType.MODULE_CONFIG.

Returns

A decorator that takes the configuration class, registers it,

and enhances its string representation.

Return type

callable

easydel.infra.factory.register_module(task_type: TaskType, config: EasyDeLBaseConfig, model_type: str, embedding_layer_names: Optional[List[str]] = None, layernorm_names: Optional[List[str]] = None) → callable

Decorator factory to register an EasyDeL module class for a specific task.

Parameters

• task_type (TaskType) – The task the module is designed for (e.g., TaskType.CAUSAL_LM).

• config (EasyDeLBaseConfig) – The configuration class associated with this module.

• model_type (str) – A unique string identifier for this model implementation (e.g., “llama”).

• embedding_layer_names (tp.Optional[tp.List[str]]) – Optional list of embedding layer names. Defaults to None.

• layernorm_names (tp.Optional[tp.List[str]]) – Optional list of LayerNorm layer names. Defaults to None.

Returns

A decorator that takes the module class, registers it with its metadata,

and sets internal _model_task and _model_type attributes on the class.

Return type

callable