Skip to main content

entrenar/config/train/batches/
loader.rs

1//! Main batch loading entry point
2
3use crate::config::schema::TrainSpec;
4use crate::error::{Error, Result};
5use crate::train::Batch;
6
7#[cfg(not(target_arch = "wasm32"))]
8use super::json::load_json_batches;
9#[cfg(all(not(target_arch = "wasm32"), feature = "parquet"))]
10use super::parquet::load_parquet_batches;
11
12/// The documented on-disk schema for `--task pretrain` JSON training data.
13///
14/// Quoted verbatim in every load failure so the user is never left guessing
15/// what the loader wanted.
16pub(crate) const JSON_SCHEMA_HINT: &str = "expected JSON of the form \
17     {\"examples\":[{\"input\":[f32,..],\"target\":[f32,..]}, ..]} \
18     or a bare array [{\"input\":[..],\"target\":[..]}, ..]";
19
20/// Load training batches from the dataset named by the config.
21///
22/// Supported formats: JSON (see [`JSON_SCHEMA_HINT`]) and, when the `parquet`
23/// feature is enabled, Parquet via alimentar.
24///
25/// # Errors
26///
27/// Returns [`Error::ConfigError`] when the dataset is missing, is in a format
28/// this build cannot read, or cannot be parsed. It NEVER substitutes synthetic
29/// data for a dataset it failed to read: a training run that silently trains on
30/// fabricated examples and reports success is worse than one that refuses to
31/// start.
32pub fn load_training_batches(spec: &TrainSpec) -> Result<Vec<Batch>> {
33    let data_path = &spec.data.train;
34    let batch_size = spec.data.batch_size;
35
36    // Check if data file exists
37    if !data_path.exists() {
38        return Err(Error::ConfigError(format!(
39            "Training data not found at '{}'. Training cannot proceed without it.",
40            data_path.display()
41        )));
42    }
43
44    // Load data using alimentar (only on non-WASM)
45    #[cfg(not(target_arch = "wasm32"))]
46    {
47        let ext = data_path.extension().and_then(|e| e.to_str()).unwrap_or("").to_lowercase();
48
49        match ext.as_str() {
50            #[cfg(feature = "parquet")]
51            "parquet" => load_parquet_batches(data_path, batch_size),
52            #[cfg(not(feature = "parquet"))]
53            "parquet" => Err(Error::ConfigError(format!(
54                "Cannot read Parquet training data '{}': this build lacks the 'parquet' feature. \
55                 Rebuild with --features parquet, or convert the dataset to JSON ({JSON_SCHEMA_HINT}).",
56                data_path.display()
57            ))),
58            "json" => load_json_batches(data_path, batch_size),
59            _ => Err(Error::ConfigError(format!(
60                "Unsupported training data format '{ext}' for '{}'. Supported: {}. \
61                 Convert the dataset to JSON — {JSON_SCHEMA_HINT}.",
62                data_path.display(),
63                supported_formats()
64            ))),
65        }
66    }
67
68    #[cfg(target_arch = "wasm32")]
69    {
70        let _ = batch_size;
71        Err(Error::ConfigError(
72            "Data loading is not available in WASM builds; training cannot proceed.".to_string(),
73        ))
74    }
75}
76
77/// Human-readable list of the dataset formats this build can actually read.
78#[cfg(not(target_arch = "wasm32"))]
79fn supported_formats() -> &'static str {
80    #[cfg(feature = "parquet")]
81    {
82        "json, parquet"
83    }
84    #[cfg(not(feature = "parquet"))]
85    {
86        "json"
87    }
88}