Make EgoSpeak Mamba dependency optional

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()

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

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)
+    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)

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

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