Fix bug that shape mismatch error in predict when changing objective to softmax and update rust-xgboost commit (#1636)

cyccbxhl · web-flow · commit 9303cb4b68fe · 2024-10-11T15:52:59.000-07:00
diff --git a/pgml-extension/Cargo.lock b/pgml-extension/Cargo.lock
diff --git a/pgml-extension/src/bindings/lightgbm.rs b/pgml-extension/src/bindings/lightgbm.rs
@@ -100,7 +100,7 @@ impl Bindings for Estimator {
     }
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized,
     {
diff --git a/pgml-extension/src/bindings/linfa.rs b/pgml-extension/src/bindings/linfa.rs
@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};
 
 use super::Bindings;
 use crate::orm::*;
+use pgrx::*;
 
 #[derive(Debug, Serialize, Deserialize)]
 pub struct LinearRegression {
@@ -58,7 +59,7 @@ impl Bindings for LinearRegression {
     }
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized,
     {
@@ -187,7 +188,7 @@ impl Bindings for LogisticRegression {
     }
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized,
     {
@@ -261,7 +262,7 @@ impl Bindings for Svm {
     }
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized,
     {
diff --git a/pgml-extension/src/bindings/mod.rs b/pgml-extension/src/bindings/mod.rs
@@ -106,7 +106,7 @@ pub trait Bindings: Send + Sync + Debug + AToAny {
     fn to_bytes(&self) -> Result<Vec<u8>>;
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized;
 }
diff --git a/pgml-extension/src/bindings/sklearn/mod.rs b/pgml-extension/src/bindings/sklearn/mod.rs
@@ -197,7 +197,7 @@ impl Bindings for Estimator {
     }
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized,
     {
diff --git a/pgml-extension/src/bindings/xgboost.rs b/pgml-extension/src/bindings/xgboost.rs
@@ -288,10 +288,18 @@ fn fit(dataset: &Dataset, hyperparams: &Hyperparams, objective: learning::Object
         Err(e) => error!("Failed to train model:\n\n{}", e),
     };
 
-    Ok(Box::new(Estimator { estimator: booster }))
+    let softmax_objective = match hyperparams.get("objective") {
+        Some(value) => match value.as_str().unwrap() {
+            "multi:softmax" => true,
+            _ => false,
+        },
+        None => false,
+    };
+    Ok(Box::new(Estimator { softmax_objective, estimator: booster }))
 }
 
 pub struct Estimator {
+    softmax_objective: bool,
     estimator: xgboost::Booster,
 }
 
@@ -308,6 +316,9 @@ impl Bindings for Estimator {
     fn predict(&self, features: &[f32], num_features: usize, num_classes: usize) -> Result<Vec<f32>> {
         let x = DMatrix::from_dense(features, features.len() / num_features)?;
         let y = self.estimator.predict(&x)?;
+        if self.softmax_objective {
+            return Ok(y);
+        }
         Ok(match num_classes {
             0 => y,
             _ => y
@@ -340,7 +351,7 @@ impl Bindings for Estimator {
     }
 
     /// Deserialize self from bytes, with additional context
-    fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>
+    fn from_bytes(bytes: &[u8], hyperparams: &JsonB) -> Result<Box<dyn Bindings>>
     where
         Self: Sized,
     {
@@ -366,6 +377,12 @@ impl Bindings for Estimator {
             .set_param("nthread", &concurrency.to_string())
             .map_err(|e| anyhow!("could not set nthread XGBoost parameter: {e}"))?;
 
-        Ok(Box::new(Estimator { estimator }))
+        let objective_opt = hyperparams.0.get("objective").and_then(|v| v.as_str());
+        let softmax_objective = match objective_opt {
+            Some("multi:softmax") => true,
+            _ => false,
+        };
+
+        Ok(Box::new(Estimator { softmax_objective, estimator }))
     }
 }
diff --git a/pgml-extension/src/orm/file.rs b/pgml-extension/src/orm/file.rs
@@ -31,6 +31,7 @@ pub fn find_deployed_estimator_by_model_id(model_id: i64) -> Result<Arc<Box<dyn
     let mut runtime: Option<String> = None;
     let mut algorithm: Option<String> = None;
     let mut task: Option<String> = None;
+    let mut hyperparams: Option<JsonB> = None;
 
     Spi::connect(|client| {
         let result = client
@@ -39,7 +40,8 @@ pub fn find_deployed_estimator_by_model_id(model_id: i64) -> Result<Arc<Box<dyn
                     data,
                     runtime::TEXT,
                     algorithm::TEXT,
-                    task::TEXT
+                    task::TEXT,
+                    hyperparams
                 FROM pgml.models
                     INNER JOIN pgml.files
                         ON models.id = files.model_id 
@@ -66,6 +68,7 @@ pub fn find_deployed_estimator_by_model_id(model_id: i64) -> Result<Arc<Box<dyn
             runtime = result.get(2).expect("Runtime for model is corrupted.");
             algorithm = result.get(3).expect("Algorithm for model is corrupted.");
             task = result.get(4).expect("Task for project is corrupted.");
+            hyperparams = result.get(5).expect("Hyperparams for model is corrupted.");
         }
     });
 
@@ -83,6 +86,7 @@ pub fn find_deployed_estimator_by_model_id(model_id: i64) -> Result<Arc<Box<dyn
     let runtime = Runtime::from_str(&runtime.unwrap()).unwrap();
     let algorithm = Algorithm::from_str(&algorithm.unwrap()).unwrap();
     let task = Task::from_str(&task.unwrap()).unwrap();
+    let hyperparams = hyperparams.unwrap();
 
     debug1!(
         "runtime = {:?}, algorithm = {:?}, task = {:?}",
@@ -94,22 +98,22 @@ pub fn find_deployed_estimator_by_model_id(model_id: i64) -> Result<Arc<Box<dyn
     let bindings: Box<dyn Bindings> = match runtime {
         Runtime::rust => {
             match algorithm {
-                Algorithm::xgboost => crate::bindings::xgboost::Estimator::from_bytes(&data)?,
-                Algorithm::lightgbm => crate::bindings::lightgbm::Estimator::from_bytes(&data)?,
+                Algorithm::xgboost => crate::bindings::xgboost::Estimator::from_bytes(&data, &hyperparams)?,
+                Algorithm::lightgbm => crate::bindings::lightgbm::Estimator::from_bytes(&data, &hyperparams)?,
                 Algorithm::linear => match task {
-                    Task::regression => crate::bindings::linfa::LinearRegression::from_bytes(&data)?,
+                    Task::regression => crate::bindings::linfa::LinearRegression::from_bytes(&data, &hyperparams)?,
                     Task::classification => {
-                        crate::bindings::linfa::LogisticRegression::from_bytes(&data)?
+                        crate::bindings::linfa::LogisticRegression::from_bytes(&data, &hyperparams)?
                     }
                     _ => error!("Rust runtime only supports `classification` and `regression` task types for linear algorithms."),
                 },
-                Algorithm::svm => crate::bindings::linfa::Svm::from_bytes(&data)?,
+                Algorithm::svm => crate::bindings::linfa::Svm::from_bytes(&data, &hyperparams)?,
                 _ => todo!(), //smartcore_load(&data, task, algorithm, &hyperparams),
             }
         }
 
         #[cfg(feature = "python")]
-        Runtime::python => crate::bindings::sklearn::Estimator::from_bytes(&data)?,
+        Runtime::python => crate::bindings::sklearn::Estimator::from_bytes(&data, &hyperparams)?,
 
         #[cfg(not(feature = "python"))]
         Runtime::python => {
diff --git a/pgml-extension/src/orm/model.rs b/pgml-extension/src/orm/model.rs
@@ -360,6 +360,7 @@ impl Model {
                 )
                 .unwrap()
                 .unwrap();
+                let hyperparams = result.get(11).unwrap().unwrap();
 
                 let bindings: Box<dyn Bindings> = match runtime {
                     Runtime::openai => {
@@ -369,27 +370,27 @@ impl Model {
                     Runtime::rust => {
                         match algorithm {
                             Algorithm::xgboost => {
-                                xgboost::Estimator::from_bytes(&data)?
+                                xgboost::Estimator::from_bytes(&data, &hyperparams)?
                             }
                             Algorithm::lightgbm => {
-                                lightgbm::Estimator::from_bytes(&data)?
+                                lightgbm::Estimator::from_bytes(&data, &hyperparams)?
                             }
                             Algorithm::linear => match project.task {
                                 Task::regression => {
-                                    linfa::LinearRegression::from_bytes(&data)?
+                                    linfa::LinearRegression::from_bytes(&data, &hyperparams)?
                                 }
                                 Task::classification => {
-                                    linfa::LogisticRegression::from_bytes(&data)?
+                                    linfa::LogisticRegression::from_bytes(&data, &hyperparams)?
                                 }
                                 _ => bail!("No default runtime available for tasks other than `classification` and `regression` when using a linear algorithm."),
                             },
-                            Algorithm::svm => linfa::Svm::from_bytes(&data)?,
+                            Algorithm::svm => linfa::Svm::from_bytes(&data, &hyperparams)?,
                             _ => todo!(), //smartcore_load(&data, task, algorithm, &hyperparams),
                         }
                     }
 
                     #[cfg(feature = "python")]
-                    Runtime::python => sklearn::Estimator::from_bytes(&data)?,
+                    Runtime::python => sklearn::Estimator::from_bytes(&data, &hyperparams)?,
 
                     #[cfg(not(feature = "python"))]
                     Runtime::python => {
@@ -409,7 +410,7 @@ impl Model {
                     snapshot_id,
                     algorithm,
                     runtime,
-                    hyperparams: result.get(6).unwrap().unwrap(),
+                    hyperparams: hyperparams,
                     status: Status::from_str(result.get(7).unwrap().unwrap()).unwrap(),
                     metrics: result.get(8).unwrap(),
                     search: result.get(9).unwrap().map(|search| Search::from_str(search).unwrap()),

Original file line number	Diff line number	Diff line change
`@@ -100,7 +100,7 @@ impl Bindings for Estimator {`
`100`	`100`	`}`
`101`	`101`
`102`	`102`	`/// Deserialize self from bytes, with additional context`
`103`		`- fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>`
	`103`	`+ fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>`
`104`	`104`	`where`
`105`	`105`	`Self: Sized,`
`106`	`106`	`{`
Original file line number	Diff line number	Diff line change
`@@ -8,6 +8,7 @@ use serde::{Deserialize, Serialize};`
`8`	`8`
`9`	`9`	`use super::Bindings;`
`10`	`10`	`use crate::orm::*;`
	`11`	`+use pgrx::*;`
`11`	`12`
`12`	`13`	`#[derive(Debug, Serialize, Deserialize)]`
`13`	`14`	`pub struct LinearRegression {`
`@@ -58,7 +59,7 @@ impl Bindings for LinearRegression {`
`58`	`59`	`}`
`59`	`60`
`60`	`61`	`/// Deserialize self from bytes, with additional context`
`61`		`- fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>`
	`62`	`+ fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>`
`62`	`63`	`where`
`63`	`64`	`Self: Sized,`
`64`	`65`	`{`
`@@ -187,7 +188,7 @@ impl Bindings for LogisticRegression {`
`187`	`188`	`}`
`188`	`189`
`189`	`190`	`/// Deserialize self from bytes, with additional context`
`190`		`- fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>`
	`191`	`+ fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>`
`191`	`192`	`where`
`192`	`193`	`Self: Sized,`
`193`	`194`	`{`
`@@ -261,7 +262,7 @@ impl Bindings for Svm {`
`261`	`262`	`}`
`262`	`263`
`263`	`264`	`/// Deserialize self from bytes, with additional context`
`264`		`- fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>`
	`265`	`+ fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>`
`265`	`266`	`where`
`266`	`267`	`Self: Sized,`
`267`	`268`	`{`
Original file line number	Diff line number	Diff line change
`@@ -106,7 +106,7 @@ pub trait Bindings: Send + Sync + Debug + AToAny {`
`106`	`106`	`fn to_bytes(&self) -> Result<Vec<u8>>;`
`107`	`107`
`108`	`108`	`/// Deserialize self from bytes, with additional context`
`109`		`- fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>`
	`109`	`+ fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>`
`110`	`110`	`where`
`111`	`111`	`Self: Sized;`
`112`	`112`	`}`
Original file line number	Diff line number	Diff line change
`@@ -197,7 +197,7 @@ impl Bindings for Estimator {`
`197`	`197`	`}`
`198`	`198`
`199`	`199`	`/// Deserialize self from bytes, with additional context`
`200`		`- fn from_bytes(bytes: &[u8]) -> Result<Box<dyn Bindings>>`
	`200`	`+ fn from_bytes(bytes: &[u8], _hyperparams: &JsonB) -> Result<Box<dyn Bindings>>`
`201`	`201`	`where`
`202`	`202`	`Self: Sized,`
`203`	`203`	`{`