Only perform one shot encoding if the estimator does not support categorical data types

Author: DivineOmega

src/Console/Commands/Train.php

@@ -12,6 +12,7 @@
 use Rubix\ML\Classifiers\KNearestNeighbors;
 use Rubix\ML\Classifiers\MultilayerPerceptron;
 use Rubix\ML\Datasets\Labeled;
+use Rubix\ML\DataType;
 use Rubix\ML\Estimator;
 use Rubix\ML\NeuralNet\Layers\Dense;
 use Rubix\ML\NeuralNet\Layers\PReLU;
@@ -131,14 +132,7 @@ public function handle()
     private function getEstimator(string $modelPath, Estimator $baseEstimator): Estimator
     {
         $estimator = new PersistentModel(
-            new Pipeline(
-                [
-                    new MissingDataImputer(),
-                    new OneHotEncoder(),
-                    new ZScaleStandardizer(),
-                ],
-                $baseEstimator
-            ),
+            new Pipeline($this->getTransformers($baseEstimator), $baseEstimator),
             new Filesystem($modelPath)
         );
 
@@ -151,12 +145,6 @@ private function getEstimator(string $modelPath, Estimator $baseEstimator): Esti
 
     private function getDefaultBaseEstimator(bool $continuous): Estimator
     {
-//        $layers = [
-//            new Dense(100),
-//            new Dense(100),
-//            new Dense(100),
-//        ];
-
         $baseEstimator = new KDNeighbors();
 
         if ($continuous) {
@@ -165,4 +153,20 @@ private function getDefaultBaseEstimator(bool $continuous): Estimator
 
         return $baseEstimator;
     }
+
+    private function getTransformers(Estimator $estimator): array
+    {
+        $dataTypes = $estimator->compatibility();
+
+        $transformers = [];
+        $transformers[] = new MissingDataImputer();
+
+        if (!in_array(DataType::categorical(), $dataTypes) && in_array(DataType::continuous(), $dataTypes)) {
+            $transformers[] = new OneHotEncoder();
+        }
+
+        $transformers[] = new ZScaleStandardizer();
+
+        return $transformers;
+    }
 }