[tab] utils/dataloader.py (python)
import jax
import numpy as np
import grain
from typing import Any, Optional
from array_record.python import array_record_module
import tensorflow as tf
import os
from pathlib import Path


def convert_tfrecords_to_arrayrecords(
    tfrecord_folder: str,
    output_folder: str,
    records_per_file: int = 1000,
    feature_description: Optional[dict] = None,
):
    """
    Converts TFRecord files to ArrayRecord format for use with Grain.
    
    Args:
        tfrecord_folder: Path to folder containing TFRecord files
        output_folder: Path to output folder for ArrayRecord files
        records_per_file: Number of records to write per ArrayRecord file
        feature_description: Dictionary describing TFRecord features. If None,
                           uses default description for video data.
    
    Returns:
        List of paths to created ArrayRecord files
    """
    if feature_description is None:
        feature_description = {
            'raw_video': tf.io.FixedLenFeature([], tf.string),
            'sequence_length': tf.io.FixedLenFeature([], tf.int64),
        }
    
    os.makedirs(output_folder, exist_ok=True)
    
    tfrecord_files = list(Path(tfrecord_folder).glob("*.tfrecord"))
    if not tfrecord_files:
        raise ValueError(f"No TFRecord files found in {tfrecord_folder}")
    
    print(f"Found {len(tfrecord_files)} TFRecord files")
    
    dataset = tf.data.TFRecordDataset([str(f) for f in tfrecord_files])
    
    def parse_tfrecord(example_proto):
        """Parse a single TFRecord example."""
        parsed_features = tf.io.parse_single_example(example_proto, feature_description)
        return {
            'raw_video': parsed_features['raw_video'].numpy(),
            'sequence_length': int(parsed_features['sequence_length'].numpy())
        }
    
    arrayrecord_files = []
    record_count = 0
    file_index = 0
    
    current_writer = None
    
    for record in dataset:
        parsed_record = parse_tfrecord(record)
        
        if current_writer is None:
            output_file = os.path.join(output_folder, f"data_{file_index:04d}.arrayrecord")
            current_writer = grain.io.ArrayRecordWriter(output_file)
            arrayrecord_files.append(output_file)
        
        current_writer.write(parsed_record)
        record_count += 1
        
        if record_count % records_per_file == 0:
            current_writer.close()
            current_writer = None
            file_index += 1
            print(f"Created ArrayRecord file {file_index}: {record_count} records processed")
    
    if current_writer is not None:
        current_writer.close()
        print(f"Created final ArrayRecord file: {record_count} total records processed")
    
    print(f"Conversion complete! Created {len(arrayrecord_files)} ArrayRecord files")
    return arrayrecord_files


class ProcessEpisodeAndSlice(grain.transforms.RandomMap):
    """
    A Grain Transformation that combines parsing, slicing, and normalizing.

    This class replaces the functionality of both `_parse_tfrecord_fn` and
    `_tf_process_episode` from the original tf.data pipeline. It operates on
    a single data record (episode).
    """

    def __init__(self, seq_len: int, image_h: int, image_w: int, image_c: int):
        """Initializes the transformation with processing parameters."""
        self.seq_len = seq_len
        self.image_h = image_h
        self.image_w = image_w
        self.image_c = image_c

    def random_map(self, element: dict, rng: np.random.Generator) -> Any:
        """
        Processes a single raw episode from the data source.

        Args:
            element: A dictionary representing one record from the DataSource.
                     Expected to contain 'raw_video' (bytes), and
                     'sequence_length' (int).
            rng: A per-record random number generator provided by the Grain sampler.

        Returns:
            A processed video sequence as a NumPy array with shape
            (seq_len, image_h, image_w, image_c) and dtype float32.
        """
        # Step 1: Parse and reshape the raw video data (replaces _parse_tfrecord_fn)
        # We assume the element from the data source is a dict.
        video_shape = (
            element["sequence_length"],
            self.image_h,
            self.image_w,
            self.image_c,
        )
        episode_tensor = np.frombuffer(element["raw_video"], dtype=np.uint8)
        episode_tensor = episode_tensor.reshape(video_shape)

        # Step 2: Randomly slice and normalize (replaces _tf_process_episode)
        current_episode_len = episode_tensor.shape[0]
        if current_episode_len < self.seq_len:
             raise ValueError(f"An episode has length {current_episode_len}, which is "
                              f"shorter than the requested sequence length {self.seq_len}.")
        
        max_start_idx = current_episode_len - self.seq_len
        
        # Use the provided Grain RNG for deterministic randomness.
        start_idx = rng.integers(0, max_start_idx + 1)

        seq = episode_tensor[start_idx : start_idx + self.seq_len]

        # Normalize to [0, 1]
        processed_sequence = seq.astype(np.float32) / 255.0

        return processed_sequence


def get_dataloader_grain(
    array_record_paths: list[str],  # List of ArrayRecord file paths
    seq_len: int,
    global_batch_size: int,
    image_h: int,
    image_w: int,
    image_c: int,
    # shuffle_buffer_size is not needed; Grain shuffles all indices globally.
    num_workers: int = 4,  # Replaces num_parallel_calls
    seed: int = 42,
):
    """
    Creates a data loading pipeline using Google Grain.

    This function sets up a data source, a sampler for sharding/shuffling,
    and a sequence of operations (transformations) to be applied to the data.
    """
    if not array_record_paths:
        raise ValueError("array_record_paths list cannot be empty.")

    process_id = jax.process_index()
    num_processes = jax.process_count()

    if global_batch_size % num_processes != 0:
        raise ValueError(
            f"Global batch size {global_batch_size} must be divisible by "
            f"the number of JAX processes {num_processes} for proper sharding."
        )
    per_process_batch_size = global_batch_size // num_processes

    # 1. DataSource: Reads records from ArrayRecord files.
    # This replaces `tf.data.TFRecordDataset`.
    source = grain.sources.ArrayRecordDataSource(array_record_paths)
    
    # 2. Sampler: Defines the order of access, sharding, shuffling, and epochs.
    # This replaces `dataset.shard`, `dataset.shuffle`, and `dataset.repeat`.
    shard_options = grain.sharding.ShardOptions(
        shard_index=process_id, shard_count=num_processes, drop_remainder=True
    )

    sampler = grain.samplers.IndexSampler(
        num_records=len(source),
        shard_options=shard_options,
        seed=seed,
        num_epochs=None,  # `None` means repeat indefinitely, like `tf.data.repeat(None)`.
        shuffle=True,  # Enables global shuffling of all record indices.
    )

    # 3. Operations: A list of transformations to apply to each record.
    # This replaces the sequence of `.map()` calls.
    operations = [
        ProcessEpisodeAndSlice(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        grain.transforms.Batch(batch_size=per_process_batch_size, drop_remainder=True),
    ]

    # 4. DataLoader: The main entry point that orchestrates the pipeline.
    # It manages worker processes for parallel data loading and transformation,
    # replacing `num_parallel_calls` and `.prefetch()`.
    dataloader = grain.DataLoader(
        data_source=source,
        sampler=sampler,
        operations=operations,
        worker_count=num_workers,
    )

    # 5. Return an iterator, which is the standard way to consume a Grain DataLoader.
    return iter(dataloader)