20 lines
704 B
Python
20 lines
704 B
Python
from typing import Dict, Set
|
|
|
|
from tha2.nn.batch_module.batch_input_module import BatchInputModule, BatchInputModuleFactory
|
|
|
|
|
|
class BatchInputModelFactory:
|
|
def __init__(self, module_factories: Dict[str, BatchInputModuleFactory]):
|
|
self.module_factories = module_factories
|
|
|
|
def get_module_names(self) -> Set[str]:
|
|
return set(self.module_factories.keys())
|
|
|
|
def create(self) -> Dict[str, BatchInputModule]:
|
|
output = {}
|
|
for name in self.module_factories:
|
|
output[name] = self.module_factories[name].create()
|
|
return output
|
|
|
|
def get_module_factory(self, module_name) -> BatchInputModuleFactory:
|
|
return self.module_factories[module_name] |