ModuleFactory (#21)

* module factory

---------

Co-authored-by: SimBe195 <[email protected]>
Co-authored-by: Albert Zeyer <[email protected]>

Author: 3 people

i6_models/config.py

@@ -17,6 +17,9 @@ class ExampleModule(Module):
 from __future__ import annotations
 from dataclasses import dataclass, fields
 import typeguard
+from torch import nn
+from typing import Generic, TypeVar, Type
+import inspect
 
 
 @dataclass
@@ -51,3 +54,37 @@ def _validate_types(self) -> None:
 
     def __post_init__(self) -> None:
         self._validate_types()
+
+
+ConfigType = TypeVar("ConfigType", bound=ModelConfiguration)
+ModuleType = TypeVar("ModuleType", bound=nn.Module)
+
+
+@dataclass
+class ModuleFactoryV1(Generic[ConfigType, ModuleType]):
+    """
+    Dataclass for a combination of a Subassembly/Part and the corresponding configuration.
+    Also provides a function to construct the corresponding object through this dataclass
+    """
+
+    module_class: Type[ModuleType]
+    cfg: ConfigType
+
+    def __call__(self) -> ModuleType:
+        """Constructs an instance of the given module class"""
+        return self.module_class(self.cfg)
+
+    def __post_init__(self) -> None:
+        # Check typing of module_class and cfg, i.e. make sure that "self.module_class(self.cfg)" is a valid call.
+        parameters = inspect.signature(self.module_class).parameters.values()
+        assert len(parameters) >= 1
+        parameter_iter = iter(parameters)
+
+        # 1. Check that the first parameter is either not annotated or the annotation matches the type of self.cfg
+        cfg_parameter = next(parameter_iter)
+        if cfg_parameter.annotation is not inspect.Parameter.empty:
+            typeguard.check_type(self.cfg, cfg_parameter.annotation)
+
+        # 2. Check that all other parameters have default values
+        for parameter in parameter_iter:
+            assert parameter.default is not inspect.Parameter.empty

tests/test_config.py

@@ -1,7 +1,8 @@
 from dataclasses import dataclass
 import pytest
 
-from i6_models.config import ModelConfiguration
+from i6_models.config import ModelConfiguration, ModuleFactoryV1
+from torch import nn
 
 
 def test_simple_configuration():
@@ -58,3 +59,31 @@ class TestConfiguration(ModelConfiguration):
     TestConfiguration(num_layers=2, hidden_dim=1)
     with pytest.raises(TypeCheckError):
         TestConfiguration(num_layers=2.0, hidden_dim="One")
+
+
+def test_module_factory():
+    @dataclass
+    class TestConfiguration(ModelConfiguration):
+        param: int = 4
+
+    @dataclass
+    class TestConfiguration2(ModelConfiguration):
+        param: float = 4.3
+
+    class TestModule(nn.Module):
+        def __init__(self, cfg: TestConfiguration):
+            super().__init__()
+            self.cfg = cfg
+
+        def forward(self):
+            pass
+
+    factory = ModuleFactoryV1(module_class=TestModule, cfg=TestConfiguration())
+    obj = factory()
+    assert type(obj) == TestModule
+    obj2 = factory()
+    assert obj != obj2
+    from typeguard import TypeCheckError
+
+    with pytest.raises(TypeCheckError):
+        ModuleFactoryV1(module_class=TestModule, cfg=TestConfiguration2())