103 lines
3.1 KiB
Python
103 lines
3.1 KiB
Python
|
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()
|