train test split

train_test_split.py

import random

def train_test_split(data, labels, test_prop=0.3):
    """roll your own train test split"""
    assert len(data) == len(labels)
    n_test = round(test_prop * len(data))
    n_train = len(data) - n_test
    combined = list(zip(data, labels))
    random.shuffle(combined)
    x_train, y_train = zip(*combined[:n_train])
    x_test, y_test = zip(*combined[-n_test:])
    return [list(i) for i in [x_train, y_train, x_test, y_test]]