bosk 0.1.0

//! CatBoost inference via the `catboost-rust` crate; loads the `.cbm` binary
//! format natively.
//!
//! The CatBoost C API does not expose the trained loss function, so this
//! backend cannot tell a binary classifier from a regression model — both
//! are single-dimensional. The caller states it explicitly via [`Output`]
//! when loading. Models this backend cannot evaluate faithfully are rejected
//! at load: multiclass (dimension > 1) and models with categorical, text, or
//! embedding features (the [`Model`] interface only carries `f64` features).

use std::path::Path;

use crate::{Error, Model, Result};

/// How to map CatBoost's raw formula value to [`Model::predict`] output.
///
/// This is the caller's declaration of what the model is — it cannot be read
/// from the `.cbm` file through the CatBoost C API.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Output {
    /// `1 / (1 + exp(-raw))` — for binary classifiers trained with
    /// `Logloss`/`CrossEntropy`; [`Model::predict`] returns the
    /// positive-class probability.
    Probability,
    /// The raw formula value as-is — for regression and ranking models.
    Raw,
}

/// A CatBoost model. `catboost_rust::Model::predict` takes `&self` and the type
/// is `Send + Sync`, so no interior mutability is needed — predictions run
/// concurrently without a lock.
pub struct CatBoostModel {
    model: catboost_rust::Model,
    output: Output,
    num_features: usize,
}

fn backend<E: std::fmt::Display>(e: E) -> Error {
    Error::Backend {
        backend: "catboost",
        message: e.to_string(),
    }
}

impl CatBoostModel {
    /// Load a model from a `.cbm` file. `output` declares whether the model
    /// is a binary classifier ([`Output::Probability`]) or a regression /
    /// ranking model ([`Output::Raw`]) — see [`Output`] for why this cannot
    /// be auto-detected.
    pub fn load(path: &Path, output: Output) -> Result<Self> {
        // `catboost_rust::Model::load` unwraps `to_str()` internally, so guard
        // non-UTF-8 paths here to return an error instead of panicking.
        let path_str =
            path.to_str().ok_or_else(|| Error::Backend {
                backend: "catboost",
                message: format!(
                    "path is not valid UTF-8: {}",
                    path.display()
                ),
            })?;
        let model =
            catboost_rust::Model::load(path_str).map_err(backend)?;

        let dims = model.get_dimensions_count();
        if dims != 1 {
            return Err(Error::Unsupported {
                message: format!(
                    "multiclass CatBoost model ({dims} dimensions); only \
                     single-output models are supported"
                ),
            });
        }
        for (kind, count) in [
            ("categorical", model.get_cat_features_count()),
            ("text", model.get_text_features_count()),
            ("embedding", model.get_embedding_features_count()),
        ] {
            if count != 0 {
                return Err(Error::Unsupported {
                    message: format!(
                        "model uses {count} {kind} features; this backend \
                         only passes numeric features"
                    ),
                });
            }
        }

        let num_features = model.get_float_features_count();
        Ok(Self {
            model,
            output,
            num_features,
        })
    }

    /// The number of float features the model expects per sample.
    #[must_use]
    pub fn num_features(&self) -> usize {
        self.num_features
    }

    /// One native call over `rows` samples, applying the declared [`Output`]
    /// transform to each raw formula value.
    fn predict_rows(&self, rows: &[&[f32]]) -> Result<Vec<f64>> {
        let feats = catboost_rust::ObjectsOrderFeatures::new()
            .with_float_features(rows);
        let preds = self.model.predict(feats).map_err(backend)?;
        if preds.len() != rows.len() {
            return Err(backend(format!(
                "expected {} predictions, got {}",
                rows.len(),
                preds.len()
            )));
        }
        Ok(preds
            .into_iter()
            .map(|raw| match self.output {
                Output::Probability => 1.0 / (1.0 + (-raw).exp()),
                Output::Raw => raw,
            })
            .collect())
    }

    fn check_features(&self, got: usize) -> Result<()> {
        if got != self.num_features {
            return Err(Error::FeatureCount {
                expected: self.num_features,
                got,
            });
        }
        Ok(())
    }
}

