Lasso (#320)

levkk · web-flow · commit f934892c55e1 · 2022-09-15T13:21:44.000-04:00
diff --git a/pgml-extension/pgml_rust/src/orm/algorithm.rs b/pgml-extension/pgml_rust/src/orm/algorithm.rs
@@ -7,6 +7,7 @@ pub enum Algorithm {
     linear,
     xgboost,
     svm,
+    lasso,
 }
 
 impl std::str::FromStr for Algorithm {
@@ -17,6 +18,7 @@ impl std::str::FromStr for Algorithm {
             "linear" => Ok(Algorithm::linear),
             "xgboost" => Ok(Algorithm::xgboost),
             "svm" => Ok(Algorithm::svm),
+            "lasso" => Ok(Algorithm::lasso),
             _ => Err(()),
         }
     }
@@ -28,6 +30,7 @@ impl std::string::ToString for Algorithm {
             Algorithm::linear => "linear".to_string(),
             Algorithm::xgboost => "xgboost".to_string(),
             Algorithm::svm => "svm".to_string(),
+            Algorithm::lasso => "lasso".to_string(),
         }
     }
 }
diff --git a/pgml-extension/pgml_rust/src/orm/estimator.rs b/pgml-extension/pgml_rust/src/orm/estimator.rs
@@ -82,6 +82,11 @@ pub fn find_deployed_estimator_by_project_name(name: &str) -> Arc<Box<dyn Estima
                 > = rmp_serde::from_read(&*data).unwrap();
                 Box::new(estimator)
             }
+            Algorithm::lasso => {
+                let estimator: smartcore::linear::lasso::Lasso<f32, Array2<f32>> =
+                    rmp_serde::from_read(&*data).unwrap();
+                Box::new(estimator)
+            }
             Algorithm::xgboost => {
                 let bst = Booster::load_buffer(&*data).unwrap();
                 Box::new(BoosterBox::new(bst))
@@ -143,6 +148,7 @@ pub fn find_deployed_estimator_by_project_name(name: &str) -> Arc<Box<dyn Estima
                 > = rmp_serde::from_read(&*data).unwrap();
                 Box::new(estimator)
             }
+            Algorithm::lasso => panic!("Lasso does not support classification"),
             Algorithm::xgboost => {
                 let bst = Booster::load_buffer(&*data).unwrap();
                 Box::new(BoosterBox::new(bst))
@@ -395,6 +401,17 @@ impl Estimator for smartcore::svm::svr::SVR<f32, Array2<f32>, smartcore::svm::RB
     }
 }
 
+#[typetag::serialize]
+impl Estimator for smartcore::linear::lasso::Lasso<f32, Array2<f32>> {
+    fn test(&self, task: Task, data: &Dataset) -> HashMap<String, f32> {
+        test_smartcore(self, task, data)
+    }
+
+    fn predict(&self, features: Vec<f32>) -> f32 {
+        predict_smartcore(self, features)
+    }
+}
+
 pub struct BoosterBox {
     contents: Box<xgboost::Booster>,
 }
diff --git a/pgml-extension/pgml_rust/src/orm/model.rs b/pgml-extension/pgml_rust/src/orm/model.rs
@@ -555,9 +555,69 @@ impl Model {
                       (PgBuiltInOids::INT8OID.oid(), self.id.into_datum()),
                       (PgBuiltInOids::BYTEAOID.oid(), bytes.into_datum()),
                   ]
-            ).unwrap();
+                ).unwrap();
                 Some(Box::new(BoosterBox::new(bst)))
             }
+
+            Algorithm::lasso => {
+                let x_train = Array2::from_shape_vec(
+                    (dataset.num_train_rows, dataset.num_features),
+                    dataset.x_train().to_vec(),
+                )
+                .unwrap();
+
+                let y_train =
+                    Array1::from_shape_vec(dataset.num_train_rows, dataset.y_train().to_vec())
+                        .unwrap();
+
+                let alpha = match hyperparams.get("alpha") {
+                    Some(alpha) => alpha.as_f64().unwrap_or(1.0) as f32,
+                    _ => 1.0,
+                };
+
+                let normalize = match hyperparams.get("normalize") {
+                    Some(normalize) => normalize.as_bool().unwrap_or(false),
+                    _ => false,
+                };
+
+                let tol = match hyperparams.get("tol") {
+                    Some(tol) => tol.as_f64().unwrap_or(1e-4) as f32,
+                    _ => 1e-4,
+                };
+
+                let max_iter = match hyperparams.get("max_iter") {
+                    Some(max_iter) => max_iter.as_u64().unwrap_or(1000) as usize,
+                    _ => 1000,
+                };
+
+                let estimator: Option<Box<dyn Estimator>> = match project.task {
+                    Task::regression => Some(Box::new(
+                        smartcore::linear::lasso::Lasso::fit(
+                            &x_train,
+                            &y_train,
+                            smartcore::linear::lasso::LassoParameters::default()
+                                .with_alpha(alpha)
+                                .with_normalize(normalize)
+                                .with_tol(tol)
+                                .with_max_iter(max_iter),
+                        )
+                        .unwrap(),
+                    )),
+
+                    Task::classification => panic!("Lasso only supports regression"),
+                };
+
+                let bytes: Vec<u8> = rmp_serde::to_vec(estimator.as_ref().unwrap()).unwrap();
+                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![
+                      (PgBuiltInOids::INT8OID.oid(), self.id.into_datum()),
+                      (PgBuiltInOids::BYTEAOID.oid(), bytes.into_datum()),
+                  ]
+                ).unwrap();
+
+                estimator
+            }
         };
     }
 

Original file line number	Diff line number	Diff line change
`@@ -7,6 +7,7 @@ pub enum Algorithm {`
`7`	`7`	`linear,`
`8`	`8`	`xgboost,`
`9`	`9`	`svm,`
	`10`	`+ lasso,`
`10`	`11`	`}`
`11`	`12`
`12`	`13`	`impl std::str::FromStr for Algorithm {`
`@@ -17,6 +18,7 @@ impl std::str::FromStr for Algorithm {`
`17`	`18`	`"linear" => Ok(Algorithm::linear),`
`18`	`19`	`"xgboost" => Ok(Algorithm::xgboost),`
`19`	`20`	`"svm" => Ok(Algorithm::svm),`
	`21`	`+ "lasso" => Ok(Algorithm::lasso),`
`20`	`22`	`_ => Err(()),`
`21`	`23`	`}`
`22`	`24`	`}`
`@@ -28,6 +30,7 @@ impl std::string::ToString for Algorithm {`
`28`	`30`	`Algorithm::linear => "linear".to_string(),`
`29`	`31`	`Algorithm::xgboost => "xgboost".to_string(),`
`30`	`32`	`Algorithm::svm => "svm".to_string(),`
	`33`	`+ Algorithm::lasso => "lasso".to_string(),`
`31`	`34`	`}`
`32`	`35`	`}`
`33`	`36`	`}`