added model training
This commit is contained in:
parent
8e9f37aaff
commit
8834036711
@ -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)
|
||||
|
||||
|
||||
|
Loading…
Reference in New Issue
Block a user