38 lines
1.1 KiB
Python
Executable File
38 lines
1.1 KiB
Python
Executable File
__all__ = ['build_model']
|
|
|
|
import importlib
|
|
|
|
from utils import Registry
|
|
|
|
META_ARCHITECTURES = Registry()
|
|
_MAMBA_MODELS = frozenset({"Mamba", "MambaA"})
|
|
|
|
|
|
def _load_optional_architecture(model_name):
|
|
if model_name not in _MAMBA_MODELS or model_name in META_ARCHITECTURES:
|
|
return
|
|
|
|
try:
|
|
importlib.import_module(".mamba", package=__package__)
|
|
except ModuleNotFoundError as error:
|
|
missing_package = error.name or "mamba_ssm"
|
|
raise ModuleNotFoundError(
|
|
f"Model {model_name!r} requires the optional Mamba backend "
|
|
f"(missing package: {missing_package!r}). Install "
|
|
"requirements-mamba.txt, or use MiniROAD/MiniROADA without it."
|
|
) from error
|
|
|
|
|
|
def build_model(cfg, device=None):
|
|
model_name = cfg["model"]
|
|
_load_optional_architecture(model_name)
|
|
|
|
if model_name not in META_ARCHITECTURES:
|
|
available = ", ".join(sorted(META_ARCHITECTURES))
|
|
raise KeyError(
|
|
f"Unknown model {model_name!r}. Available architectures: {available}"
|
|
)
|
|
|
|
model = META_ARCHITECTURES[model_name](cfg)
|
|
return model.to(device)
|