Make EgoSpeak Mamba dependency optional
This commit is contained in:
@@ -1,9 +1,37 @@
|
||||
__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 = META_ARCHITECTURES[cfg["model"]](cfg)
|
||||
return model.to(device)
|
||||
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)
|
||||
|
||||
Reference in New Issue
Block a user