cursus.processing.datasets.pipeline_iterable_datasets

Streaming dataset implementation using PyTorch’s IterableDataset.

This module provides PipelineIterableDataset, a memory-efficient streaming alternative to PipelineDataset that loads data incrementally from shards.

Key Features: - Fixed memory usage (loads one shard at a time) - Multi-GPU/multi-worker support (automatic shard distribution) - Same pipeline injection API as PipelineDataset - Drop-in replacement with minimal code changes

Example

>>> from processing.datasets.pipeline_iterable_datasets import PipelineIterableDataset
>>>
>>> # Create streaming dataset
>>> dataset = PipelineIterableDataset(
...     config=config,
...     file_dir="/data/train",  # Directory with part-*.parquet shards
... )
>>>
>>> # Add pipelines (same API as PipelineDataset)
>>> dataset.add_pipeline("dialogue", text_pipeline)
>>> dataset.add_pipeline("customer_id", categorical_pipeline)
>>>
>>> # Use with DataLoader (same as regular dataset)
>>> loader = DataLoader(dataset, batch_size=32, collate_fn=collate_batch)

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

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

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

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.)

set_text_field_name(text_name)

Set text field name(s).

set_label_field_name(label_name)

Set label field name(s).

set_cat_field_list(cat_field_list)

Set categorical field list.

set_full_field_list(full_field_list)

Set full field list.