diff --git a/create_dataset.py b/create_dataset.py
index 22d9dbf..068a6e4 100644
--- a/create_dataset.py
+++ b/create_dataset.py
@@ -49,5 +49,14 @@ water_test, water_dev = train_test_split(water_test, train_size=0.66, random_sta
 
 #water["is_safe"].value_counts().plot(kind="bar")
 
-waterNorm.to_csv('waterQuality.csv', index=False)
+
+
+directory ='dane'
+if not os.path.exists(directory):
+    os.makedirs('dane')
+
+waterNorm.to_csv('waterQuality.csv', index=False)
+water_train.to_csv('dane/diamonds_train.csv', index=False)
+water_test.to_csv('dane/diamonds_test.csv', index=False)
+water_dev.to_csv('dane/diamonds_dev.csv', index=False)