dvc/train.py

import argparse
import pandas as pd
from sklearn import tree
import pickle

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--train_dataset", type=str, required=True, help="Path to train dataset")
    parser.add_argument("--test_dataset", type=str, required=True, help="Path to train dataset")
    parser.add_argument("--target_column", type=str, required=True, help="Column with classname")
    parser.add_argument("--model_path", type=str, required=True, help="Path where the model will be saved")
    parser.add_argument("--predictions_path", type=str, required=True, help="Path where model predictions will be saved")
    parser.add_argument("--max_depth", type=int, default=1, help="Max tree depth")
    args = parser.parse_args()

    # Load train dataset
    df_train = pd.read_csv(args.train_dataset)
    y_train = df_train[args.target_column]
    xs_train = df_train.drop(args.target_column, axis=1)

    # Load test dataset
    df_test = pd.read_csv(args.test_dataset)
    xs_test = df_test.drop(args.target_column, axis=1)

    # Train
    clf = tree.DecisionTreeClassifier(max_depth=args.max_depth)
    clf.fit(xs_train, y_train)

    # Save model
    with open(args.model_path, 'wb') as f:
        pickle.dump(clf, f)

    # Save predictions
    predictions = pd.DataFrame()
    predictions[args.target_column] = clf.predict(xs_test)
    predictions.to_csv(args.predictions_path, index=False)