Struct TrainerBuilder 

pub struct TrainerBuilder { /* private fields */ }

Builder for constructing a Trainer with sensible defaults.

Implementations

impl TrainerBuilder

pub fn vocab_size(self, size: usize) -> Self

Sets the vocabulary size (required).

This is the number of distinct token IDs your tokenizer produces.

pub fn budget(self, budget: ParameterBudget) -> Self

Sets the parameter budget. Defaults to ParameterBudget::Params10M.

pub fn device(self, device: Device) -> Self

Sets the compute device. Defaults to default_device().

pub fn batch_size(self, size: usize) -> Self

Sets the batch size. Defaults to 4.

pub fn seq_len(self, len: usize) -> Self

Sets the sequence length. Defaults to 128.

pub fn steps(self, steps: usize) -> Self

Sets the number of training steps. Defaults to 1000.

pub fn learning_rate(self, lr: f64) -> Self

Sets the learning rate. Defaults to 2e-4.

pub fn weight_decay(self, wd: f64) -> Self

Sets the weight decay. Defaults to 0.01.

pub fn grad_clip_norm(self, norm: Option<f64>) -> Self

Sets gradient clip norm. Defaults to Some(1.0).

pub fn checkpoint_dir(self, dir: impl Into<String>) -> Self

Sets the directory where checkpoints are saved.

pub fn checkpoint_interval(self, steps: usize) -> Self

Sets how often (in steps) to save checkpoints. Defaults to 1000.

pub fn run_dir(self, dir: impl Into<String>) -> Self

Overrides the run directory.

pub fn build(self) -> Result<Trainer>

Builds the Trainer.

This constructs the model config from the parameter budget + vocab size, and creates the model.

Trait Implementations

impl Default for TrainerBuilder

fn default() -> Self

Returns the “default value” for a type.