A dataset of elements of type 'a. Datasets are lazy, composable, and abstract. Use creation functions to build them and transformations to modify.

Generalized dataset of tensors, parameterized over element, kind, and device

Function type for pluggable tokenizers

Built-in whitespace tokenizer

from_array arr creates a dataset from an in-memory array

from_list lst creates a dataset from a list

from_seq seq creates a dataset from a sequence

from_tensor tensor creates a dataset where each element is a slice of the first dimension

from_tensors (x, y) creates a dataset of (input, target) pairs

from_file parser path creates a dataset from a file, parsing each line with parser

from_text_file ?encoding ?chunk_size path creates a memory-mapped text dataset yielding lines as strings.

• encoding: Text encoding (default: UTF8)
• chunk_size: Size of chunks to read at once (default: 64KB) The file is memory-mapped and read lazily in chunks.

from_text_files paths creates a dataset from multiple text files. Files are processed sequentially without loading all into memory.

from_jsonl ?field path reads a JSONL file where each line is a JSON object.

• field: Extract text from this field (default: "text") Example JSONL format:

  {"text": "First document", "label": 0}
  {"text": "Second document", "label": 1}

sliding_window ~block_size ~tokenize texts creates a dataset of sliding window context/target pairs for language modeling.

• parameter block_size

  Size of the context window

• parameter tokenize

  Function to convert text to token indices

• parameter texts

  List of input texts (e.g., names for character-level modeling)

• returns

  Dataset of (context, target) tensor pairs

Creates all possible sliding windows of size block_size from the input texts, where each window predicts the next token. Automatically handles padding with a special token.

Example:

  let dataset =
    sliding_window ~block_size:3
      ~tokenize:(fun s -> encode_chars ~vocab s)
      [ "hello"; "world" ]
  (* Generates windows like: "...h" -> "e" "..he" -> "l" ".hel" -> "l"
     "hell" -> "o" etc. *)

from_csv ?separator ?text_column ?label_column ?has_header path reads CSV data.

• separator: Field separator (default: ',')
• text_column: Column index for text (default: 0)
• label_column: Optional column index for labels
• has_header: Skip first row if true (default: true)

from_text ~tokenizer path reads a text file and returns a dataset of token ID arrays. The file is read as a single document and tokenized. This is useful for language modeling tasks where you want the entire document as a sequence of tokens.

map f dataset applies function f to each element

filter pred dataset keeps only elements satisfying pred

flat_map f dataset maps and flattens nested datasets

zip ds1 ds2 pairs corresponding elements. Stops at shorter dataset.

concatenate ds1 ds2 appends ds2 after ds1

interleave datasets alternates between datasets in round-robin fashion

enumerate dataset adds indices to elements, starting from 0

tokenize tokenizer ?max_length ?padding ?truncation dataset tokenizes text data using the provided tokenizer.

• max_length: Maximum sequence length
• padding: Padding strategy
• truncation: Whether to truncate long sequences
• add_special_tokens: Add <bos>, <eos> tokens

normalize ?lowercase ?remove_punctuation ?collapse_whitespace dataset applies text normalization

batch ?drop_remainder size dataset groups tensor pairs into batches and automatically stacks them along the batch dimension.

• drop_remainder: Drop final batch if incomplete (default: false)

This is the primary batching function for ML workflows where datasets contain (input, target) tensor pairs. The tensors are automatically stacked using Rune.stack ~axis:0.

batch_map ?drop_remainder size f dataset groups elements into batches and applies function f to each batch.

This is useful for custom batching logic that can't be handled by batch or batch_array.

bucket_by_length ?boundaries ?batch_sizes length_fn dataset groups elements into buckets by length for efficient padding. Example:

  bucket_by_length ~boundaries:[ 10; 20; 30 ] ~batch_sizes:[ 32; 16; 8; 4 ]
    (fun text -> String.length text)
    dataset

Creates 4 buckets: <10, 10-20, 20-30, >30 with different batch sizes

shuffle ?rng ?buffer_size dataset randomly shuffles elements.

• rng: Random state for reproducibility (default: self-init)
• buffer_size: Size of shuffle buffer (default: 10000) Uses a buffer to shuffle without loading entire dataset in memory.

sample ?rng ?replacement n dataset randomly samples n elements

weighted_sample ?rng ~weights n dataset samples with given weights

take n dataset takes first n elements

skip n dataset skips first n elements

repeat ?count dataset repeats dataset. Infinite if count not specified.

window ?shift ?stride ?drop_remainder size dataset creates sliding windows.

• shift: How many elements to shift window (default: size)
• stride: Stride within window (default: 1) Example: window ~shift:1 3 dataset creates overlapping windows of size 3

cache ?directory dataset caches dataset elements.

• directory: Directory for file cache, in-memory if not specified

prefetch ?buffer_size dataset pre-fetches elements in background.

• buffer_size: Number of elements to prefetch (default: 2) Uses a separate thread to prepare next elements while current is processed.

parallel_map ?num_workers f dataset applies f using multiple workers.

• num_workers: Number of parallel workers (default: CPU count)

parallel_interleave ?num_workers ?block_length f dataset applies f in parallel and interleaves results

prepare ?shuffle_buffer ?batch_size ?prefetch ?cache ?drop_remainder dataset applies common preprocessing pipeline for tensor datasets: 1. Cache (if enabled) 2. Shuffle (if buffer size provided) 3. Batch with automatic tensor stacking (if batch size provided) 4. Prefetch (if prefetch count provided)

This is the primary pipeline function for ML training data.

iter f dataset applies f to each element for side effects

fold f init dataset folds over dataset elements

to_seq dataset converts to a sequence for lazy iteration

to_list dataset materializes dataset as list. Warning: loads all into memory.

to_array dataset materializes dataset as array. Warning: loads all into memory.

cardinality dataset returns the cardinality (finite length, unknown, or infinite)

element_spec dataset returns a structured description of element types

reset dataset resets the dataset to its initial state if supported. This makes it possible to iterate a dataset multiple times (e.g., across training epochs). If the dataset does not support reset, this is a no-op.

Pre-configured pipeline for text classification tasks. Returns batched token tensors ready for embedding layers.

Pre-configured pipeline for language modeling. Returns batched (input, target) tensor pairs ready for training.