data preparation

This commit is contained in:
Mateusz Kuc 2023-06-30 19:16:11 +02:00
parent 7032df2b65
commit 51282fd1c1
2 changed files with 27 additions and 0 deletions

5
get-data.sh Normal file
View File

@ -0,0 +1,5 @@
#!/bin/bash
kaggle datasets download zynicide/wine-reviews
mkdir -p data
unzip -o wine-reviews.zip -d data
rm wine-reviews.zip

22
prepare_dataset.py Normal file
View File

@ -0,0 +1,22 @@
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)