HuggingFace Datasets Tutorial

Metadata	Value
Level	Intermediate
Runtime	~30 min
Prerequisites	HuggingFace Quick Reference, Pipeline Tutorial
Format	Python + Jupyter

Overview

This tutorial provides a detailed guide to using HuggingFace Datasets with Datarax. You'll learn to work with different data modalities, configure advanced options like field filtering and shuffle buffers, and build production-ready training pipelines.

What You'll Learn

1. Load different dataset types (images, text, audio) from HuggingFace Hub
2. Configure field filtering with include_keys and exclude_keys
3. Set up shuffling with proper buffer configuration for streaming datasets
4. Build complete training pipelines with preprocessing and augmentation
5. Handle streaming vs downloaded modes effectively
6. Use named RNG streams for reproducible data loading

Coming from PyTorch?

If you're familiar with PyTorch's dataset ecosystem, here's how Datarax + HuggingFace compares:

PyTorch	Datarax
datasets.load_dataset('mnist', split='train')	HFEagerSource(HFEagerConfig(name='mnist', split='train'))
DataLoader(shuffle=True, num_workers=4)	HFEagerSource with shuffle=True, seed=N
datasets.set_format('torch')	Automatic JAX array conversion
Manual field selection in __getitem__	include_keys / exclude_keys in config
IterableDataset for streaming	from_hf(name, split, streaming=True) / HFStreamingSource

Key difference: Datarax uses JAX arrays and provides declarative configuration instead of imperative code.

Coming from TensorFlow?

TensorFlow tf.data	Datarax
tfds.load('mnist', split='train')	HFEagerSource(HFEagerConfig(name='mnist', split='train'))
dataset.shuffle(buffer_size=1000)	shuffle=True in config (O(1) index shuffle, no buffer)
dataset.take(1000)	split='train[:1000]' syntax
dataset.skip(1000)	split='train[1000:]' syntax
dataset.map(fn).filter(pred)	Chain operators by passing them in the stages=[...] list

Files

• Python Script: examples/integration/huggingface/02_hf_tutorial.py
• Jupyter Notebook: examples/integration/huggingface/02_hf_tutorial.ipynb

Quick Start

# Run the Python script
python examples/integration/huggingface/02_hf_tutorial.py

# Or launch the Jupyter notebook
jupyter lab examples/integration/huggingface/02_hf_tutorial.ipynb

Part 1: Understanding HFEagerSource Configuration

HFEagerConfig provides extensive options for loading HuggingFace datasets.

Note: You can also use the factory function from_hf(name, split, ...) which auto-selects between eager and streaming modes.

Key Configuration Parameters

Parameter	Description	Default	Example
name	Dataset identifier on HF Hub	Required	"mnist", "stanfordnlp/imdb"
split	Which split to use	Required	"train", "test[:1000]"
shuffle	Enable shuffling	False	True for training
include_keys	Only include these fields	None	{"image", "label"}
exclude_keys	Exclude these fields	None	{"metadata", "id"}
seed	Integer seed for shuffling	42	0, 123

Auto-derived internals: stochastic and stream_name are not user-set knobs. When shuffle=True, the source uses the RNG stream named "shuffle" (with stochastic=True); otherwise stream_name defaults to None. Eager shuffling is an O(1) Feistel index shuffle, so there is no shuffle buffer in eager mode.

Basic Configuration Example

import jax
import jax.numpy as jnp
from flax import nnx

from datarax.pipeline import Pipeline
from datarax.operators import ElementOperator, ElementOperatorConfig
from datarax.sources import HFEagerConfig, HFEagerSource

# Basic configuration for MNIST
basic_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:1000]",  # Load first 1000 samples
)

basic_source = HFEagerSource(basic_config, rngs=nnx.Rngs(0))
print(f"Basic MNIST source: {len(basic_source)} samples")

Terminal Output:

Basic MNIST source: 1000 samples

Part 2: Field Filtering

Use include_keys or exclude_keys to control which fields are loaded and returned.

Benefits of Field Filtering

• Reduces memory usage by excluding unnecessary fields
• Simplifies downstream processing
• Faster iteration when working with large metadata fields
• Cleaner batch dictionaries

Include Keys Example

