Add stub

src/training/decision_tree/main.py

@@ -0,0 +1,71 @@
+import os
+
+from sklearn.ensemble import HistGradientBoostingRegressor, GradientBoostingRegressor, RandomForestRegressor
+from sklearn.model_selection import train_test_split, GroupShuffleSplit
+from sklearn.tree import DecisionTreeRegressor
+
+from src.training.decision_tree import preprocessor
+from src.training.decision_tree.model import Model
+from src.training.decision_tree.visualizer import Visualizer
+
+
+def create_split(x, y, game_ids):
+
+    splitter = GroupShuffleSplit(
+        n_splits=1,
+        test_size=0.10,
+        random_state=42
+    )
+
+    train_indices, test_indices = next(splitter.split(x, y, groups=game_ids))
+
+    x_train = x[train_indices]
+    x_test = x[test_indices]
+
+    y_train = y[train_indices]
+    y_test = y[test_indices]
+
+    train_games = set(game_ids[train_indices])
+    test_games = set(game_ids[test_indices])
+
+    print("Train positions:", len(x_train))
+    print("Test positions:", len(x_test))
+    print("Train games:", len(train_games))
+    print("Test games:", len(test_games))
+    print("Overlapping games:", len(train_games & test_games))
+
+    return x_train, x_test, y_train, y_test
+
+if __name__ == "__main__":
+
+    model_path = "C:/Users/Laurent/Desktop/chess-engine/src/training/decision_tree/cache/model.joblib"
+
+    if os.path.exists(model_path):
+        model = Model.load(model_path)
+
+    else:
+        x, y, game_ids = preprocessor.run("C:/Users/Laurent/Desktop/chess-engine/dataset/v1/output.csv")
+
+        split = create_split(x, y, game_ids)
+
+        depth = 35
+        leafs = int(depth * 3.5)
+
+        model = Model(
+            tree=RandomForestRegressor(
+                n_estimators=25,
+                max_depth=depth,
+                max_features=0.35,
+                min_samples_leaf=leafs,
+                random_state=42
+            ),
+            feature_names=preprocessor.get_feature_names(),
+            train_test_split=split
+        )
+
+        model.train()
+        model.evaluate()
+        #model.save_model(model_path)
+
+    visualizer = Visualizer(model=model)
+    visualizer.show_feature_importance(15)