83 lines
2.2 KiB
Python
83 lines
2.2 KiB
Python
import os
|
|
import joblib
|
|
from sklearn.metrics import mean_absolute_error, mean_squared_error
|
|
|
|
|
|
class Model:
|
|
def __init__(self, tree=None, feature_names=None, train_test_split=None):
|
|
self._model = tree
|
|
self._feature_names = feature_names
|
|
self._x_train = self._x_test = None
|
|
self._y_train = self._y_test = None
|
|
|
|
if train_test_split is not None:
|
|
self._x_train, self._x_test, \
|
|
self._y_train, self._y_test = train_test_split
|
|
|
|
@classmethod
|
|
def load(cls, path):
|
|
data = joblib.load(path)
|
|
|
|
model = cls()
|
|
model._model = data["model"]
|
|
model._feature_names = data["feature_names"]
|
|
|
|
print(f"Loaded model from {path}")
|
|
|
|
return model
|
|
|
|
def get_feature_names(self):
|
|
return self._feature_names
|
|
|
|
def get_feature_importance(self):
|
|
importances = self._model.feature_importances_
|
|
|
|
return sorted(
|
|
zip(self._feature_names, importances),
|
|
key=lambda x: x[1],
|
|
reverse=True
|
|
)
|
|
|
|
def train(self):
|
|
print("Start training...")
|
|
self._model.fit(self._x_train, self._y_train)
|
|
print("Training complete.")
|
|
|
|
def evaluate(self):
|
|
print("Evaluating model...")
|
|
|
|
train_predictions = self._model.predict(self._x_train)
|
|
test_predictions = self._model.predict(self._x_test)
|
|
|
|
train_mae = mean_absolute_error(self._y_train, train_predictions)
|
|
train_mse = mean_squared_error(self._y_train, train_predictions)
|
|
|
|
test_mae = mean_absolute_error(self._y_test, test_predictions)
|
|
test_mse = mean_squared_error(self._y_test, test_predictions)
|
|
|
|
print()
|
|
print("Training:")
|
|
print("MAE:", train_mae)
|
|
print("MSE:", train_mse)
|
|
|
|
print()
|
|
print("Test:")
|
|
print("MAE:", test_mae)
|
|
print("MSE:", test_mse)
|
|
|
|
return [train_mae, train_mse, test_mae, test_mse]
|
|
|
|
def release_training_data(self):
|
|
self._x_train = None
|
|
self._x_test = None
|
|
self._y_train = None
|
|
self._y_test = None
|
|
|
|
|
|
def save_model(self, path):
|
|
joblib.dump({
|
|
"model": self._model,
|
|
"feature_names": self._feature_names
|
|
}, path)
|
|
|
|
print(f"Saved model to {path}") |