Skip to main content

silvic_core/model/
mod.rs

1//! A trained model in LightGBM 3.3.5's text format (`GBDT::SaveModelToString`), which LightGBM 3.3.5 and 4.x load
2//! and lleaves compiles, and the predictor over it.
3
4mod predict;
5mod text;
6
7use crate::tree::Tree;
8
9#[derive(Debug, Clone, PartialEq)]
10pub enum ModelObjective {
11    Regression,
12    /// `binary sigmoid:<sigmoid>`.
13    Binary {
14        sigmoid: f64,
15    },
16    LambdaRank,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ImportanceType {
21    /// Times a feature is split on.
22    Split,
23    /// Total gain of a feature's splits.
24    Gain,
25}
26
27#[derive(Debug, Clone, PartialEq)]
28pub struct Model {
29    pub objective: ModelObjective,
30    pub feature_names: Vec<String>,
31    /// `BinMapper::info_string` per feature.
32    pub feature_infos: Vec<String>,
33    pub trees: Vec<Tree>,
34    /// `[key: value]` lines of the `parameters:` block.
35    pub parameters: Vec<(String, String)>,
36    /// The JSON after `pandas_categorical:` (LightGBM's Python package stores pandas category levels there).
37    pub pandas_categorical: Option<String>,
38}
39
40impl Model {
41    pub fn num_features(&self) -> usize {
42        self.feature_names.len()
43    }
44
45    pub fn num_trees(&self) -> usize {
46        self.trees.len()
47    }
48}