# Include only specific fields
filtered_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:500]",
    include_keys={"image", "label"},  # Only return these fields
)

filtered_source = HFEagerSource(filtered_config, rngs=nnx.Rngs(1))

# Check what fields are available
pipeline = Pipeline(source=filtered_source, stages=[], batch_size=1, rngs=nnx.Rngs(0))
batch = next(iter(pipeline))

print("Filtered fields:")
for key in batch.keys():
    print(f"  - {key}: shape={batch[key].shape}")

Terminal Output:

Filtered fields:
  - image: shape=(1, 28, 28, 1)
  - label: shape=(1,)

Exclude Keys Example

# Exclude metadata fields
exclude_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:500]",
    exclude_keys={"id"},  # Exclude ID field
)

exclude_source = HFEagerSource(exclude_config, rngs=nnx.Rngs(2))

Terminal Output:

Excluded 'id' field from dataset
Remaining fields: image, label

Part 3: Shuffling Configuration

Shuffling is essential for training ML models. In eager mode, HFEagerSource shuffles with an O(1) Feistel index shuffle - there is no shuffle buffer, so you only need shuffle=True and an integer seed.

Shuffle Modes

Mode	When to Use	Configuration
No shuffle	Testing, evaluation	shuffle=False
Index shuffle	Eager (downloaded) datasets	shuffle=True, seed=N
Buffer shuffle	Streaming datasets (large)	HFStreamingConfig(shuffle=True, shuffle_buffer_size=N)

Eager Shuffle Example

# Configure shuffling with a reproducible seed
shuffle_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:2000]",
    shuffle=True,
    seed=42,  # Integer seed for Grain's index_shuffle
)

# Create source with explicit RNG for reproducibility
shuffle_source = HFEagerSource(
    shuffle_config,
    rngs=nnx.Rngs(42),
)

print("Shuffle configuration:")
print(f"  Seed: {shuffle_config.seed}")

Terminal Output:

Shuffle configuration:
  Seed: 42

Part 4: Streaming vs Downloaded Mode

Downloaded / Eager Mode (HFEagerSource)

• Full dataset downloaded and cached locally
• Random access to any sample
• Faster iteration after initial download
• Requires disk space

# Downloaded (eager) mode
downloaded_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:1000]",
)
downloaded_source = HFEagerSource(downloaded_config, rngs=nnx.Rngs(0))

print(f"Downloaded mode length: {len(downloaded_source)}")

Terminal Output:

Downloaded mode length: 1000

Streaming Mode (HFStreamingSource)

• Data loaded on-the-fly from HuggingFace servers
• No disk storage required
• Ideal for large datasets (ImageNet, Common Crawl)
• Cannot seek to specific indices
• Dataset length may not be available
• Shuffling uses a buffer (shuffle_buffer_size) rather than an index shuffle

from datarax.sources import HFStreamingConfig, HFStreamingSource

# Streaming mode with buffer-based shuffle
streaming_config = HFStreamingConfig(
    name="ylecun/mnist",
    split="train",
    streaming=True,
    shuffle=True,
    shuffle_buffer_size=1000,
)
streaming_source = HFStreamingSource(streaming_config, rngs=nnx.Rngs(0))

try:
    print(f"Streaming mode length: {len(streaming_source)}")
except (NotImplementedError, TypeError):
    print("Streaming mode length: N/A (not available in streaming)")

Terminal Output:

Streaming mode length: N/A (not available in streaming)

Tip: The from_hf(name, split, ...) factory auto-selects HFEagerSource for datasets under ~1GB and HFStreamingSource for larger ones. Pass streaming=True to force streaming regardless of size.

Mode Comparison Table

Aspect	Streaming	Downloaded
Disk usage	Minimal	Full dataset
First iteration	Immediate	After download
Subsequent iterations	Network dependent	Fast (local)
Random access	No	Yes
Length available	Usually no	Yes
Best for	>10GB datasets	<1GB datasets

Part 5: Building Complete Training Pipeline

Combine HFEagerSource with operators for a production-ready pipeline.

Define Preprocessing Operators

