Skip to main content

entrenar/config/train/batches/
loader.rs

1//! Main batch loading entry point
2
3use super::super::demo::create_demo_batches;
4use crate::config::schema::TrainSpec;
5use crate::error::Result;
6use crate::train::Batch;
7
8#[cfg(not(target_arch = "wasm32"))]
9use super::json::load_json_batches;
10#[cfg(all(not(target_arch = "wasm32"), feature = "parquet"))]
11use super::parquet::load_parquet_batches;
12
13/// Load training batches from data file using alimentar
14///
15/// Supports parquet, JSON, and CSV formats via alimentar.
16/// Falls back to demo data if the file doesn't exist (for testing).
17pub fn load_training_batches(spec: &TrainSpec) -> Result<Vec<Batch>> {
18    let data_path = &spec.data.train;
19    let batch_size = spec.data.batch_size;
20
21    // Check if data file exists
22    if !data_path.exists() {
23        eprintln!("Warning: Training data not found at '{}', using demo data", data_path.display());
24        return Ok(create_demo_batches(batch_size));
25    }
26
27    // Load data using alimentar (only on non-WASM)
28    #[cfg(not(target_arch = "wasm32"))]
29    {
30        let ext = data_path.extension().and_then(|e| e.to_str()).unwrap_or("").to_lowercase();
31
32        match ext.as_str() {
33            #[cfg(feature = "parquet")]
34            "parquet" => load_parquet_batches(data_path, batch_size),
35            #[cfg(not(feature = "parquet"))]
36            "parquet" => {
37                eprintln!(
38                    "Warning: Parquet support requires the 'parquet' feature, using demo data"
39                );
40                Ok(create_demo_batches(batch_size))
41            }
42            "json" => load_json_batches(data_path, batch_size),
43            _ => {
44                eprintln!("Warning: Unsupported data format '{ext}', using demo data");
45                Ok(create_demo_batches(batch_size))
46            }
47        }
48    }
49
50    #[cfg(target_arch = "wasm32")]
51    {
52        eprintln!("Warning: Data loading not available in WASM, using demo data");
53        Ok(create_demo_batches(batch_size))
54    }
55}