postgresml · levkk · Sep 20, 2022 · Sep 20, 2022 · Sep 20, 2022 · Sep 20, 2022
diff --git a/pgml-extension/pgml_rust/sql/schema.sql b/pgml-extension/pgml_rust/sql/schema.sql
@@ -81,7 +81,7 @@ CREATE TABLE IF NOT EXISTS pgml_rust.models(
 	hyperparams JSONB NOT NULL,
 	status TEXT NOT NULL,
 	metrics JSONB,
-	search pgml_rust.search,
+	search TEXT,
 	search_params JSONB NOT NULL,
 	search_args JSONB NOT NULL,
 	created_at TIMESTAMP WITHOUT TIME ZONE NOT NULL DEFAULT clock_timestamp(),

diff --git a/pgml-extension/pgml_rust/src/engines/mod.rs b/pgml-extension/pgml_rust/src/engines/mod.rs
@@ -2,3 +2,7 @@ pub mod engine;
 pub mod sklearn;
 pub mod smartcore;
 pub mod xgboost;
+
+use serde_json;
+
+pub type Hyperparams = serde_json::Map<std::string::String, serde_json::Value>;
diff --git a/pgml-extension/pgml_rust/src/engines/sklearn.rs b/pgml-extension/pgml_rust/src/engines/sklearn.rs
@@ -10,9 +10,11 @@
 use pyo3::prelude::*;
 use pyo3::types::PyTuple;
 
+use crate::engines::Hyperparams;
 use crate::orm::algorithm::Algorithm;
 use crate::orm::dataset::Dataset;
 use crate::orm::estimator::SklearnBox;
+use crate::orm::search::Search;
 use crate::orm::task::Task;
 
 use pgx::*;
@@ -171,3 +173,60 @@ pub fn sklearn_load(data: &Vec<u8>) -> SklearnBox {
         SklearnBox::new(estimator)
     })
 }
