diff --git a/wines_train.py b/wines_train.py
index 48c8464..6d43e1f 100644
--- a/wines_train.py
+++ b/wines_train.py
@@ -1,46 +1,28 @@
-import pandas as pd
-from sklearn.model_selection import train_test_split
 import tensorflow as tf
+import pandas as pd
 import numpy as np
 
-# Wczytanie danych z pliku CSV
+# Wczytanie danych treningowych
 data = pd.read_csv('data/winemag-data_first150k.csv')
 
-# Podział danych na zbiory treningowe, deweloperskie (validation) i testowe
-train_data, temp_data = train_test_split(data, test_size=0.2, random_state=42)
-dev_data, test_data = train_test_split(temp_data, test_size=0.5, random_state=42)
+# Usunięcie rekordów zawierających wartości NaN w kolumnach 'points' i 'price'
+data = data.dropna(subset=['points', 'price'])
 
-# Zapisanie danych do osobnych plików CSV
-train_data.to_csv('train_data.csv', index=False)
-dev_data.to_csv('dev_data.csv', index=False)
-test_data.to_csv('test_data.csv', index=False)
-
-
-
-# # Wczytanie danych treningowych
-# train_data = pd.read_csv('train_data.csv')
-# dev_data = pd.read_csv('dev_data.csv')
-
-# Przygotowanie danych wejściowych i wyjściowych dla treningu
-train_input_data = train_data[['price', 'points']]  # Wejście modelu (kolumny 'price' i 'points')
-train_output_data = train_data['price']  # Oczekiwane wyjście modelu (kolumna 'price')
-
-# Przygotowanie danych wejściowych dla predykcji
-dev_input_data = dev_data[['price', 'points']]  # Wejście modelu dla predykcji (kolumny 'price' i 'points')
+# Przygotowanie danych wejściowych i wyjściowych
+input_data = data[['price', 'points']]  # Wejście modelu (kolumny 'price' i 'points')
+output_data = data['price']  # Oczekiwane wyjście modelu (kolumna 'price')
 
 # Konwersja danych do numpy arrays
-train_input_array = train_input_data.values
-train_output_array = train_output_data.values
-dev_input_array = dev_input_data.values
+input_array = input_data.values
+output_array = output_data.values
 
 # Normalizacja danych wejściowych
-input_mean = np.mean(train_input_array, axis=0)
-input_std = np.std(train_input_array, axis=0)
-train_input_array = (train_input_array - input_mean) / input_std
-dev_input_array = (dev_input_array - input_mean) / input_std
+input_mean = np.mean(input_array, axis=0)
+input_std = np.std(input_array, axis=0)
+input_array = (input_array - input_mean) / input_std
 
 # Utworzenie modelu przy użyciu TensorFlow
-input_size = train_input_array.shape[1]
+input_size = input_array.shape[1]
 output_size = 1  # Jeden wymiar dla wyjścia modelu (kolumna 'price')
 
 model = tf.keras.models.Sequential([
@@ -53,6 +35,7 @@ model = tf.keras.models.Sequential([
 model.compile(optimizer='adam', loss='mse')
 
 # Uczenie modelu
-model.fit(train_input_array, train_output_array, epochs=100, batch_size=32, verbose=1)
+history = model.fit(input_array, output_array, epochs=10, batch_size=32, verbose=1)
 
+# Zapisanie wytrenowanego modelu do pliku
 model.save('trained_model_tf.h5')