1mod predict;
5mod text;
6
7use crate::tree::Tree;
8
9#[derive(Debug, Clone, PartialEq)]
10pub enum ModelObjective {
11 Regression,
12 Binary {
14 sigmoid: f64,
15 },
16 LambdaRank,
17}
18
19#[derive(Debug, Clone, Copy, PartialEq, Eq)]
20pub enum ImportanceType {
21 Split,
23 Gain,
25}
26
27#[derive(Debug, Clone, PartialEq)]
28pub struct Model {
29 pub objective: ModelObjective,
30 pub feature_names: Vec<String>,
31 pub feature_infos: Vec<String>,
33 pub trees: Vec<Tree>,
34 pub parameters: Vec<(String, String)>,
36 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}