36 lines
1.3 KiB
Python
36 lines
1.3 KiB
Python
from pandas.core.frame import DataFrame
|
|
from sklearn.model_selection import train_test_split
|
|
from sklearn.neural_network import MLPRegressor
|
|
from sklearn.metrics import mean_squared_error
|
|
|
|
import math
|
|
|
|
class PredictionModelTrainer:
|
|
def __init__(self, preprocessed_data : DataFrame) -> None:
|
|
self.data_ : DataFrame = preprocessed_data
|
|
self.trained_model_ : MLPRegressor = None
|
|
|
|
def train(self) -> None:
|
|
print("Training home pricing model with MLPRegressor")
|
|
X = self.data_.drop(columns=['Price'])
|
|
y = self.data_['Price']
|
|
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=5)
|
|
self.trained_model_ = MLPRegressor(
|
|
activation="relu",
|
|
hidden_layer_sizes=(10, 80, 200),
|
|
max_iter=2000,
|
|
random_state=5,
|
|
alpha=0.01,
|
|
solver="lbfgs"
|
|
)
|
|
self.trained_model_.fit(X_train, y_train)
|
|
y_pred = self.trained_model_.predict(X_test)
|
|
mse = mean_squared_error(y_test, y_pred)
|
|
print("Mean Squared Error:", mse)
|
|
rmse = math.sqrt(mse)
|
|
print(f"Root mean squared error: {rmse}")
|
|
print(f"Model on avaerage is wrong by {round(rmse, 2)} PLN")
|
|
|
|
def get_trained_model(self) -> MLPRegressor:
|
|
return self.trained_model_
|