bosk 0.1.0

//! ONNX Runtime inference via the `ort` crate. Supports any model exported to
//! `.onnx` (LightGBM, CatBoost, PyTorch, …).
//!
//! The graph must take a single rank-2 `f32` tensor input
//! (`[batch, n_features]` — the layout every GBDT converter produces); this
//! is validated at load, so anything else is an [`Error::Unsupported`] up
//! front rather than an opaque runtime error at predict time.
//!
//! The model must produce a plain float tensor (`f32` or `f64`) with one or
//! two columns (raw score, or `[negative, positive]` class probabilities).
//! Classifiers exported by sklearn-onnx / onnxmltools wrap probabilities in
//! a `ZipMap` sequence by default — export with `zipmap=False` (or
//! `options={"zipmap": False}`) so the output stays a tensor. Multiclass
//! models (three or more columns) are rejected at predict time rather than
//! silently reduced to one column.

use std::path::Path;
use std::sync::Mutex;

use ort::session::builder::GraphOptimizationLevel;
use ort::session::Session;
use ort::value::{TensorElementType, Value, ValueType};

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

/// An ONNX model. `Session::run` needs `&mut self`, so the session is behind a
/// `Mutex` to keep [`Model`] object-safe and `Sync`; batch through
/// [`Model::predict_batch`] to amortise the lock over one native run.
pub struct OnnxModel {
    session: Mutex<Session>,
    /// Feature count from the validated graph input (see [`validate_input`]);
    /// `None` when the feature dimension is dynamic.
    num_features: Option<usize>,
}

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

/// Check that the graph declares the single rank-2 `f32` tensor input this
/// backend feeds, returning the declared feature count when the feature
/// dimension is static (`-1` marks a dynamic dimension).
fn validate_input(session: &Session) -> Result<Option<usize>> {
    let unsupported =
        |message: String| Error::Unsupported { message };
    let [input] = session.inputs() else {
        return Err(unsupported(format!(
            "model declares {} inputs; this backend feeds exactly one \
             [batch, n_features] tensor",
            session.inputs().len()
        )));
    };
    let ValueType::Tensor { ty, shape, .. } = input.dtype() else {
        return Err(unsupported(format!(
            "model input '{}' is not a tensor",
            input.name()
        )));
    };
    if *ty != TensorElementType::Float32 {
        return Err(unsupported(format!(
            "model input is {ty}, not float32; re-export the model \
             with a float32 input (e.g. FloatTensorType)"
        )));
    }
    if shape.len() != 2 {
        return Err(unsupported(format!(
            "model input has rank {}, expected [batch, n_features]",
            shape.len()
        )));
    }
    Ok((shape[1] > 0).then_some(shape[1] as usize))
}

impl OnnxModel {
    /// Load a model from an `.onnx` file.
    ///
    /// The graph's input is validated here: it must be a single rank-2
    /// `f32` tensor `[batch, n_features]`, anything else is refused with
    /// [`Error::Unsupported`]. When the feature dimension is static, a
    /// wrong-length feature vector is later refused with
    /// [`Error::FeatureCount`]; when dynamic, that pre-check is impossible
    /// and a shape mismatch surfaces from ONNX Runtime as
    /// [`Error::Backend`].
    pub fn load(path: &Path) -> Result<Self> {
        let session = Session::builder()
            .map_err(backend)?
            .with_optimization_level(GraphOptimizationLevel::Level3)
            .map_err(backend)?
            .with_intra_threads(1)
            .map_err(backend)?
            .commit_from_file(path)
            .map_err(backend)?;
        let num_features = validate_input(&session)?;
        Ok(Self {
            session: Mutex::new(session),
            num_features,
        })
    }

    /// The feature count declared by the graph's input, or `None` when the
    /// input's feature dimension is dynamic.
    #[must_use]
    pub fn num_features(&self) -> Option<usize> {
        self.num_features
    }

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

