pymllm.models

Model registry for pymllm.

Maps HuggingFace config.architectures[0] strings to pymllm model classes. Models are imported lazily via importlib so that heavy dependencies (torch, numpy, etc.) are only loaded when a model is actually requested.

Submodules

Attributes

Functions

get_model_class(architecture)

Look up a pymllm model class by HuggingFace architecture string.

Package Contents

pymllm.models.logger
pymllm.models.get_model_class(architecture)

Look up a pymllm model class by HuggingFace architecture string.

Returns None if the architecture is not registered or cannot be imported. The caller is responsible for raising an appropriate error.

Parameters:

architecture (str)

Return type:

Optional[Type[torch.nn.Module]]