Add stub

src/training/decision_tree/model.py

@@ -0,0 +1,75 @@
+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)
+
+
+    def save_model(self, path):
+        joblib.dump({
+            "model": self._model,
+            "feature_names": self._feature_names
+        }, path)
+
+        print(f"Saved model to {path}")