easydel.infra.elarge_model.utils

Utility functions for ELM configuration handling.

This module provides utility functions for configuration manipulation, type coercion, and file I/O operations for the ELM system.

easydel.infra.elarge_model.utils.as_map(cfg: Any) → dict[str, Any]

Convert configuration object to dictionary.

Supports dataclasses and Mapping types, pruning None values from dataclasses.

Parameters

cfg – Configuration object (dataclass or Mapping)

Returns

Dictionary representation of the configuration

Raises

TypeError – If cfg is not a dataclass or Mapping

Example

>>> from dataclasses import dataclass
>>> @dataclass
... class Config:
...     value: int = 1
...     optional: str | None = None
>>> as_map(Config())
{'value': 1}

easydel.infra.elarge_model.utils.coerce_dtype(x: Optional[Union[str, dtype, type, Literal['fp8', 'bf16', 'fp16', 'fp32']]]) → dtype

Convert dtype-like value to JAX dtype.

Supports string representations (e.g., “bf16”, “fp32”), JAX dtypes, and various FP8 formats. Returns float32 as default.

Parameters

x – Dtype specification (string, jnp.dtype, or None)

Returns

JAX dtype object

Example

>>> coerce_dtype("bf16")
dtype('bfloat16')
>>> coerce_dtype("fp8_e4m3")
dtype('float8_e4m3')
>>> coerce_dtype(None)
dtype('float32')

easydel.infra.elarge_model.utils.coerce_precision(p: Union[str, Precision, None, Literal['HIGH', 'DEFAULT', 'HIGHEST']]) → jax._src.lax.lax.Precision | None

Convert precision-like value to JAX Precision.

Parameters

p – Precision specification (string, jax.lax.Precision, or None)

Returns

JAX Precision object or None

Example

>>> coerce_precision("HIGH")
<Precision.HIGH: 1>
>>> coerce_precision(None)
None

easydel.infra.elarge_model.utils.deep_merge(base: dict[str, Any], overlay: dict[str, Any]) → dict[str, Any]

Deep merge two dictionaries, with overlay values taking precedence.

Recursively merges nested dictionaries. Non-dict values in overlay replace corresponding values in base.

Parameters

• base – Base dictionary

• overlay – Dictionary to merge into base

Returns

New dictionary with merged values

Example

>>> base = {"a": 1, "b": {"c": 2, "d": 3}}
>>> overlay = {"b": {"c": 20, "e": 4}, "f": 5}
>>> deep_merge(base, overlay)
{'a': 1, 'b': {'c': 20, 'd': 3, 'e': 4}, 'f': 5}

easydel.infra.elarge_model.utils.infer_task_from_hf_config(model_name_or_path: str) → easydel.infra.factory.TaskType | None

Infer task type from HuggingFace model config without downloading the model.

Fetches the config.json from HuggingFace Hub and determines the task type based on the model architecture. Supports gated models through HF authentication.

Parameters

model_name_or_path – HuggingFace model ID or local path

Returns

Inferred TaskType, or None if unable to determine (will trigger fallback to CAUSAL_LM)

Example

>>> infer_task_from_hf_config("meta-llama/Llama-2-7b")
<TaskType.CAUSAL_LM: 'causal-language-model'>
>>> infer_task_from_hf_config("Qwen/Qwen2-VL-7B")
<TaskType.IMAGE_TEXT_TO_TEXT: 'image-text-to-text'>

easydel.infra.elarge_model.utils.load_elm_config(json_file_path: str | os.PathLike | eformer.paths.GCSPath | eformer.paths.LocalPath | eformer.paths.MLUtilPath) → ELMConfig

Load an ELMConfig from a JSON file.

Parameters

json_file_path – Path to the JSON file to load

Returns

The loaded and normalized configuration

Return type

ELMConfig

Example

>>> config = load_elm_config("my_config.json")
>>> model = build_model(config)

easydel.infra.elarge_model.utils.normalize_task(t: easydel.infra.factory.TaskType | str | None) → easydel.infra.factory.TaskType | None

Normalize task type specification to TaskType enum.

Handles string aliases, case variations, and hyphen/underscore differences.

Parameters

t – Task type specification (TaskType, string alias, or None)

Returns

Normalized TaskType or None if not recognized

Example

>>> normalize_task("causal-lm")
<TaskType.CAUSAL_LM: 'causal_lm'>
>>> normalize_task("LM")
<TaskType.CAUSAL_LM: 'causal_lm'>

easydel.infra.elarge_model.utils.prune_nones(obj: Any) → Any

Recursively remove None values from nested data structures.

Parameters

obj – Object to prune (dict, list, tuple, or any value)

Returns

Object with None values removed from dicts and preserved structure

Example

>>> data = {"a": 1, "b": None, "c": {"d": 2, "e": None}}
>>> prune_nones(data)
{'a': 1, 'c': {'d': 2}}

easydel.infra.elarge_model.utils.save_elm_config(config: easydel.infra.elarge_model.types.ELMConfig | collections.abc.Mapping[str, Any], json_file_path: str | os.PathLike | eformer.paths.GCSPath | eformer.paths.LocalPath | eformer.paths.MLUtilPath) → None

Save an ELMConfig to a JSON file.

Parameters

• config – The ELMConfig or config dict to save

• json_file_path – Path to the JSON file where the config will be saved

Example

>>> config = {"model": {"name_or_path": "meta-llama/Llama-2-7b"}}
>>> save_elm_config(config, "my_config.json")