and then have a wrapper that gets the model data and selects which way to construct and deserialise the model class.
def get_model(task_id, model_name): task = Task.get_task(task_id) try: model_data = next(model for model in task.models['output'] if model.name == model_name) except StopIteration as ex: raise ValueError(f'Model {model_name} not found in: {[model.name for model in task.models["output"]]}') filename = model_data.get_local_copy() model_type = model_data.config_dict['model_type'] if model_type == 'XGBClassifier': model = XGBClassifier() model.load_model(filename) elif model_type == 'BaseEstimator': model = joblib.load(filename) else: raise ValueError(f'Unknown model type {model_type}') return model