| _model_entrypoints = {} | |
| def build_model(config, **kwargs): | |
| model_name = config['MODEL']['NAME'] | |
| if not is_model(model_name): | |
| raise ValueError(f'Unkown model: {model_name}') | |
| return model_entrypoints(model_name)(config, **kwargs) | |
| def register_model(fn): | |
| module_name_split = fn.__module__.split('.') | |
| model_name = module_name_split[-1] | |
| _model_entrypoints[model_name] = fn | |
| return fn | |
| def model_entrypoints(model_name): | |
| return _model_entrypoints[model_name] | |
| def is_model(model_name): | |
| return model_name in _model_entrypoints |