feat(automl): Add AutoKeras classifier wrapper and tests

This commit introduces a scikit-learn compatible wrapper for the AutoKeras library, allowing it to be integrated into the ml_grid framework.

Key additions include:

- `AutoKerasClassifierWrapper`: A new wrapper class that provides a standard sklearn interface (`fit`, `predict`, `predict_proba`) for `autokeras.StructuredDataClassifier`.
  - Automatically handles the conversion of pandas DataFrames to NumPy arrays.
  - Manages the creation and cleanup of temporary directories required by AutoKeras.
  - Includes robust error handling to raise a `RuntimeError` if AutoKeras fails to find a usable model within the given trials.
  - Correctly reshapes the output of `predict_proba` for binary classification to the standard `(n_samples, 2)` format.

- `AutoKerasClassifierClass`: A configuration class that defines the parameter spaces for the wrapper. It supports different configurations for grid search, Bayesian optimization, and a minimal `test_mode` for faster testing.

- `tests/test_auto_keras_classifier.py`: Comprehensive unit tests for both the wrapper and the configuration class. These tests use `unittest.mock` to validate the integration logic without requiring actual model training, ensuring fast and reliable test execution.

Author: SamoraHunter

ml_grid/model_classes/AutoKerasClassifierWrapper.py

@@ -0,0 +1,137 @@
+"""AutoKeras Classifier Wrapper.
+
+This module provides a scikit-learn compatible wrapper for AutoKeras StructuredDataClassifier.
+"""
+
+import logging
+import os
+import shutil
+import tempfile
+from typing import Optional, Union
+
+import numpy as np
+import pandas as pd
+from sklearn.base import BaseEstimator, ClassifierMixin
+from sklearn.utils.validation import check_is_fitted
+
+# Attempt to import AutoKeras and TensorFlow
+try:
+    import autokeras as ak
+    import tensorflow as tf
+except ImportError:
+    ak = None
+    tf = None
+
+logger = logging.getLogger(__name__)
+
+
+class AutoKerasClassifierWrapper(BaseEstimator, ClassifierMixin):
+    """A scikit-learn compatible wrapper for AutoKeras StructuredDataClassifier."""
+
+    def __init__(
+        self,
+        max_trials: int = 3,
+        epochs: int = 10,
+        validation_split: float = 0.2,
+        directory: Optional[str] = None,
+        seed: int = 42,
+        verbose: int = 1,
+        overwrite: bool = True,
+    ):
+        self.max_trials = max_trials
+        self.epochs = epochs
+        self.validation_split = validation_split
+        self.directory = directory
+        self.seed = seed
+        self.verbose = verbose
+        self.overwrite = overwrite
+
+        self.model_ = None
+        self._temp_dir = None
+
+    def fit(
+        self,
+        X: Union[np.ndarray, pd.DataFrame],
+        y: Union[np.ndarray, pd.Series],
+        **kwargs,
+    ) -> "AutoKerasClassifierWrapper":
+        if ak is None:
+            raise ImportError(
+                "AutoKeras is not installed. Please install it to use AutoKerasClassifierWrapper."
+            )
+
+        # Ensure input is numpy array to avoid AutoKeras ValueError with DataFrames
+        if isinstance(X, pd.DataFrame):
+            X = X.values
+        if isinstance(y, (pd.Series, pd.DataFrame)):
+            y = y.values
+
+        # Handle directory
+        if self.directory is None:
+            self._temp_dir = tempfile.mkdtemp(prefix="autokeras_")
+            dir_path = self._temp_dir
+        else:
+            dir_path = self.directory
+
+        if tf:
+            tf.random.set_seed(self.seed)
+
+        self.model_ = ak.StructuredDataClassifier(
+            max_trials=self.max_trials,
+            directory=dir_path,
+            seed=self.seed,
+            overwrite=self.overwrite,
+        )
+
+        self.model_.fit(
+            x=X,
+            y=y,
+            epochs=self.epochs,
+            validation_split=self.validation_split,
+            verbose=self.verbose,
+            **kwargs,
+        )
+
+        # Check if a model was actually found and can be exported.
+        try:
+            self.model_.export_model()
+        except Exception as e:
+            # This typically happens if max_trials is too low and no model is found.
+            msg = f"AutoKeras failed to find a usable model (max_trials={self.max_trials}, epochs={self.epochs}). Original error: {e}"
+            logger.error(msg)
+            raise RuntimeError(msg)
+
+        # AutoKeras does not explicitly expose classes_, so we infer them from y
+        if isinstance(y, (pd.Series, pd.DataFrame)):
+            self.classes_ = np.unique(y.values)
+        else:
+            self.classes_ = np.unique(y)
+
+        return self
+
+    def predict(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
+        check_is_fitted(self, "model_")
+        if isinstance(X, pd.DataFrame):
+            X = X.values
+        return self.model_.predict(X).flatten()
+
+    def predict_proba(self, X: Union[np.ndarray, pd.DataFrame]) -> np.ndarray:
+        check_is_fitted(self, "model_")
+        if isinstance(X, pd.DataFrame):
+            X = X.values
+        # Export the underlying Keras model to get probabilities
+        keras_model = self.model_.export_model()
+        probs = keras_model.predict(X, verbose=0)
+
+        # Handle binary classification case where output is (N, 1)
+        if probs.shape[1] == 1:
+            return np.hstack([1 - probs, probs])
+        return probs
+
+    def __del__(self):
+        # Cleanup temporary directory
+        if self._temp_dir and os.path.exists(self._temp_dir):
+            try:
+                shutil.rmtree(self._temp_dir)
+            except Exception:
+                pass

ml_grid/model_classes/auto_keras_classifier_class.py

@@ -0,0 +1,50 @@
+"""AutoKeras Classifier Configuration.
+
+This module contains the AutoKerasClassifierClass, which is a configuration
+class for the AutoKerasClassifierWrapper.
+"""
+
+import logging
+from typing import Any, Dict, List, Optional, Union
+
+import pandas as pd
+from skopt.space import Integer
+
+from ml_grid.model_classes.AutoKerasClassifierWrapper import AutoKerasClassifierWrapper
+from ml_grid.util.global_params import global_parameters
+
+logger = logging.getLogger(__name__)
+
+
+class AutoKerasClassifierClass:
+    """Configuration class for AutoKerasClassifierWrapper."""
+
+    def __init__(
+        self,
+        X: Optional[pd.DataFrame] = None,
+        y: Optional[pd.Series] = None,
+        parameter_space_size: Optional[str] = None,
+    ):
+        self.X = X
+        self.y = y
+        self.algorithm_implementation = AutoKerasClassifierWrapper()
+        self.method_name = "AutoKerasClassifier"
+
+        self.parameter_space: Union[List[Dict[str, Any]], Dict[str, Any]]
+
+        if getattr(global_parameters, "test_mode", False):
+            # Extremely small parameter space for fast unit testing
+            logger.info("Using test_mode parameter space for AutoKerasClassifier")
+            self.parameter_space = [{"max_trials": [2], "epochs": [3]}]
+        elif global_parameters.bayessearch:
+            self.parameter_space = {
+                "max_trials": Integer(3, 10),
+                "epochs": Integer(10, 30),
+            }
+        else:
+            self.parameter_space = [
+                {
+                    "max_trials": [3, 5],
+                    "epochs": [10, 20],
+                }
+            ]

tests/test_auto_keras_classifier.py

@@ -0,0 +1,126 @@
+import unittest
+from unittest.mock import MagicMock, patch
+import pandas as pd
+import numpy as np
+
+from ml_grid.model_classes.AutoKerasClassifierWrapper import AutoKerasClassifierWrapper
+from ml_grid.model_classes.auto_keras_classifier_class import AutoKerasClassifierClass
+
+
+class TestAutoKerasClassifier(unittest.TestCase):
+    def setUp(self):
+        self.X = pd.DataFrame(
+            {"feature_0": [1.0, 2.0, 3.0, 4.0], "feature_1": [4.0, 3.0, 2.0, 1.0]}
+        )
+        self.y = pd.Series([0, 1, 0, 1], name="target")
+
+    def test_init(self):
+        clf = AutoKerasClassifierWrapper(max_trials=5, epochs=20)
+        self.assertEqual(clf.max_trials, 5)
+        self.assertEqual(clf.epochs, 20)
+        self.assertIsNone(clf.model_)
+
+    @patch(
+        "ml_grid.model_classes.AutoKerasClassifierWrapper.ak.StructuredDataClassifier"
+    )
+    @patch("ml_grid.model_classes.AutoKerasClassifierWrapper.tempfile.mkdtemp")
+    @patch("ml_grid.model_classes.AutoKerasClassifierWrapper.shutil.rmtree")
+    def test_fit(self, mock_rmtree, mock_mkdtemp, mock_ak_cls):
+        # Setup mocks
+        mock_mkdtemp.return_value = "/tmp/mock_autokeras_dir"
+        mock_ak_instance = MagicMock()
+        mock_ak_cls.return_value = mock_ak_instance
+
+        clf = AutoKerasClassifierWrapper(max_trials=2)
+
+        # Test fit
+        clf.fit(self.X, self.y)
+
+        # Verify StructuredDataClassifier init
+        mock_ak_cls.assert_called_once()
+        _, kwargs = mock_ak_cls.call_args
+        self.assertEqual(kwargs["max_trials"], 2)
+        self.assertEqual(kwargs["directory"], "/tmp/mock_autokeras_dir")
+
+        # Verify fit call
+        mock_ak_instance.fit.assert_called_once()
+        _, fit_kwargs = mock_ak_instance.fit.call_args
+        self.assertEqual(fit_kwargs["epochs"], 10)
+        np.testing.assert_array_equal(fit_kwargs["x"], self.X.values)
+        np.testing.assert_array_equal(fit_kwargs["y"], self.y.values)
+
+        # Verify attributes set
+        self.assertIsNotNone(clf.model_)
+
+    @patch(
+        "ml_grid.model_classes.AutoKerasClassifierWrapper.ak.StructuredDataClassifier"
+    )
+    def test_predict(self, mock_ak_cls):
+        # Setup mock
+        mock_ak_instance = MagicMock()
+        mock_ak_cls.return_value = mock_ak_instance
+
+        # Mock predict return (AutoKeras returns array of shape (N, 1))
+        mock_ak_instance.predict.return_value = np.array([[0], [1], [0], [1]])
+
+        clf = AutoKerasClassifierWrapper()
+        clf.fit(self.X, self.y)
+
+        preds = clf.predict(self.X)
+
+        self.assertIsInstance(preds, np.ndarray)
+        self.assertEqual(preds.shape, (4,))
+        np.testing.assert_array_equal(preds, np.array([0, 1, 0, 1]))
+
+    @patch(
+        "ml_grid.model_classes.AutoKerasClassifierWrapper.ak.StructuredDataClassifier"
+    )
+    def test_predict_proba(self, mock_ak_cls):
+        # Setup mock
+        mock_ak_instance = MagicMock()
+        mock_ak_cls.return_value = mock_ak_instance
+
+        # Mock export_model and its predict method
+        mock_keras_model = MagicMock()
+        mock_ak_instance.export_model.return_value = mock_keras_model
+        # Return (N, 1) probabilities for binary classification
+        mock_keras_model.predict.return_value = np.array([[0.1], [0.9], [0.2], [0.8]])
+
+        clf = AutoKerasClassifierWrapper()
+        clf.fit(self.X, self.y)
+
+        probas = clf.predict_proba(self.X)
+
+        self.assertIsInstance(probas, np.ndarray)
+        self.assertEqual(probas.shape, (4, 2))  # Should be converted to (N, 2)
+        self.assertAlmostEqual(probas[0, 0], 0.9)  # 1 - 0.1
+        self.assertAlmostEqual(probas[0, 1], 0.1)
+
+        # Verify predict called on internal keras model with numpy array
+        mock_keras_model.predict.assert_called_once()
+        np.testing.assert_array_equal(
+            mock_keras_model.predict.call_args[0][0], self.X.values
+        )
+
+
+class TestAutoKerasClassifierClass(unittest.TestCase):
+    def test_structure(self):
+        with patch(
+            "ml_grid.model_classes.auto_keras_classifier_class.global_parameters"
+        ) as mock_globals:
+            # Case 1: Grid Search
+            mock_globals.test_mode = False
+            mock_globals.bayessearch = False
+            config = AutoKerasClassifierClass()
+            self.assertEqual(config.method_name, "AutoKerasClassifier")
+            self.assertIsInstance(config.parameter_space, list)
+
+            # Case 2: Bayes Search
+            mock_globals.bayessearch = True
+            config = AutoKerasClassifierClass()
+            self.assertIsInstance(config.parameter_space, dict)
+            self.assertIn("max_trials", config.parameter_space)
+
+
+if __name__ == "__main__":
+    unittest.main()