From 8834036711e001921c229defdcf7c551fcfb0375 Mon Sep 17 00:00:00 2001 From: s464953 Date: Thu, 9 May 2024 00:23:19 +0200 Subject: [PATCH] added model training --- model_creator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/model_creator.py b/model_creator.py index 9d5521a..f35b62f 100644 --- a/model_creator.py +++ b/model_creator.py @@ -80,7 +80,7 @@ check_datasets_presence() result_df = datasets_preparation() Y = result_df[['playlist_genre']] X = result_df.drop(columns='playlist_genre') -X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=sys.argv[1], random_state=42) +X_train, X_test, Y_train, Y_test = train_test_split(X, Y, test_size=float(sys.argv[1]), random_state=42) Y_train = np.ravel(Y_train) @@ -91,7 +91,7 @@ numeric_columns = X_train.select_dtypes(include=['int', 'float']).columns X_train_scaled = scaler.fit_transform(X_train[numeric_columns]) X_test_scaled = scaler.transform(X_test[numeric_columns]) -model = LogisticRegression(max_iter=sys.argv[2]) +model = LogisticRegression(max_iter=int(sys.argv[2])) model.fit(X_train_scaled, Y_train)