def normalize_image(element, key=None):
    """Normalize image to [0, 1] and ensure proper shape."""
    image = element.data.get("image")
    if image is not None and hasattr(image, "dtype"):
        # Normalize to [0, 1]
        normalized = image.astype(jnp.float32) / 255.0
        # Add channel dimension if needed (for grayscale)
        if normalized.ndim == 2:
            normalized = normalized[..., None]
        return element.update_data({"image": normalized})
    return element

def random_flip(element, key):
    """Randomly flip image horizontally."""
    flip_key, _ = jax.random.split(key)
    should_flip = jax.random.bernoulli(flip_key, 0.5)

    image = element.data.get("image")
    if image is not None:
        flipped = jax.lax.cond(
            should_flip,
            lambda x: jnp.flip(x, axis=1),  # Flip width axis
            lambda x: x,
            image,
        )
        return element.update_data({"image": flipped})
    return element

# Create operators
normalizer = ElementOperator(
    ElementOperatorConfig(stochastic=False),
    fn=normalize_image,
    rngs=nnx.Rngs(0),
)

flipper = ElementOperator(
    ElementOperatorConfig(stochastic=True, stream_name="flip"),
    fn=random_flip,
    rngs=nnx.Rngs(flip=42),
)

Terminal Output:

Created operators: normalizer (deterministic), flipper (stochastic)

Build Complete Pipeline

from datarax.operators.composite_operator import (
    CompositeOperatorConfig,
    CompositeOperatorModule,
    CompositionStrategy,
)

# Create composite augmentation
augmentation = CompositeOperatorModule(
    CompositeOperatorConfig(
        strategy=CompositionStrategy.SEQUENTIAL,
        operators=[normalizer, flipper],
        stochastic=True,
        stream_name="augment",
    ),
    rngs=nnx.Rngs(augment=999),
)

# Build the complete pipeline
train_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:5000]",
    shuffle=True,
    include_keys={"image", "label"},
    seed=42,
)

train_source = HFEagerSource(train_config, rngs=nnx.Rngs(0))

# Chain: Source -> Augmentation -> Output
training_pipeline = Pipeline(source=train_source, stages=[augmentation], batch_size=64, rngs=nnx.Rngs(0))

print("Training pipeline:")
print("  HFEagerSource(mnist) -> Normalize -> RandomFlip -> Output")
print("  Batch size: 64")

Terminal Output:

Training pipeline:
  HFEagerSource(mnist) -> Normalize -> RandomFlip -> Output
  Batch size: 64

Process Training Data

print("\nProcessing training batches:")
stats = {"batches": 0, "samples": 0}

for i, batch in enumerate(training_pipeline):
    if i >= 5:  # Process 5 batches for demo
        break

    image_batch = batch["image"]
    label_batch = batch["label"]

    stats["batches"] += 1
    stats["samples"] += image_batch.shape[0]

    if i == 0:  # Print details for first batch
        print(f"Batch {i}:")
        print(f"  Image: shape={image_batch.shape}, dtype={image_batch.dtype}")
        img_min, img_max = float(image_batch.min()), float(image_batch.max())
        print(f"  Image range: [{img_min:.3f}, {img_max:.3f}]")
        print(f"  Label: shape={label_batch.shape}")

print(f"\nProcessed {stats['batches']} batches, {stats['samples']} samples")

Terminal Output:

Processing training batches:
Batch 0:
  Image: shape=(64, 28, 28, 1), dtype=float32
  Image range: [0.000, 1.000]
  Label: shape=(64,)

Processed 5 batches, 320 samples

Part 6: Working with Different Datasets

HuggingFace Hub hosts thousands of datasets across different modalities.

Common Dataset Examples

Dataset	Type	Split Syntax	Use Case
mnist	Image	split="train"	Computer vision basics
cifar10	Image	split="train"	Image classification
imagenet-1k	Image	split="train"	Large-scale vision
stanfordnlp/imdb	Text	split="train"	Sentiment analysis
squad	QA	split="train"	Question answering
librispeech_asr	Audio	split="train.clean.100"	Speech recognition

Split Syntax Examples

print("Split syntax examples:")
print("  'train' - Full training set")
print("  'train[:1000]' - First 1000 samples")
print("  'train[1000:2000]' - Samples 1000-2000")
print("  'train[:10%]' - First 10% of data")
print("  'train[10%:20%]' - Second 10% of data")
print("  'train+test' - Combined splits")

