File size: 567 Bytes
814a594 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 |
_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 |