diff --git a/src/lib.rs b/src/lib.rs index 6a14167..b979eac 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -8,14 +8,14 @@ //! //! - dense in-memory `f64` input //! - tree-ensemble prediction -//! - official upstream `model.json` loading for a narrow regression subset +//! - official upstream `model.json` loading for a narrow inference subset //! - explicit tree construction for tests and adapters //! //! Current official model support is intentionally narrow: //! //! - `booster=gbtree` -//! - `objective=reg:squarederror` -//! - single-target regression +//! - `objective=reg:squarederror` or `objective=binary:logistic` +//! - single-target prediction //! - numerical splits only //! //! # Example diff --git a/src/model.rs b/src/model.rs index 5046650..bc1a42f 100644 --- a/src/model.rs +++ b/src/model.rs @@ -8,6 +8,12 @@ use crate::official_model; use crate::predict; use crate::tree::RegressionTree; +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub(crate) enum PredictionTask { + Regression, + BinaryLogistic, +} + /// Inference-only gradient-boosted tree model. /// /// This type is the stable public model surface. Callers can either load a @@ -16,8 +22,9 @@ use crate::tree::RegressionTree; #[derive(Debug, Clone, PartialEq)] pub struct XGBModel { trees: Vec, - base_score: f64, + base_margin: f64, n_features: usize, + task: PredictionTask, } impl XGBModel { @@ -28,7 +35,16 @@ impl XGBModel { /// Returns [`XGBError::InvalidParameter`] if `base_score` is not finite or /// if `n_features == 0`. pub fn new(base_score: f64, n_features: usize, trees: Vec) -> Result { - if !base_score.is_finite() { + Self::from_parts(PredictionTask::Regression, base_score, n_features, trees) + } + + pub(crate) fn from_parts( + task: PredictionTask, + base_margin: f64, + n_features: usize, + trees: Vec, + ) -> Result { + if !base_margin.is_finite() { return Err(XGBError::InvalidParameter { name: "base_score", reason: "must be finite", @@ -44,15 +60,19 @@ impl XGBModel { Ok(Self { trees, - base_score, + base_margin, n_features, + task, }) } - /// Return the global bias added before tree contributions. + /// Return the serialized `XGBoost` `base_score` for supported objectives. #[must_use] pub fn base_score(&self) -> f64 { - self.base_score + match self.task { + PredictionTask::Regression => self.base_margin, + PredictionTask::BinaryLogistic => predict::sigmoid(self.base_margin), + } } /// Return the number of feature columns expected by this model. @@ -67,14 +87,25 @@ impl XGBModel { &self.trees } - /// Predict regression values for a dense feature matrix. + /// Predict task outputs for a dense feature matrix. + /// + /// For supported official JSON models this returns: + /// + /// - regression values for `reg:squarederror` + /// - positive-class probabilities for `binary:logistic` /// /// # Errors /// /// Returns [`XGBError::FeatureCountMismatch`] if the feature count differs /// from the model expectation. pub fn predict_dense(&self, features: &DenseMatrix) -> Result> { - predict::predict_ensemble(self.base_score, &self.trees, features, self.n_features) + predict::predict_ensemble( + self.task, + self.base_margin, + &self.trees, + features, + self.n_features, + ) } /// Load an official upstream `XGBoost` `model.json` file. @@ -82,8 +113,8 @@ impl XGBModel { /// Currently supported model scope: /// /// - `booster=gbtree` - /// - `objective=reg:squarederror` - /// - single-target regression + /// - `objective=reg:squarederror` or `objective=binary:logistic` + /// - single-target prediction /// - numerical splits only /// /// # Errors diff --git a/src/official_model/loader.rs b/src/official_model/loader.rs index 4e9ac2a..d115901 100644 --- a/src/official_model/loader.rs +++ b/src/official_model/loader.rs @@ -2,7 +2,7 @@ use std::fs; use std::path::Path; use crate::error::{Result, XGBError}; -use crate::model::XGBModel; +use crate::model::{PredictionTask, XGBModel}; use crate::tree::{RegressionTree, TreeNode}; use super::schema::{JsonModel, JsonTree}; @@ -21,12 +21,7 @@ pub(super) fn convert_model(model: JsonModel) -> Result { }); } - if model.learner.objective.name != "reg:squarederror" { - return Err(XGBError::UnsupportedModel { - context: "objective", - value: model.learner.objective.name, - }); - } + let task = parse_prediction_task(&model.learner.objective.name)?; let num_target = parse_usize_field( &model.learner.model_param.num_target, @@ -61,7 +56,7 @@ pub(super) fn convert_model(model: JsonModel) -> Result { }); } - let base_score = parse_base_score(&model.learner.model_param.base_score)?; + let base_margin = parse_base_score(&model.learner.model_param.base_score, task)?; let n_features = parse_usize_field( &model.learner.model_param.num_feature, "learner_model_param.num_feature", @@ -76,7 +71,18 @@ pub(super) fn convert_model(model: JsonModel) -> Result { .map(|(tree_idx, tree)| convert_tree(tree_idx, tree, n_features)) .collect::>>()?; - XGBModel::new(base_score, n_features, trees) + XGBModel::from_parts(task, base_margin, n_features, trees) +} + +fn parse_prediction_task(objective: &str) -> Result { + match objective { + "reg:squarederror" => Ok(PredictionTask::Regression), + "binary:logistic" => Ok(PredictionTask::BinaryLogistic), + _ => Err(XGBError::UnsupportedModel { + context: "objective", + value: objective.to_owned(), + }), + } } fn convert_tree( @@ -168,14 +174,14 @@ fn convert_node( }) } -fn parse_base_score(raw: &str) -> Result { +fn parse_base_score(raw: &str, task: PredictionTask) -> Result { if let Ok(value) = raw.parse::() { - return Ok(xgb_f32(value)); + return parse_base_score_value(value, task); } if let Ok(values) = serde_json::from_str::>(raw) { if values.len() == 1 { - return Ok(xgb_f32(values[0])); + return parse_base_score_value(values[0], task); } } @@ -184,6 +190,14 @@ fn parse_base_score(raw: &str) -> Result { )) } +fn parse_base_score_value(value: f64, task: PredictionTask) -> Result { + let base_score = xgb_f32(value); + match task { + PredictionTask::Regression => Ok(base_score), + PredictionTask::BinaryLogistic => logit(base_score), + } +} + fn parse_usize_field(raw: &str, context: &'static str) -> Result { raw.parse::() .map_err(|_| XGBError::InvalidModelFormat(context)) @@ -235,6 +249,8 @@ fn validate_reachable_subtree( visit_state: &mut [VisitState], indegree: &mut [usize], ) -> Result<()> { + // TODO: Replace this recursive walk with an explicit stack so valid but + // highly unbalanced trees cannot overflow the loader's call stack. match visit_state[node_idx] { VisitState::Unvisited => {} VisitState::Visiting => { @@ -312,3 +328,13 @@ fn xgb_f32(value: f64) -> f64 { )] f64::from(value as f32) } + +fn logit(probability: f64) -> Result { + if !(0.0..1.0).contains(&probability) { + return Err(XGBError::InvalidModelFormat( + "binary:logistic base_score must be strictly between 0 and 1", + )); + } + + Ok((probability / (1.0 - probability)).ln()) +} diff --git a/src/predict.rs b/src/predict.rs index b48f57c..f746876 100644 --- a/src/predict.rs +++ b/src/predict.rs @@ -2,18 +2,21 @@ use crate::dataset::DenseMatrix; use crate::error::{Result, XGBError}; +use crate::model::PredictionTask; use crate::tree::RegressionTree; /// Predict a batch of rows using a fitted tree ensemble. /// -/// Predictions start from `base_score` and then add the contribution of each tree. +/// Predictions start from `base_margin`, then add the contribution of each tree, +/// and finally apply the supported task-specific output transform. /// /// # Errors /// /// Returns [`XGBError::FeatureCountMismatch`] if `features` does not match the /// expected number of feature columns. pub fn predict_ensemble( - base_score: f64, + task: PredictionTask, + base_margin: f64, trees: &[RegressionTree], features: &DenseMatrix, expected_feature_count: usize, @@ -25,11 +28,16 @@ pub fn predict_ensemble( }); } - let mut predictions = vec![base_score; features.n_rows()]; + let mut predictions = vec![base_margin; features.n_rows()]; for (row_idx, prediction) in predictions.iter_mut().enumerate() { for tree in trees { *prediction += predict_tree(tree, features, row_idx); } + + *prediction = match task { + PredictionTask::Regression => *prediction, + PredictionTask::BinaryLogistic => sigmoid(*prediction), + }; } Ok(predictions) @@ -83,3 +91,14 @@ pub fn predict_tree(tree: &RegressionTree, features: &DenseMatrix, row_idx: usiz fn xgb_less_than(feature_value: f64, split_value: f64) -> bool { (feature_value as f32) < (split_value as f32) } + +#[must_use] +pub(crate) fn sigmoid(value: f64) -> f64 { + if value >= 0.0 { + let exp_neg = (-value).exp(); + 1.0 / (1.0 + exp_neg) + } else { + let exp_pos = value.exp(); + exp_pos / (1.0 + exp_pos) + } +} diff --git a/tests/official_model_json.rs b/tests/official_model_json.rs index 2fa8081..e996480 100644 --- a/tests/official_model_json.rs +++ b/tests/official_model_json.rs @@ -40,16 +40,33 @@ fn loads_official_regressor_json_and_predicts() { } #[test] -fn rejects_unsupported_official_objective() { - let error = XGBModel::load_json(fixture_path("classifier.json")).unwrap_err(); +fn loads_official_binary_classifier_json_and_predicts_probabilities() { + let model = XGBModel::load_json(fixture_path("classifier.json")).unwrap(); + let features = DenseMatrix::from_shape_vec( + 4, + 3, + vec![ + 0.0, 0.0, 0.0, // left-left, left-right + 0.0, 1.0, 0.0, // left-right, left-right + 0.0, -2.0, 1.0, // right-left, left-left + 0.0, 0.0, 1.0, // right-left, right-right + ], + ) + .unwrap(); - assert!(matches!( - error, - XGBError::UnsupportedModel { - context: "objective", - .. - } - )); + let predictions = model.predict_dense(&features).unwrap(); + + assert_vec_close( + &predictions, + &[0.469_360_77, 0.522_249_44, 0.553_298_23, 0.617_173_93], + ); +} + +#[test] +fn binary_classifier_reports_probability_base_score() { + let model = XGBModel::load_json(fixture_path("classifier.json")).unwrap(); + + assert!((model.base_score() - 0.52).abs() < 1.0e-6); } #[test]