cursus.processing.datasets

Dataset classes for PyTorch training pipelines.

Provides both regular (map-style) and streaming (iterable-style) datasets with identical pipeline injection APIs.

class PipelineDataset(*args, **kwargs)

Bases: Dataset

Custom dataset for multimodal input supporting text, tabular, categorical, and Parquet/CSV/TSV file formats with per-column processing pipelines.

add_pipeline(field_name, processor_pipeline)

Adds a processing pipeline for a specified field. The pipeline is built by composing Processors via the >> operator. For example, for the text field ‘dialogue’, you might have:

pipeline = (HTMLNormalizerProcessor() >> EmojiRemoverProcessor() >>

TextNormalizationProcessor() >> DialogueSplitterProcessor() >> DialogueChunkerProcessor(tokenizer, max_tokens=512) >> TokenizationProcessor(tokenizer, add_special_tokens=True))

which you then add via:

dataset.add_pipeline(“dialogue”, pipeline)

fill_missing_value(**kwargs)

Ensure all fields are numeric or categorical with default fill values.

load_data(**kwargs)

Load data from file (CSV, TSV, or Parquet).

load_dataframe(dataframe)

Load data directly from a provided DataFrame.

set_cat_field_list(cat_field_list)

set_full_field_list(full_field_list)

set_label_field_name(label_name)

set_text_field_name(text_name)

class PipelineIterableDataset(*args, **kwargs)

Bases: IterableDataset

Streaming dataset for multimodal input with distributed training support.

Memory-efficient alternative to PipelineDataset that loads data incrementally from multiple shard files. Maintains the same pipeline injection API for backward compatibility.

Distributed Training Support:

The dataset implements two-tier sharding for distributed training:

1. Rank-based sharding: Shards are distributed across GPU ranks using round-robin assignment (rank 0 gets shards [0, world_size, 2*world_size, …]). This ensures each GPU processes unique data in FSDP/DDP training.

2. Worker-based sharding: Within each rank, shards are further distributed across DataLoader workers for parallel loading.

Example Usage:

Single GPU training:

>>> dataset = PipelineIterableDataset(
...     config=config,
...     file_dir="/data/train",
... )
>>> loader = DataLoader(dataset, batch_size=32, num_workers=4)

Distributed training (FSDP):

>>> # Same code! Distribution happens automatically
>>> dataset = PipelineIterableDataset(
...     config=config,
...     file_dir="/data/train",
... )
>>> loader = DataLoader(dataset, batch_size=32, num_workers=4)
>>>
>>> trainer = pl.Trainer(strategy=FSDPStrategy(), devices=8)
>>> trainer.fit(model, loader)

Epoch-aware shuffling:

>>> for epoch in range(num_epochs):
...     dataset.set_epoch(epoch)  # Important for deterministic shuffling
...     for batch in loader:
...         # Training step

config

Configuration dictionary (same as PipelineDataset)

processor_pipelines

Dictionary mapping field names to Processor pipelines

shard_files

List of shard file paths to stream through

shuffle_shards

Whether to shuffle shard order per epoch

Key Differences from PipelineDataset:

• Inherits from IterableDataset (not Dataset)

• Implements __iter__() instead of __getitem__()

• Loads shards incrementally (not all at once)

• No __len__() by default (optional estimate available)

• Automatic multi-GPU and multi-worker shard distribution

add_pipeline(field_name, processor_pipeline)

Add a processing pipeline for a specified field.

IDENTICAL API to PipelineDataset.add_pipeline() for drop-in compatibility.

The pipeline is built by composing Processors via the >> operator. For example:

pipeline = (HTMLNormalizerProcessor() >>

EmojiRemoverProcessor() >> TextNormalizationProcessor() >> DialogueSplitterProcessor() >> DialogueChunkerProcessor(tokenizer, max_tokens=512) >> TokenizationProcessor(tokenizer))

Parameters:

• field_name (str) – Name of the field to process

• processor_pipeline (Processor) – Processor or ComposedProcessor to apply

Raises:

TypeError – If arguments are not of expected types

fill_missing_value(**kwargs)

Configure missing value handling rules.

IDENTICAL API to PipelineDataset.fill_missing_value() for compatibility.

Note: For IterableDataset, this sets rules to apply during iteration. The actual filling happens in _postprocess_dataframe() when each shard is loaded.

Parameters:

**kwargs – Configuration updates (label_name, cat_field_list, etc.)

get_shard_distribution_info()

Get diagnostic information about shard distribution.

Returns:

• total_shards: Total number of shards

• shards_per_rank: Number of shards assigned to this rank

• shards_per_worker: Number of shards per worker

• rank: Current rank

• world_size: Total number of ranks

• worker_id: Current worker ID (if available)

• num_workers: Total number of workers per rank

• assigned_shards: List of shard files assigned to this rank/worker

Return type:

Dict containing

set_cat_field_list(cat_field_list)

Set categorical field list.

set_epoch(epoch)

Set current epoch for deterministic shuffling.

Should be called at the start of each epoch, similar to DistributedSampler.set_epoch().

Parameters:

epoch (int) – Current epoch number

Example

>>> dataset.set_epoch(epoch)
>>> for batch in dataloader:
>>>     # Training step

set_full_field_list(full_field_list)

Set full field list.

set_label_field_name(label_name)

Set label field name(s).

set_text_field_name(text_name)

Set text field name(s).

Modules

pipeline_datasets
pipeline_iterable_datasets	Streaming dataset implementation using PyTorch's IterableDataset.