treeboost

Quick Start

Add the library to a Rust project:

cargo add treeboost

Train a baseline from a CSV file and save it:

use treeboost::auto_train_csv;

fn main() -> treeboost::Result<()> {
    let model = auto_train_csv("data.csv", "target")?;
    println!("{}", model.summary());

    model.save("model.rkyv")?;
    model.save_config("config.json")?;
    model.save_trb("model.trb", "initial training")?;
    Ok(())
}

Training writes three files:

Models

TreeBoost has three API levels. You can move between them without changing your dataset format.

Use this workflow:

1. Train with AutoModel for a baseline and a training report.
2. Read the report and config.json to see the chosen mode.
3. Refine with UniversalModel or GBDTModel for accuracy, constraints, or incremental updates.

Inference

Load a saved model and predict on a DataFrame with the training feature columns:

use treeboost::UniversalModel;

let model = UniversalModel::load("model.rkyv")?;
let predictions = model.predict_df(&df)?;

Load a TRB file with UniversalModel::load_trb("model.trb")? instead.

A model saves to rkyv or TRB and predicts through TreeBoost. There is no ONNX or XGBoost export.

UniversalModel

UniversalModel::train takes a BinnedDataset. DatasetLoader creates it from Parquet or CSV.

use treeboost::dataset::DatasetLoader;
use treeboost::loss::MseLoss;
use treeboost::{BoostingMode, UniversalConfig, UniversalModel};

let loader = DatasetLoader::new(255);
let dataset = loader.load_parquet("data.parquet", "target", None)?;

let config = UniversalConfig::new()
    .with_mode(BoostingMode::LinearThenTree)
    .with_num_rounds(100)
    .with_linear_rounds(10)
    .with_learning_rate(0.1)?;

let model = UniversalModel::train(&dataset, config, &MseLoss)?;
let predictions = model.predict(&dataset);

Choose the mode from the shape of your data:

GBDTModel

GBDTModel::train takes a flat f32 slice. Each row holds num_features values.

use treeboost::{GBDTConfig, GBDTModel};

let config = GBDTConfig::new()
    .with_num_rounds(200)
    .with_max_depth(6)
    .with_learning_rate(0.05);

let model = GBDTModel::train(&features, num_features, &targets, config, None)?;

Multi-label classification

For samples with several independent labels, use AutoModel::train_multilabel. See docs/features.md for the example.

Backends

TreeBoost selects a backend automatically. Set BackendType to override the choice.

use treeboost::backend::BackendType;
use treeboost::{GBDTConfig, GBDTModel};

let config = GBDTConfig::new().with_backend(BackendType::Scalar);
let model = GBDTModel::train(&features, num_features, &targets, config, None)?;

BackendType::Auto is the default. It selects the first available backend in this order:

CLI

Install the treeboost binary:

cargo install treeboost

Train, predict, and inspect a model:

treeboost train --data data.csv --target price --output model.rkyv
treeboost predict --model model.rkyv --data test.csv --output predictions.json
treeboost info --model model.rkyv --importances

The train and predict commands accept Parquet or CSV input. The train command writes .trb output when --output ends in .trb.

Run treeboost <command> --help for all flags.

Python

TreeBoost ships Python bindings for GBDTConfig, GBDTModel, AutoTuner, and the dataset pipeline types. Build them from source with the python feature. See docs/python.md for the build steps and an example.

Features

• AutoML mode selection: AutoModel runs probes and reports the mode it chose.
• Hybrid linear and tree mode: LinearThenTree fits trends and trees together.
• Preprocessing with the model: encoders, scalers, and imputers save with the model.
• Conformal prediction: prediction intervals from split conformal calibration.
• Incremental learning: TRB files store updates, and drift checks run before each update.

Installation and Feature Flags

cargo add treeboost --features gpu

Documentation

• API Reference: types, methods, and examples
• Advanced Features: feature matrix and incremental learning
• Saved Config: fields of the saved config.json
• Python Bindings: build steps and example
• Preset Guide: choosing a preset

Contributing

Contributions are welcome. See CONTRIBUTING.md for conventions and checks.

Questions and design discussion happen on Discord.

License

Licensed under Apache-2.0. See LICENSE.