diff --git a/test_egospeak_optional_mamba.py b/test_egospeak_optional_mamba.py new file mode 100644 index 0000000..5b6d1e5 --- /dev/null +++ b/test_egospeak_optional_mamba.py @@ -0,0 +1,102 @@ +import ast +import importlib.util +import sys +import types +import unittest +from pathlib import Path +from unittest import mock + + +ROOT = Path(__file__).resolve().parent +EGOSPEAK_ROOT = ROOT / "third_party" / "EgoSpeak" +MINIROAD_MODEL = EGOSPEAK_ROOT / "MiniROAD" / "model" + + +class _Registry(dict): + def register(self, name, module=None): + if module is not None: + self[name] = module + return None + + def decorator(value): + self[name] = value + return value + + return decorator + + +def _load_model_builder(): + package_name = "_egospeak_optional_mamba_test_model" + module_name = f"{package_name}.model_builder" + + package = types.ModuleType(package_name) + package.__path__ = [str(MINIROAD_MODEL)] + fake_utils = types.ModuleType("utils") + fake_utils.Registry = _Registry + + spec = importlib.util.spec_from_file_location( + module_name, MINIROAD_MODEL / "model_builder.py" + ) + module = importlib.util.module_from_spec(spec) + with mock.patch.dict( + sys.modules, + {package_name: package, module_name: module, "utils": fake_utils}, + ): + assert spec.loader is not None + spec.loader.exec_module(module) + return module + + +class EgoSpeakOptionalMambaTest(unittest.TestCase): + def test_default_requirements_do_not_install_mamba(self): + requirements = (EGOSPEAK_ROOT / "requirements.txt").read_text() + packages = { + line.partition("==")[0].strip().lower() + for line in requirements.splitlines() + if line.strip() and not line.lstrip().startswith(("#", "-")) + } + self.assertNotIn("mamba-ssm", packages) + + def test_model_package_does_not_eagerly_import_mamba(self): + tree = ast.parse((MINIROAD_MODEL / "__init__.py").read_text()) + imports = { + node.module + for node in ast.walk(tree) + if isinstance(node, ast.ImportFrom) + } + self.assertNotIn("mamba", imports) + + def test_miniroad_build_does_not_load_mamba(self): + module = _load_model_builder() + + class DummyMiniROAD: + def __init__(self, cfg): + self.cfg = cfg + + def to(self, device): + self.device = device + return self + + module.META_ARCHITECTURES["MiniROADA"] = DummyMiniROAD + with mock.patch.object(module.importlib, "import_module") as import_module: + model = module.build_model({"model": "MiniROADA"}, "cuda") + + import_module.assert_not_called() + self.assertEqual(model.device, "cuda") + + def test_mamba_model_reports_optional_dependency(self): + module = _load_model_builder() + missing = ModuleNotFoundError( + "No module named 'mamba_ssm'", name="mamba_ssm" + ) + with mock.patch.object( + module.importlib, "import_module", side_effect=missing + ): + with self.assertRaisesRegex( + ModuleNotFoundError, "optional Mamba backend" + ): + module.build_model({"model": "MambaA"}, "cuda") + + +if __name__ == "__main__": + unittest.main() diff --git a/third_party/EgoSpeak/MiniROAD/model/__init__.py b/third_party/EgoSpeak/MiniROAD/model/__init__.py index feef232..38001d2 100755 --- a/third_party/EgoSpeak/MiniROAD/model/__init__.py +++ b/third_party/EgoSpeak/MiniROAD/model/__init__.py @@ -1,4 +1,3 @@ -from .mamba import OADMamba from .model_builder import build_model -from .rnn import MROAD +from .rnn import MROAD, MROADA from .transformer_models import ViTEnc diff --git a/third_party/EgoSpeak/MiniROAD/model/model_builder.py b/third_party/EgoSpeak/MiniROAD/model/model_builder.py index 9350a84..91a2087 100755 --- a/third_party/EgoSpeak/MiniROAD/model/model_builder.py +++ b/third_party/EgoSpeak/MiniROAD/model/model_builder.py @@ -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) \ No newline at end of file + 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) diff --git a/third_party/EgoSpeak/requirements-mamba.txt b/third_party/EgoSpeak/requirements-mamba.txt new file mode 100644 index 0000000..975be5a --- /dev/null +++ b/third_party/EgoSpeak/requirements-mamba.txt @@ -0,0 +1,4 @@ +-r requirements.txt + +# Optional: only required by the MiniROAD Mamba/MambaA baselines. +mamba-ssm==2.2.4 diff --git a/third_party/EgoSpeak/requirements.txt b/third_party/EgoSpeak/requirements.txt index bcfb222..d5ee39a 100644 --- a/third_party/EgoSpeak/requirements.txt +++ b/third_party/EgoSpeak/requirements.txt @@ -16,7 +16,6 @@ huggingface-hub==0.28.1 idna==3.10 Jinja2==3.1.5 joblib==1.4.2 -mamba-ssm==2.2.4 Markdown==3.7 MarkupSafe==3.0.2 -e git+ssh://git@github.com/jun297/EgoSpeak-CAM.git@dcbfd115abd8ce6ecfe37cf4a40e71d0089c98a1#egg=MiniROAD&subdirectory=MiniROAD