20 lines
692 B
Python
20 lines
692 B
Python
import pandas as pd
|
|
import os
|
|
from sklearn.model_selection import train_test_split
|
|
|
|
CUTOFF = int(os.environ['CUTOFF'])
|
|
adults = pd.read_csv('adult.csv')
|
|
adults = adults.dropna()
|
|
|
|
adults = adults.sample(CUTOFF)
|
|
|
|
adult_X, adult_Y = adults, adults
|
|
adult_X_train, adult_X_temp, adult_Y_train, adult_Y_temp = train_test_split(adult_X, adult_Y, test_size=0.3,
|
|
random_state=1)
|
|
adult_X_dev, adult_X_test, adult_Y_dev, adult_Y_test = train_test_split(adult_X_temp, adult_Y_temp, test_size=0.3,
|
|
random_state=1)
|
|
|
|
adult_X_train.to_csv('adult_train.csv', index=False)
|
|
adult_X_dev.to_csv('adult_dev.csv', index=False)
|
|
adult_X_test.to_csv('adult_test.csv', index=False)
|