tiny_recursive_rs/data/mod.rs
1/// Data loading modules for TRM training
2pub mod numpy_dataset;
3
4pub use numpy_dataset::{NumpyDataset, NumpyDataLoader, DatasetMetadata};
5
6use candle_core::{Result, Tensor, Device};
7
8/// Generic data loader trait
9pub trait BatchDataLoader {
10 /// Get next batch of (input, target) tensors
11 fn next_batch(&mut self, device: &Device) -> Result<Option<(Tensor, Tensor)>>;
12
13 /// Reset loader for new epoch
14 fn reset(&mut self);
15
16 /// Get total number of batches
17 fn num_batches(&self) -> usize;
18}