ium_z444439/create-dataset.py
2023-04-19 19:46:29 +02:00

20 lines
753 B
Python

import pandas
import os
from sklearn.model_selection import train_test_split
CUTOFF = int(os.environ['CUTOFF'])
adults = pandas.read_csv('./ium_z444439/adult.csv', engine='python', encoding='ISO-8859-1', sep=',')
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)