dvc/preprocess.py

import argparse
import pandas as pd
import numpy as np

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument("--dataset", type=str, required=True, help="Path to train dataset")
    parser.add_argument("--out_train", type=str, required=True, help="Column with classname")
    parser.add_argument("--out_test", type=str, required=True, help="Column with classname")
    args = parser.parse_args()

    df = pd.read_csv(args.dataset)
    train, test = np.split(df.sample(frac=1), [int(.6*len(df))])

    train.to_csv(args.out_train, index=False)
    test.to_csv(args.out_test, index=False)