+
+/// Hyperparameter search using Scikit's
+/// RandomizedSearchCV or GridSearchCV.
+pub fn sklearn_search(
+    task: Task,
+    algorithm: Algorithm,
+    search: Search,
+    dataset: &Dataset,
+    hyperparams: &Hyperparams,
+    search_params: &Hyperparams,
+) -> (SklearnBox, Hyperparams) {
+    let module = include_str!(concat!(
+        env!("CARGO_MANIFEST_DIR"),
+        "/src/engines/wrappers.py"
+    ));
+
+    let algorithm_name = match task {
+        Task::regression => match algorithm {
+            Algorithm::linear => "linear_regression",
+            _ => todo!(),
+        },
+
+        Task::classification => match algorithm {
+            Algorithm::linear => "linear_classification",
+            _ => todo!(),
+        },
+    };
+
+    Python::with_gil(|py| -> (SklearnBox, Hyperparams) {
+        let module = PyModule::from_code(py, module, "", "").unwrap();
+        let estimator_search = module.getattr("estimator_search").unwrap();
+        let train = estimator_search
+            .call1(PyTuple::new(
+                py,
+                &[
+                    algorithm_name.into_py(py),
+                    dataset.num_features.into_py(py),
+                    serde_json::to_string(hyperparams).unwrap().into_py(py),
+                    serde_json::to_string(search_params).unwrap().into_py(py),
+                    search.to_string().into_py(py),
+                    None::<String>.into_py(py),
+                ],
+            ))
+            .unwrap();
+
+        let (estimator, hyperparams): (Py<PyAny>, String) = train
+            .call1(PyTuple::new(py, &[dataset.x_train(), dataset.y_train()]))
+            .unwrap()
+            .extract()
+            .unwrap();
+
+        let estimator = SklearnBox::new(estimator);
+        let hyperparams: Hyperparams = serde_json::from_str::<Hyperparams>(&hyperparams).unwrap();
+
+        (estimator, hyperparams)
+    })
+}
diff --git a/pgml-extension/pgml_rust/src/engines/wrappers.py b/pgml-extension/pgml_rust/src/engines/wrappers.py
@@ -75,7 +75,7 @@ def estimator_joint(algorithm_name, num_features, num_targets, hyperparams):
     """Returns the correct estimator based on algorithm names we defined
     internally (see dict above).
 
-    
+
     Parameters:
         - algorithm_name: The human-readable name of the algorithm (see dict above).
         - num_features: The number of features in X.
@@ -101,6 +101,70 @@ def train(X_train, y_train):
     return train
 
 
+def estimator_search_joint(algorithm_name, num_features, num_targets, hyperparams, search_params, search, search_args):
+    """Hyperparameter search.
+
+    Parameters:
+        - algorithm_name: The human-readable name of the algorithm (see dict above).
+        - num_features: The number of features in X.
+        - num_targets: For joint training (more than one y target).
+        - hyperparams: JSON of hyperparameters.
+        - search_params: Hyperparameters to search (see Scikit docs for examples).
+        - search: Type of search to do, grid or random.
+        - search_args: See Scikit docs for examples.
+
+    Return:
+        A tuple of Estimator and chosen hyperparameters.
+    """
+    if search_args is None:
+        search_args = {}
+    else:
+        search_args = json.loads(search_args)
+
+    if search is None:
+        search = "grid"
+
+    search_params = json.loads(search_params)
+    hyperparams = json.loads(hyperparams)
+
+    if search == "random":
+        algorithm = sklearn.model_selection.RandomizedSearchCV(
+            _ALGORITHM_MAP[algorithm_name](**hyperparams),
+            search_params,
+        )
+    elif search == "grid":
+        algorithm = sklearn.model_selection.GridSearchCV(
+            _ALGORITHM_MAP[algorithm_name](**hyperparams),
+            search_params,
+        )
+    else:
+        raise Exception(f"search can be 'grid' or 'random', got: '{search}'")
+
+    def train(X_train, y_train):
+        X_train = np.asarray(X_train).reshape((-1, num_features))
+        y_train = np.asarray(y_train).reshape((-1, num_targets))
+
+        algorithm.fit(X_train, y_train)
+
+        return (algorithm.best_estimator_, json.dumps(algorithm.best_params_))
+
+    return train
+
+
+def estimator_search(algorithm_name, num_features, hyperparams, search_params, search, search_args):
+    """Hyperparameter search.
+
+    Parameters:
+        - algorithm_name: The human-readable name of the algorithm (see dict above).
+        - num_features: The number of features in X.
+        - hyperparams: JSON of hyperparameters.
+        - search_params: Hyperparameters to search (see Scikit docs for examples).
+        - search: Type of search to do, grid or random.
+        - search_args: See Scikit docs for examples.
+    """
+    return estimator_search_joint(algorithm_name, num_features, 1, hyperparams, search_params, search, search_args)
+
+
 def test(estimator, X_test):
     """Validate the estimator using the test dataset.
 
@@ -134,6 +198,7 @@ def predictor_joint(estimator, num_features, num_targets):
         - num_features: The number of features in X.
         - num_targets: Used in joint models (more than 1 y target).
     """
+
     def predict(X):
         X = np.asarray(X).reshape((-1, num_features))
         y_hat = estimator.predict(X)
@@ -149,7 +214,7 @@ def predict(X):
 
 def save(estimator):
     """Save the estimtator as bytes (pickle).
-    
+
     Parameters:
         - estimator: Scikit-Learn estimator, instantiated.
 

diff --git a/pgml-extension/pgml_rust/src/orm/model.rs b/pgml-extension/pgml_rust/src/orm/model.rs
@@ -11,7 +11,7 @@ use crate::orm::Project;
 use crate::orm::Search;
 use crate::orm::Snapshot;
 
-use crate::engines::sklearn::{sklearn_save, sklearn_train};
+use crate::engines::sklearn::{sklearn_save, sklearn_search, sklearn_train};
 use crate::engines::smartcore::{smartcore_save, smartcore_train};
 use crate::engines::xgboost::{xgboost_save, xgboost_train};
 
