Example code snippet for Naive Bayes fundamentals article, part [2]

bayes_train_test_split.py

# [2] now split into train/test set
# create our mask (70%)
mask = np.random.rand(len(dataset)) < 0.7

train = dataset[mask]  # get 70% of samples from mask indices
test = dataset[~mask]  # get other 30% of samples

# we also need to split the data based on whether person earns
# more than or less than 50K
less = train[train['income'] == '<=50K']
more = train[train['income'] == '>50K']