Train a model with a config file using Metaflow

config_train.py

from metaflow import FlowSpec, step, IncludeFile

def dataset_wine():
    from sklearn import datasets
    return datasets.load_wine(return_X_y=True)
    
def model_knn(train_data, train_labels):
    from sklearn.neighbors import KNeighborsClassifier
    model = KNeighborsClassifier()
    model.fit(train_data, train_labels)
    return model

def model_svm(train_data, train_labels):
    from sklearn import svm
    model = svm.SVC(kernel='poly')
    model.fit(train_data, train_labels)
    return model

MODELS = {'knn': model_knn,
          'svm': model_svm}

DATASETS = {'wine': dataset_wine}

class TrainWithConfigFlow(FlowSpec):

    config_file = IncludeFile('config', default='config.json')

    @step
    def start(self):
        import json
        self.config = json.loads(self.config_file)
        self.next(self.load_data)

    @step
    def load_data(self):
        from sklearn.model_selection import train_test_split
        print('Loading dataset', self.config['dataset'])
        X, y = DATASETS[self.config['dataset']]()

        self.train_data,\
        self.test_data,\
        self.train_labels,\
        self.test_labels = train_test_split(X, y, test_size=0.2, random_state=0)
        self.next(self.train)

    @step
    def train(self):
        print("Training model", self.config['model'])
        self.model = MODELS[self.config['model']](self.train_data,
                                                  self.train_labels)
        self.next(self.end)

    @step
    def end(self):
        self.score = self.model.score(self.test_data, self.test_labels)
        print('Eval score', self.score)

if __name__ == '__main__':
    TrainWithConfigFlow()