23 lines
692 B
Python
23 lines
692 B
Python
|
import pandas as pd
|
||
|
|
import os
|
||
|
|
from sklearn.model_selection import train_test_split
|
||
|
|
|
||
|
|
CUTOFF = int(os.environ['CUTOFF'])
|
||
|
|
|
||
|
|
wines = pd.read_csv('data/winemag-data_first150k.csv', engine='python', encoding='ISO-8859-1', sep=',')
|
||
|
|
|
||
|
|
wines = wines.dropna()
|
||
|
|
|
||
|
|
wines = wines.sample(100)
|
||
|
|
|
||
|
|
X, Y = wines, wines
|
||
|
|
|
||
|
|
# SPLIT BETWEEN DEV, TRAINS, AND TEST
|
||
|
|
wines_train, wines_temp, wines_train, wines_temp = train_test_split(X, Y, test_size=0.2, random_state=1)
|
||
|
|
wines_dev, wines_test, wines_dev, wines_test = train_test_split(wines_temp, wines_temp, test_size=0.2)
|
||
|
|
|
||
|
|
wines_train.to_csv('wines_train.csv', index=False)
|
||
|
|
wines_dev.to_csv('wines_dev.csv', index=False)
|
||
|
|
wines_test.to_csv('wines_test.csv', index=False)
|
||
|
|
|