impl Model for CatBoostModel {
    fn predict(&self, features: &[f64]) -> Result<f64> {
        self.check_features(features.len())?;
        let row: Vec<f32> =
            features.iter().map(|&v| v as f32).collect();
        let out = self.predict_rows(&[row.as_slice()])?;
        Ok(out[0])
    }

    /// Runs the whole batch as one native CatBoost call rather than a call
    /// per row.
    fn predict_batch(
        &self,
        flat: &[f64],
        n_features: usize,
    ) -> Result<Vec<f64>> {
        crate::check_batch_shape(flat, n_features)?;
        self.check_features(n_features)?;
        if flat.is_empty() {
            return Ok(Vec::new());
        }
        let flat32: Vec<f32> =
            flat.iter().map(|&v| v as f32).collect();
        let rows: Vec<&[f32]> =
            flat32.chunks_exact(n_features).collect();
        self.predict_rows(&rows)
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// Feature rows shared by both parity tests. `BINARY_EXPECTED` /
    /// `REG_EXPECTED` are CatBoost 1.2.10 outputs for these rows — see
    /// `tests/fixtures/README.md`.
    const CASES: [[f64; 5]; 3] = [
        [0.5, -0.2, 0.7, 1.1, -0.9],
        [-1.0, 0.3, -0.4, 0.2, 0.6],
        [0.0, 0.0, 0.0, 0.0, 0.0],
    ];
    const BINARY_EXPECTED: [f64; 3] =
        [0.89125654109023, 0.08371516497497586, 0.5937930112700515];
    const REG_EXPECTED: [f64; 3] = [
        0.9239112738939165,
        -1.572984464485203,
        -0.039614635195669365,
    ];

    fn parity(model: &CatBoostModel, expected: &[f64; 3]) {
        for (features, expected) in CASES.iter().zip(expected) {
            let pred = model.predict(features).unwrap();
            let diff = (pred - expected).abs();
            assert!(
                diff < 1e-9,
                "mismatch on {features:?}: rust={pred}, python={expected}, diff={diff:.2e}"
            );
        }
        // The batched path must agree with the per-sample path.
        let flat: Vec<f64> =
            CASES.iter().flatten().copied().collect();
        let batch = model.predict_batch(&flat, 5).unwrap();
        for (b, expected) in batch.iter().zip(expected) {
            assert!((b - expected).abs() < 1e-9, "batch mismatch");
        }
    }

    /// Parity against CatBoost's own `Probability` output for a Logloss
    /// binary classifier ([`Output::Probability`] applies the sigmoid).
    #[test]
    fn test_catboost_binary_parity() {
        let model = CatBoostModel::load(
            Path::new("tests/fixtures/tiny_binary.cbm"),
            Output::Probability,
        )
        .expect("Failed to load CatBoost binary fixture");
        assert_eq!(model.num_features(), 5);
        parity(&model, &BINARY_EXPECTED);
    }

    /// Parity against `RawFormulaVal` for an RMSE regressor: [`Output::Raw`]
    /// must pass the value through untransformed (the regression guard for
    /// the unconditional-sigmoid bug).
    #[test]
    fn test_catboost_regression_raw_parity() {
        let model = CatBoostModel::load(
            Path::new("tests/fixtures/tiny_reg.cbm"),
            Output::Raw,
        )
        .expect("Failed to load CatBoost regression fixture");
        parity(&model, &REG_EXPECTED);
    }

    /// A 3-class MultiClass model has dimension 3 and must be refused at
    /// load instead of silently returning one class's score.
    #[test]
    fn test_catboost_multiclass_is_rejected() {
        let Err(err) = CatBoostModel::load(
            Path::new("tests/fixtures/tiny_multi.cbm"),
            Output::Probability,
        ) else {
            panic!("expected Unsupported, got Ok");
        };
        assert!(
            matches!(err, Error::Unsupported { .. }),
            "got {err:?}"
        );
    }

    #[test]
    fn test_catboost_feature_count_is_checked() {
        let model = CatBoostModel::load(
            Path::new("tests/fixtures/tiny_binary.cbm"),
            Output::Probability,
        )
        .expect("Failed to load CatBoost binary fixture");
        let err = model.predict(&[0.5, -0.2, 0.7, 1.1]).unwrap_err();
        assert!(
            matches!(
                err,
                Error::FeatureCount {
                    expected: 5,
                    got: 4
                }
            ),
            "got {err:?}"
        );
    }
}