import math
from collections.abc import Iterator
from typing import Any
import polars as pl
import torch
import torch.distributed as dist
from loguru import logger
from torch.utils.data import IterableDataset
from sequifier.helpers import (
PANDAS_TO_TORCH_TYPES,
configured_window_stride,
numpy_storage_to_pytorch,
read_data,
resolve_window_sampling_plan,
)
from sequifier.io.batch import SequifierBatch
from sequifier.io.config import global_training
from sequifier.io.iteration_state import (
read_shared_int,
resolve_resume_worker,
shared_int,
skip_samples_for_batches,
write_shared_int,
)
from sequifier.io.sample_order import (
SampleOrderPlan,
curriculum_positions_from_parquet,
curriculum_sample_positions,
logical_sample_positions,
)
from sequifier.io.window_sampling import (
build_window_batch,
target_valid_from_offsets,
validate_split_bounds_available,
)
from sequifier.typechecking import beartype
[docs]class SequifierDatasetFromFile(IterableDataset):
"""Eager single-file dataset yielding pre-collated batches."""
@beartype
def __init__(self, data_path: str, config: Any, shuffle: bool = True):
super().__init__()
self.config = config
self.batch_size = global_training(config).batch_size
self.shuffle = shuffle
self._epoch_state = shared_int(0)
self._start_batch_state = shared_int(0)
all_columns = sorted(list(set(config.input_columns + config.target_columns)))
logger.info(f"Loading dataset into memory from '{data_path}'...")
data_df = (
read_data(data_path, config.read_format)
if config.read_format != "pt"
else None
)
self.depth_valid_masks = {}
self.start_item_positions = None
self.split_start_item_positions = None
self.split_end_item_positions = None
stored_sample_positions = (
curriculum_positions_from_parquet(config, data_df, str(data_path))
if data_df is not None
else None
)
curriculum_columns: tuple[str, ...] = ()
column_data_types = {
col: PANDAS_TO_TORCH_TYPES[config.column_data_types[col]]
for col in config.column_data_types
}
sampling_plan = resolve_window_sampling_plan(
config.storage_layout,
config.window_view,
configured_window_stride(config),
)
if config.read_format == "pt":
from sequifier.io.pt_payload import load_pt_payload
part = getattr(config, "part", None)
metadata = part.metadata if part is not None else None
payload = load_pt_payload(
data_path,
layouts=metadata.depth_layouts
if metadata is not None
else config.depth_layouts,
n_classes=metadata.n_classes
if metadata is not None
else config.n_classes,
)
all_tensors = {column: payload.sequences[column] for column in all_columns}
left_pad_lengths = payload.left_pad_lengths
stored_sample_positions = payload.sample_positions
curriculum_columns = payload.curriculum_columns
self.depth_valid_masks = {
name: payload.depth_valid_masks[name]
for name in config.depth_layouts.root
}
self.start_item_positions = payload.start_item_positions
self.split_start_item_positions = payload.split_start_item_positions
self.split_end_item_positions = payload.split_end_item_positions
else:
assert data_df is not None
all_tensors, left_pad_lengths = numpy_storage_to_pytorch(
data=data_df,
column_data_types=column_data_types,
all_columns=all_columns,
window_length=config.storage_layout.window_length,
)
if {
"startItemPosition",
"splitStartItemPosition",
"splitEndItemPosition",
} <= set(data_df.columns):
positions = (
data_df.group_by(["sequenceId", "subsequenceId"])
.agg(
pl.col("startItemPosition").first(),
pl.col("splitStartItemPosition").first(),
pl.col("splitEndItemPosition").first(),
)
.sort(["sequenceId", "subsequenceId"])
)
self.start_item_positions = torch.tensor(
positions["startItemPosition"].to_numpy(), dtype=torch.int64
)
self.split_start_item_positions = torch.tensor(
positions["splitStartItemPosition"].to_numpy(), dtype=torch.int64
)
self.split_end_item_positions = torch.tensor(
positions["splitEndItemPosition"].to_numpy(), dtype=torch.int64
)
validate_split_bounds_available(
self.start_item_positions,
self.split_start_item_positions,
self.split_end_item_positions,
str(data_path),
)
assert self.start_item_positions is not None
assert self.split_start_item_positions is not None
self.sample_index = sampling_plan.build_index(
left_pad_lengths,
target_valid_from_offsets(
left_pad_lengths,
self.start_item_positions,
self.split_start_item_positions,
),
)
self.n_samples = len(self.sample_index)
if self.n_samples == 0:
raise ValueError("No usable model windows were found in the dataset.")
self.sample_order = SampleOrderPlan.build(
self.n_samples,
logical_sample_positions(
curriculum_sample_positions(
config,
stored_sample_positions,
str(data_path),
curriculum_columns,
),
self.sample_index,
),
)
del data_df
self.sequences = all_tensors
for mask in self.depth_valid_masks.values():
mask.share_memory_()
if config.device.startswith("cuda"):
for key in self.sequences:
self.sequences[key] = self.sequences[key].pin_memory()
logger.info(f"Dataset loaded with {self.n_samples} samples.")
[docs] @beartype
def set_epoch(self, epoch: int):
"""Set the shuffle epoch."""
write_shared_int(self._epoch_state, epoch)
[docs] @beartype
def set_start_batch(self, start_batch: int):
"""Set the first global batch to yield on the next iteration."""
write_shared_int(self._start_batch_state, start_batch)
@beartype
def __len__(self) -> int:
num_workers = max(1, global_training(self.config).num_workers)
total_batches = 0
for worker_id in range(num_workers):
worker_samples = self.n_samples // num_workers + (
1 if worker_id < self.n_samples % num_workers else 0
)
total_batches += math.ceil(worker_samples / self.batch_size)
return total_batches
@beartype
def __iter__(
self,
) -> Iterator[SequifierBatch]:
worker_info = torch.utils.data.get_worker_info()
world_size = dist.get_world_size() if dist.is_initialized() else 1
rank = dist.get_rank() if dist.is_initialized() else 0
if worker_info is None:
# Single-process data loading
physical_worker_id = 0
num_workers = 1
else:
# Multi-process data loading
physical_worker_id = worker_info.id
num_workers = worker_info.num_workers
epoch = read_shared_int(self._epoch_state)
start_batch = read_shared_int(self._start_batch_state)
indices = self.sample_order.indices_for_epoch(
seed=self.config.seed, epoch=epoch, shuffle=self.shuffle
)
indices_for_rank = indices[rank::world_size]
worker_batch_counts = [
math.ceil(len(indices_for_rank[i::num_workers]) / self.batch_size)
for i in range(num_workers)
]
worker_id, skip_batches = resolve_resume_worker(
start_batch,
physical_worker_id,
num_workers,
worker_batch_counts,
)
indices_for_worker = indices_for_rank[worker_id::num_workers]
skipped_samples = skip_samples_for_batches(
skip_batches, self.batch_size, len(indices_for_worker)
)
indices_for_worker = indices_for_worker[skipped_samples:]
for i in range(0, len(indices_for_worker), self.batch_size):
batch_end = i + self.batch_size
batch_indices = indices_for_worker[i:batch_end]
yield build_window_batch(
self.sequences,
self.config.input_columns,
self.config.target_columns,
self.sample_index,
batch_indices,
depth_valid_masks=self.depth_valid_masks,
start_item_positions=self.start_item_positions,
split_start_item_positions=self.split_start_item_positions,
split_end_item_positions=self.split_end_item_positions,
)