klib/ml/base/
mod.rs

1//! Base types for machine learning.
2
3// Add an allow for the [`Config`] derive on [`TrainConfig`].
4#![allow(clippy::too_many_arguments)]
5
6pub mod data;
7#[cfg(all(feature = "ml_sample_gather", feature = "analyze_mic"))]
8pub mod gather;
9pub mod helpers;
10pub mod model;
11pub mod precision;
12#[cfg(feature = "ml_sample_process")]
13pub mod process;
14
15pub use precision::StorePrecisionSettings;
16#[cfg(feature = "ml_train")]
17pub use precision::{PrecisionElement, PRECISION_DTYPE};
18
19use burn::config::Config;
20use std::path::PathBuf;
21
22/// The standard frequency space size to use across all ML operations.
23///
24/// This covers up to C9, which is beyond the range of a standard 88-key piano (C8).
25pub const FREQUENCY_SPACE_SIZE: usize = 8192;
26
27/// The number of MIDI-indexed note bins used throughout the pipeline.
28pub const NOTE_SIGNATURE_SIZE: usize = 128;
29
30/// The number of pitch classes in a single octave (C through B).
31pub const PITCH_CLASS_COUNT: usize = 12;
32
33/// The deterministic guess vector mirrors the MIDI note signature.
34pub const DETERMINISTIC_GUESS_SIZE: usize = NOTE_SIGNATURE_SIZE;
35
36/// The standard mel space size to use across all ML operations.
37pub const MEL_SPACE_SIZE: usize = 512;
38
39/// Ensure exactly one primary sample loader feature is enabled when ML base is compiled.
40#[cfg(any(
41    all(
42        feature = "ml_loader_note_binned_convolution",
43        any(feature = "ml_loader_mel", feature = "ml_loader_frequency", feature = "ml_loader_frequency_pooled")
44    ),
45    all(feature = "ml_loader_mel", any(feature = "ml_loader_frequency", feature = "ml_loader_frequency_pooled")),
46    all(feature = "ml_loader_frequency", feature = "ml_loader_frequency_pooled"),
47))]
48compile_error!(
49    "Multiple ml_loader_* features enabled; enable exactly one of: \
50     ml_loader_note_binned_convolution, ml_loader_mel, ml_loader_frequency, ml_loader_frequency_pooled."
51);
52
53#[cfg(not(any(
54    feature = "ml_loader_note_binned_convolution",
55    feature = "ml_loader_mel",
56    feature = "ml_loader_frequency",
57    feature = "ml_loader_frequency_pooled",
58)))]
59compile_error!(
60    "No ml_loader_* feature enabled; enable exactly one of: \
61     ml_loader_note_binned_convolution, ml_loader_mel, ml_loader_frequency, ml_loader_frequency_pooled."
62);
63
64/// The base dimensionality of the sample tensor produced by `kord_item_to_sample_tensor`.
65#[cfg(feature = "ml_loader_note_binned_convolution")]
66const INPUT_BASE_SIZE: usize = NOTE_SIGNATURE_SIZE;
67
68/// The base dimensionality of the sample tensor produced by `kord_item_to_sample_tensor`.
69#[cfg(feature = "ml_loader_mel")]
70const INPUT_BASE_SIZE: usize = MEL_SPACE_SIZE;
71
72/// The base dimensionality of the sample tensor produced by `kord_item_to_sample_tensor`.
73#[cfg(feature = "ml_loader_frequency")]
74const INPUT_BASE_SIZE: usize = FREQUENCY_SPACE_SIZE;
75
76/// The frequency pooling factor applied when using the pooled loader variant.
77#[cfg(feature = "ml_loader_frequency_pooled")]
78pub const FREQUENCY_POOL_FACTOR: usize = 16;
79
80/// The dimensionality of the pooled frequency space representation.
81#[cfg(feature = "ml_loader_frequency_pooled")]
82pub const FREQUENCY_SPACE_POOLED_SIZE: usize = FREQUENCY_SPACE_SIZE / FREQUENCY_POOL_FACTOR;
83
84/// The base dimensionality of the sample tensor produced by `kord_item_to_sample_tensor`.
85#[cfg(feature = "ml_loader_frequency_pooled")]
86const INPUT_BASE_SIZE: usize = FREQUENCY_SPACE_POOLED_SIZE;
87
88/// The dimensionality of the sample tensor produced by `kord_item_to_sample_tensor`.
89#[cfg(feature = "ml_loader_include_deterministic_guess")]
90pub const INPUT_SPACE_SIZE: usize = INPUT_BASE_SIZE + DETERMINISTIC_GUESS_SIZE;
91
92/// The dimensionality of the sample tensor produced by `kord_item_to_sample_tensor`.
93#[cfg(not(feature = "ml_loader_include_deterministic_guess"))]
94pub const INPUT_SPACE_SIZE: usize = INPUT_BASE_SIZE;
95
96/// Ensure exactly one target encoding feature is enabled when ML base is compiled.
97#[cfg(not(any(feature = "ml_target_full", feature = "ml_target_folded", feature = "ml_target_folded_bass")))]
98compile_error!("No ml_target_* feature enabled; enable exactly one of: ml_target_full, ml_target_folded, ml_target_folded_bass.");
99
100#[cfg(any(
101    all(feature = "ml_target_full", feature = "ml_target_folded"),
102    all(feature = "ml_target_full", feature = "ml_target_folded_bass"),
103    all(feature = "ml_target_folded", feature = "ml_target_folded_bass"),
104))]
105compile_error!("Multiple ml_target_* features enabled; select exactly one of: ml_target_full, ml_target_folded, ml_target_folded_bass.");
106
107/// The dimensionality of the target tensor produced by `kord_item_to_target_tensor`.
108#[cfg(feature = "ml_target_full")]
109pub const TARGET_SPACE_SIZE: usize = NOTE_SIGNATURE_SIZE;
110
111/// The dimensionality of the target tensor produced by `kord_item_to_target_tensor`.
112#[cfg(feature = "ml_target_folded")]
113pub const TARGET_SPACE_SIZE: usize = PITCH_CLASS_COUNT;
114
115/// The dimensionality of the target tensor produced by `kord_item_to_target_tensor`.
116#[cfg(feature = "ml_target_folded_bass")]
117pub const TARGET_SPACE_SIZE: usize = 2 * PITCH_CLASS_COUNT;
118
119/// Backward-compatible alias for target dimensionality.
120pub const NUM_CLASSES: usize = TARGET_SPACE_SIZE;
121
122/// Offset where the folded-bass categorical head begins within the output vector.
123#[cfg(feature = "ml_target_folded_bass")]
124pub const TARGET_FOLDED_BASS_OFFSET: usize = 0;
125
126/// Offset where the folded-bass note mask begins relative to the bass categorical head.
127#[cfg(feature = "ml_target_folded_bass")]
128pub const TARGET_FOLDED_BASS_NOTE_OFFSET: usize = PITCH_CLASS_COUNT;
129
130// Training configuration.
131
132/// The training configuration used for all training, inference, and hyper parameter tuning.
133#[derive(Debug, Config)]
134pub struct TrainConfig {
135    /// The source directory for the noise assets used to generate simulated items.
136    pub noise_asset_root: String,
137    /// The directories (or the special `sim` source) that provide training samples. Must contain at least one entry.
138    pub training_sources: Vec<String>,
139    /// Optional directories (or `sim`) that provide validation samples. An empty list triggers an 80/20 split of the training pool.
140    pub validation_sources: Vec<String>,
141    /// The destination directory for the trained model.
142    pub destination: String,
143    /// The log directory for training.
144    pub log: String,
145
146    /// Simulation data set size.
147    pub simulation_size: usize,
148    /// Simulation peak radius.
149    pub simulation_peak_radius: f32,
150    /// Simulation harmonic decay.
151    pub simulation_harmonic_decay: f32,
152    /// Simulation frequency wobble.
153    pub simulation_frequency_wobble: f32,
154
155    /// The number of times to replicate captured samples when constructing datasets.
156    pub captured_oversample_factor: usize,
157
158    /// The number of Multi Head Attention (MHA) heads.
159    pub mha_heads: usize,
160    /// The dropout rate applied to attention and trunk layers.
161    pub dropout: f64,
162
163    /// The hidden size of the model's MLP trunk.
164    pub trunk_hidden_size: usize,
165
166    /// The number of epochs to train for.
167    pub model_epochs: usize,
168    /// The number of samples to use per epoch.
169    pub model_batch_size: usize,
170    /// The number of workers to use for training.
171    pub model_workers: usize,
172    /// The seed used for training.
173    pub model_seed: u64,
174
175    /// The Adam optimizer learning rate.
176    pub adam_learning_rate: f64,
177    /// The Adam optimizer weight decay.
178    pub adam_weight_decay: f32,
179    /// The Adam optimizer beta1.
180    pub adam_beta1: f32,
181    /// The Adam optimizer beta2.
182    pub adam_beta2: f32,
183    /// The Adam optimizer epsilon.
184    pub adam_epsilon: f32,
185
186    /// Suppresses the training plots.
187    pub no_plots: bool,
188}
189
190/// A single kord sample.
191///
192/// This is a single sample of a kord, which is a set of notes played together.
193#[derive(Clone, Debug)]
194pub struct KordItem {
195    /// The path to the sample.
196    pub path: PathBuf,
197    /// The frequency space of the sample.
198    pub frequency_space: [f32; FREQUENCY_SPACE_SIZE],
199    /// The label of the sample.
200    pub label: u128,
201}
202
203impl Default for KordItem {
204    fn default() -> Self {
205        Self {
206            path: PathBuf::new(),
207            frequency_space: [0.0; FREQUENCY_SPACE_SIZE],
208            label: 0,
209        }
210    }
211}
klib/ml/base/mod.rs

klib/ml/base/
mod.rs