    /// Run one native inference over `rows` samples of `n_features` each (row
    /// major in `flat`), returning the positive-class column per row.
    fn run(
        &self,
        flat: Vec<f32>,
        rows: usize,
        n_features: usize,
    ) -> Result<Vec<f64>> {
        let input = Value::from_array(([rows, n_features], flat))
            .map_err(backend)?;

        let mut session =
            self.session.lock().expect("ONNX session mutex poisoned");
        let outputs =
            session.run(ort::inputs![input]).map_err(backend)?;

        // Classifiers: the last output holds probabilities, shape [rows, C]
        // (C = 2 for binary, or 1 for a single-score output).
        let n_outputs = outputs.len();
        if n_outputs == 0 {
            return Err(Error::Backend {
                backend: "onnx",
                message: "model produced no outputs".into(),
            });
        }
        let output = &outputs[n_outputs - 1];
        // Converters differ on the output element type; accept both floats.
        let data: Vec<f64> = match output.try_extract_tensor::<f32>()
        {
            Ok((_, data)) => {
                data.iter().map(|&v| f64::from(v)).collect()
            }
            Err(_) => {
                let (_, data) = output
                    .try_extract_tensor::<f64>()
                    .map_err(|e| {
                        backend(format!(
                            "output is neither an f32 nor an f64 \
                             tensor (classifiers exported with ZipMap \
                             are not supported — re-export with \
                             zipmap=False): {e}"
                        ))
                    })?;
                data.to_vec()
            }
        };

        if rows == 0 || data.is_empty() || data.len() % rows != 0 {
            return Err(Error::Backend {
                backend: "onnx",
                message: format!(
                    "unexpected output length {} for {rows} rows",
                    data.len()
                ),
            });
        }
        let cols = data.len() / rows;
        // One column is a raw score, two is [negative, positive] class
        // probabilities. Anything wider is multiclass — refuse instead of
        // silently returning one class's probability.
        let col = match cols {
            1 => 0,
            2 => 1,
            c => {
                return Err(Error::Backend {
                    backend: "onnx",
                    message: format!(
                        "model outputs {c} columns per row; only \
                         single-output and binary models are supported"
                    ),
                })
            }
        };
        Ok((0..rows).map(|r| data[r * cols + col]).collect())
    }
}

impl Model for OnnxModel {
    fn predict(&self, features: &[f64]) -> Result<f64> {
        self.check_features(features.len())?;
        let flat: Vec<f32> =
            features.iter().map(|&v| v as f32).collect();
        let n = flat.len();
        let out = self.run(flat, 1, n)?;
        out.into_iter().next().ok_or_else(|| Error::Backend {
            backend: "onnx",
            message: "empty prediction".into(),
        })
    }

    /// Runs the whole batch as a single `[N, n_features]` inference rather than
    /// `N` locked calls to [`predict`](Model::predict).
    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();
        self.run(flat32, flat.len() / n_features, n_features)
    }
}

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

    /// Checks the ONNX backend against the same LightGBM reference the
    /// pure-Rust parser is tested against; `tiny_binary.onnx` is the fixture
    /// booster exported with `zipmap=False` (see `tests/fixtures/README.md`).
    /// The tolerance is loose (`1e-3`) because the ONNX graph computes in
    /// f32. Skips cleanly if the fixture is removed.
    #[test]
    fn test_onnx_parity() {
        let model_path = Path::new("tests/fixtures/tiny_binary.onnx");
        if !model_path.exists() {
            eprintln!(
                "ONNX fixture not found, skipping test_onnx_parity"
            );
            return;
        }
        let model = OnnxModel::load(model_path)
            .expect("Failed to load ONNX model");

        // The fixture declares a [None, 5] input, so the feature count is
        // known and a wrong-length vector must be refused up front.
        assert_eq!(model.num_features(), Some(5));
        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:?}"
        );

        let features = [0.5, -0.2, 0.7, 1.1, -0.9];
        let pred = model.predict(&features).unwrap();
        let expected = 0.879687246542221;
        let diff = (pred - expected).abs();
        assert!(
            diff < 1e-3,
            "ONNX prediction mismatch: {pred} vs {expected}, diff={diff:.2e}"
        );

        // The single-native-run batch path must agree with predict().
        let other = [-1.0, 0.3, -0.4, 0.2, 0.6];
        let flat: Vec<f64> =
            features.iter().chain(other.iter()).copied().collect();
        let batch = model.predict_batch(&flat, 5).unwrap();
        assert_eq!(batch.len(), 2);
        assert_eq!(batch[0], pred);
        assert_eq!(batch[1], model.predict(&other).unwrap());
    }
}