Fix lightgbm regression (#339)

levkk · web-flow · commit fe558f4bcc59 · 2022-09-22T10:05:31.000-07:00
diff --git a/pgml-extension/pgml_rust/Cargo.toml b/pgml-extension/pgml_rust/Cargo.toml
@@ -18,7 +18,7 @@ pg_test = []
 [dependencies]
 pgx = { git="https://github.com/postgresml/pgx.git", branch="master" }
 xgboost = { git="https://github.com/postgresml/rust-xgboost.git" }
-smartcore = { git="https://github.com/smartcorelib/smartcore.git", branch="main", features = ["serde", "ndarray-bindings"] }
+smartcore = { git="https://github.com/smartcorelib/smartcore.git", branch="development", features = ["serde", "ndarray-bindings"] }
 once_cell = "1"
 rand = "0.8"
 ndarray = { version = "0.15.6", features = ["serde", "blas"] }
diff --git a/pgml-extension/pgml_rust/src/engines/lightgbm.rs b/pgml-extension/pgml_rust/src/engines/lightgbm.rs
@@ -9,29 +9,36 @@ use serde_json::json;
 pub fn lightgbm_train(task: Task, dataset: &Dataset, hyperparams: &Hyperparams) -> LightgbmBox {
     let x_train = dataset.x_train();
     let y_train = dataset.y_train();
-    let objective = match task {
-        Task::regression => "regression",
+    let mut hyperparams = hyperparams.clone();
+    match task {
+        Task::regression => {
+            hyperparams.insert(
+                "objective".to_string(),
+                serde_json::Value::from("regression"),
+            );
+        }
         Task::classification => {
             let distinct_labels = dataset.distinct_labels();
 
             if distinct_labels > 2 {
-                "multiclass"
+                hyperparams.insert(
+                    "objective".to_string(),
+                    serde_json::Value::from("multiclass"),
+                );
+                hyperparams.insert(
+                    "num_class".to_string(),
+                    serde_json::Value::from(dataset.distinct_labels()),
+                ); // [0, num_class)
             } else {
-                "binary"
+                hyperparams.insert("objective".to_string(), serde_json::Value::from("binary"));
             }
         }
     };
 
     let dataset =
         lightgbm::Dataset::from_vec(x_train, y_train, dataset.num_features as i32).unwrap();
 
-    let bst = lightgbm::Booster::train(
-        dataset,
-        &json! {{
-            "objective": objective,
-        }},
-    )
-    .unwrap();
+    let bst = lightgbm::Booster::train(dataset, &json! {hyperparams}).unwrap();
 
     LightgbmBox::new(bst)
 }
@@ -67,10 +74,12 @@ pub fn lightgbm_test(estimator: &LightgbmBox, dataset: &Dataset) -> Vec<f32> {
     let x_test = dataset.x_test();
     let num_features = dataset.num_features;
 
-    estimator.predict(&x_test, num_features as i32).unwrap()
+    let y_hat = estimator.predict(&x_test, num_features as i32).unwrap();
+    let y_hat: Vec<f32> = y_hat.into_iter().map(|y| y as f32).collect();
+    y_hat
 }
 
 /// Predict a novel datapoint using the LightGBM estimator.
 pub fn lightgbm_predict(estimator: &LightgbmBox, x: &[f32]) -> f32 {
-    estimator.predict(&x, x.len() as i32).unwrap()[0]
+    estimator.predict(&x, x.len() as i32).unwrap()[0] as f32
 }
diff --git a/pgml-extension/pgml_rust/src/engines/xgboost.rs b/pgml-extension/pgml_rust/src/engines/xgboost.rs
@@ -38,6 +38,7 @@ pub fn xgboost_train(
             Task::regression => xgboost::parameters::learning::Objective::RegLinear,
             Task::classification => {
                 xgboost::parameters::learning::Objective::MultiSoftmax(dataset.distinct_labels())
+                // [0, num_class)
             }
         })
         .build()
diff --git a/pgml-extension/pgml_rust/src/orm/snapshot.rs b/pgml-extension/pgml_rust/src/orm/snapshot.rs
@@ -396,11 +396,13 @@ impl Snapshot {
                     }
                 }
             });
+
             let num_test_rows = if self.test_size > 1.0 {
                 self.test_size as usize
             } else {
                 (num_rows as f32 * self.test_size).round() as usize
             };
+
             let num_train_rows = num_rows - num_test_rows;
             if num_train_rows == 0 {
                 error!(

Original file line number	Diff line number	Diff line change
`@@ -38,6 +38,7 @@ pub fn xgboost_train(`
`38`	`38`	`Task::regression => xgboost::parameters::learning::Objective::RegLinear,`
`39`	`39`	`Task::classification => {`
`40`	`40`	`xgboost::parameters::learning::Objective::MultiSoftmax(dataset.distinct_labels())`
	`41`	`+ // [0, num_class)`
`41`	`42`	`}`
`42`	`43`	`})`
`43`	`44`	`.build()`
Original file line number	Diff line number	Diff line change
`@@ -396,11 +396,13 @@ impl Snapshot {`
`396`	`396`	`}`
`397`	`397`	`}`
`398`	`398`	`});`
	`399`	`+`
`399`	`400`	`let num_test_rows = if self.test_size > 1.0 {`
`400`	`401`	`self.test_size as usize`
`401`	`402`	`} else {`
`402`	`403`	`(num_rows as f32 * self.test_size).round() as usize`
`403`	`404`	`};`
	`405`	`+`
`404`	`406`	`let num_train_rows = num_rows - num_test_rows;`
`405`	`407`	`if num_train_rows == 0 {`
`406`	`408`	`error!(`