Terminal Output:

Split syntax examples:
  'train' - Full training set
  'train[:1000]' - First 1000 samples
  'train[1000:2000]' - Samples 1000-2000
  'train[:10%]' - First 10% of data
  'train[10%:20%]' - Second 10% of data
  'train+test' - Combined splits

Dataset Discovery

# List available datasets programmatically
from huggingface_hub import list_datasets

datasets_list = list(list_datasets(limit=100))
print(f"Datasets fetched: {len(datasets_list)}")
print(f"Example datasets: {datasets_list[:5]}")

# Get dataset info
from datasets import load_dataset_builder

builder = load_dataset_builder("mnist")
print(f"\nMNIST info:")
print(f"  Description: {builder.info.description[:100]}...")
print(f"  Features: {builder.info.features}")

Terminal Output:

Datasets fetched: 100
Example datasets: ['mnist', 'cifar10', 'imdb', 'squad', 'glue']

MNIST info:
  Description: The MNIST database of handwritten digits...
  Features: {'image': Image(shape=(28, 28, 1), dtype=uint8), 'label': ClassLabel(num_classes=10)}

Architecture Diagram

flowchart TB
    subgraph HF["HuggingFace Hub"]
        Hub[Dataset Repository<br/>85,000+ datasets]
    end

    subgraph Config["Configuration"]
        Cfg[HFEagerConfig<br/>name, split, shuffle<br/>include/exclude keys]
    end

    subgraph Source["HFEagerSource"]
        Stream{Streaming?}
        Download[Download & Cache]
        StreamLoad[Stream from Hub]
    end

    subgraph Pipeline["Datarax Pipeline"]
        Ops[Operators<br/>Normalize, Augment, etc.]
        Batch[Batching]
    end

    subgraph Output["Output"]
        JAX[JAX Arrays<br/>Ready for training]
    end

    Hub --> Cfg
    Cfg --> Stream
    Stream -->|No| Download --> Ops
    Stream -->|Yes| StreamLoad --> Ops
    Ops --> Batch --> JAX

    style HF fill:#e1f5ff
    style Config fill:#fff4e1
    style Source fill:#ffe1e1
    style Pipeline fill:#f0ffe1
    style Output fill:#e1ffe1

Results Summary

Configuration Best Practices

Feature	Recommendation	Rationale
Large datasets	HFStreamingSource (or from_hf(..., streaming=True))	Avoid memory/disk issues
Training	shuffle=True with a fixed seed	Essential for SGD convergence
Streaming shuffle	shuffle_buffer_size on HFStreamingConfig	Better shuffle quality when streaming
Field filtering	Use include_keys	Reduce memory overhead
Reproducibility	Fixed integer seed	Deterministic index shuffle
Development	split="train[:1000]"	Fast iteration

Performance Characteristics

Operation	Streaming	Downloaded
First batch latency	~2-5s	~0.1s
Throughput	Network limited	Disk limited
Memory overhead	Minimal	Full dataset
Reproducibility	Buffer-based	Perfect

Common Patterns

# Pattern 1: Development (small subset, fast iteration)
dev_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train[:100]",
    shuffle=False,
)

# Pattern 2: Training (full data, shuffled)
train_config = HFEagerConfig(
    name="ylecun/mnist",
    split="train",
    shuffle=True,
    seed=42,
)

# Pattern 3: Large dataset streaming (forces HFStreamingSource)
large_source = from_hf(
    "imagenet-1k",
    "train",
    streaming=True,
    shuffle=True,
    rngs=nnx.Rngs(0),
)

# Pattern 4: Evaluation (deterministic, no shuffle)
eval_config = HFEagerConfig(
    name="ylecun/mnist",
    split="test",
    shuffle=False,
)

Next Steps

• Image augmentation: Operators Tutorial for advanced transformations
• TFDS alternative: TFDS Integration for TensorFlow Datasets
• Text processing: IMDB Quick Reference for NLP workflows
• Distributed training: Sharding Guide for multi-device training
• HuggingFace Hub: Browse datasets at https://huggingface.co/datasets
• API Reference: HFEagerSource Documentation