@@ -67,7 +67,7 @@ impl Model {
         Spi::connect(|client| {
             let result = client.select("
           INSERT INTO pgml_rust.models (project_id, snapshot_id, algorithm, hyperparams, status, search, search_params, search_args, engine) 
-          VALUES ($1, $2, $3, $4, $5, $6::pgml_rust.search, $7, $8, $9) 
+          VALUES ($1, $2, $3, $4, $5, $6, $7, $8, $9) 
           RETURNING id, project_id, snapshot_id, algorithm, hyperparams, status, metrics, search, search_params, search_args, created_at, updated_at;",
               Some(1),
               Some(vec![
@@ -76,7 +76,10 @@ impl Model {
                   (PgBuiltInOids::TEXTOID.oid(), algorithm.to_string().into_datum()),
                   (PgBuiltInOids::JSONBOID.oid(), hyperparams.into_datum()),
                   (PgBuiltInOids::TEXTOID.oid(), "new".to_string().into_datum()),
-                  (PgBuiltInOids::TEXTOID.oid(), search.into_datum()),
+                  (PgBuiltInOids::TEXTOID.oid(), match search {
+                    Some(search) => Some(search.to_string()),
+                    None => None,
+                  }.into_datum()),
                   (PgBuiltInOids::JSONBOID.oid(), search_params.into_datum()),
                   (PgBuiltInOids::JSONBOID.oid(), search_args.into_datum()),
                   (PgBuiltInOids::TEXTOID.oid(), engine.to_string().into_datum()),
@@ -117,13 +120,30 @@ impl Model {
     fn fit(&mut self, project: &Project, dataset: &Dataset) {
         // Get the hyperparameters.
         let hyperparams: &serde_json::Value = &self.hyperparams.0;
-        let hyperparams = hyperparams.as_object().unwrap();
+        let mut hyperparams = hyperparams.as_object().unwrap().clone();
 
         // Train the estimator. We are getting the estimator struct and
         // it's serialized form to save into the `models` table.
         let (estimator, bytes): (Box<dyn Estimator>, Vec<u8>) = match self.engine {
             Engine::sklearn => {
-                let estimator = sklearn_train(project.task, self.algorithm, dataset, &hyperparams);
+                let estimator = match self.search {
+                    Some(search) => {
+                        let (estimator, chosen_hyperparams) = sklearn_search(
+                            project.task,
+                            self.algorithm,
+                            search,
+                            dataset,
+                            &hyperparams,
+                            &self.search_params.0.as_object().unwrap(),
+                        );
+
+                        hyperparams.extend(chosen_hyperparams);
+
+                        estimator
+                    }
+
+                    None => sklearn_train(project.task, self.algorithm, dataset, &hyperparams),
+                };
 
                 let bytes = sklearn_save(&estimator);
 
@@ -150,7 +170,7 @@ impl Model {
             _ => todo!(),
         };
 
-        // Save the estimator
+        // Save the estimator.
         Spi::get_one_with_args::<i64>(
           "INSERT INTO pgml_rust.files (model_id, path, part, data) VALUES($1, 'estimator.rmp', 0, $2) RETURNING id",
           vec![
@@ -159,6 +179,19 @@ impl Model {
           ]
         ).unwrap();
 
+        // Save the hyperparams after search
+        Spi::get_one_with_args::<i64>(
+            "UPDATE pgml_rust.models SET hyperparams = $1::jsonb WHERE id = $2 RETURNING id",
+            vec![
+                (
+                    PgBuiltInOids::TEXTOID.oid(),
+                    serde_json::to_string(&hyperparams).unwrap().into_datum(),
+                ),
+                (PgBuiltInOids::INT8OID.oid(), self.id.into_datum()),
+            ],
+        )
+        .unwrap();
+
         self.estimator = Some(estimator);
     }