Data API

This section documents the data loading and processing modules.

Core Dataset

Structured Data Support

JSONDataset

Load data from JSON files with flexible key-based access:

from jepa.data import JSONDataset

# Single array with temporal offset
dataset = JSONDataset(
    json_path="data.json",
    data_key="timeseries",
    time_offset=1
)

# Separate arrays format
dataset = JSONDataset(
    json_path="data.json",
    data_t_key="context_data",
    data_t1_key="target_data"
)

CSVDataset

Load data from CSV files with column-based organization:

from jepa.data import CSVDataset

# Single columns with temporal offset
dataset = CSVDataset(
    csv_path="data.csv", 
    data_columns=["feature1", "feature2", "feature3"],
    time_offset=1
)

PickleDataset

Load data from Python pickle files:

from jepa.data import PickleDataset

dataset = PickleDataset(
    pickle_path="data.pkl",
    data_key="sequences",
    time_offset=1
)

Transforms

Base Transform Classes

Image Transforms

Text Transforms

Time Series Transforms

Structured Data Transforms

Transform structured data with preprocessing pipeline:

from jepa.data.transforms import StructuredDataTransform

transform = StructuredDataTransform([
    ("normalize", {"method": "z_score"}),
    ("handle_missing", {"strategy": "interpolate"}),
    ("feature_selection", {"top_k": 20})
])

Data Utilities

Core Utilities

Data Loading Helpers

Preprocessing Utilities

Factory Functions

Dataset Creation

Create datasets with automatic format detection:

from jepa.data import create_dataset

# Automatic format detection
dataset = create_dataset(
    data_path="data.json",  # Format inferred from extension
    data_key="features",
    temporal_offset=1
)

# Explicit format specification
dataset = create_dataset(
    data_type="csv",
    data_path="measurements.csv",
    data_columns=["sensor_1", "sensor_2"],
    temporal_offset=2
)

DataLoader Creation

Create optimized DataLoaders for JEPA training:

from jepa.data.utils import create_jepa_dataloader

dataloader = create_jepa_dataloader(
    dataset,
    batch_size=64,
    shuffle=True,
    num_workers=4,
    pin_memory=True
)

Hugging Face Integration

HF Dataset Wrapper

Use Hugging Face datasets with JEPA:

from datasets import load_dataset
from jepa.data.hf_compatibility import HFDatasetWrapper

# Load HF dataset
hf_dataset = load_dataset("imagenet-1k", split="train")

# Wrap for JEPA
dataset = HFDatasetWrapper(
    hf_dataset,
    image_column="image",
    transform=JEPATransform()
)

Custom HF Adapters

Configuration

Data Configuration Classes

Configure data loading through YAML or programmatically:

data:
  dataset_type: "csv"
  dataset_path: "data/train.csv"
  data_columns: ["feature1", "feature2", "feature3"]
  temporal_offset: 1
  batch_size: 64
  num_workers: 4
  transforms:
    normalize: true
    augment: true

Examples

Basic Data Loading

from jepa.data import JSONDataset, create_dataloader

# Create dataset
dataset = JSONDataset(
    json_path="timeseries.json",
    data_key="measurements",
    temporal_offset=1
)

# Create dataloader
dataloader = create_dataloader(
    dataset,
    batch_size=32,
    shuffle=True
)

# Use in training
for batch in dataloader:
    context, target = batch
    # Process batch...

Custom Dataset Implementation

from jepa.data.dataset import BaseJEPADataset

class CustomDataset(BaseJEPADataset):
    def __init__(self, data_path, **kwargs):
        super().__init__(**kwargs)
        self.data = self.load_custom_data(data_path)
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        sample = self.data[idx]
        if self.transform:
            sample = self.transform(sample)
        return sample

Advanced Transform Pipeline

from jepa.data.transforms import Compose, Normalize, Augment

# Create transform pipeline
transform = Compose([
    Normalize(method="z_score"),
    Augment(noise_std=0.1),
    JEPATransform(mask_ratio=0.15)
])

# Apply to dataset
dataset = CSVDataset(
    csv_path="data.csv",
    data_columns=["x", "y", "z"],
    temporal_offset=1,
    transform=transform
)

For more examples and detailed usage, see the Data Loading Guide and Examples.