`nemo_automodel._transformers.registry`#

Module Contents#

Classes#

`_LazyArchMapping`	Lazy-loading mapping from architecture name to model class.
`_ModelRegistry`

Functions#

get_registry

Data#

`logger`
`MODEL_ARCH_MAPPING`
`ModelRegistry`

API#

nemo_automodel._transformers.registry.logger#: ‘getLogger(…)’

nemo_automodel._transformers.registry.MODEL_ARCH_MAPPING#: ‘OrderedDict(…)’

class nemo_automodel._transformers.registry._LazyArchMapping( auto_map: Union[collections.OrderedDict, Dict[str, Tuple[str, str]], None] = None, )#

Lazy-loading mapping from architecture name to model class.

Inspired by HuggingFace transformers’ _LazyAutoMapping. Entries from the static auto_map are imported on first access and cached. Additional entries can be added at runtime via register.

Initialization

_load(key: str) → Type[torch.nn.Module]#

__contains__(key: str) → bool#

__getitem__(key: str) → Type[torch.nn.Module]#

__setitem__(key: str, value: Type[torch.nn.Module]) → None#

register( key: str, value: Type[torch.nn.Module], exist_ok: bool = False, ) → None#: Register a model class under the given architecture name.

keys()#

__len__() → int#

__repr__() → str#

class nemo_automodel._transformers.registry._ModelRegistry#

model_arch_name_to_cls: nemo_automodel._transformers.registry._LazyArchMapping#: ‘field(…)’

__post_init__()#

property supported_models#

get_model_cls_from_model_arch( model_arch: str, ) → Type[torch.nn.Module]#

register( arch_name: str, model_cls: Type[torch.nn.Module], exist_ok: bool = False, ) → None#: Register a custom model class for a given architecture name.

nemo_automodel._transformers.registry.get_registry()#

nemo_automodel._transformers.registry.ModelRegistry#: ‘get_registry(…)’

nemo_automodel._transformers.registry#