[tab] generate_arrayrecord_dataset.py (python)
#!/usr/bin/env python3
"""
Generate ArrayRecord dataset compatible with train_tokenizer.py

This script creates synthetic video data and saves it in ArrayRecord format
that can be used by the tokenizer training script.
"""

import os
import pickle
import numpy as np
import grain
from array_record.python import array_record_module
import argparse


def generate_synthetic_video(
    seq_len: int = 16,
    height: int = 90,
    width: int = 160,
    channels: int = 3
) -> np.ndarray:
    """
    Generate synthetic video data with random frames for training.
    
    Args:
        seq_len: Number of frames in the video sequence
        height: Height of each frame
        width: Width of each frame
        channels: Number of color channels
        
    Returns:
        Video array of shape (seq_len, height, width, channels)
    """
    video = np.random.rand(seq_len, height, width, channels).astype(np.float32)
    
    video = np.clip(video, 0.0, 1.0)
    
    return video


def create_arrayrecord_dataset(
    output_dir: str,
    num_videos: int = 1000,
    seq_len: int = 16,
    height: int = 90,
    width: int = 160,
    channels: int = 3,
    records_per_file: int = 100,
    seed: int = 42
):
    """
    Create ArrayRecord dataset with synthetic video data.
    
    Args:
        output_dir: Directory to save ArrayRecord files
        num_videos: Total number of videos to generate
        seq_len: Number of frames per video
        height: Frame height
        width: Frame width
        channels: Number of color channels
        records_per_file: Number of records per ArrayRecord file
        seed: Random seed for reproducibility
    """
    np.random.seed(seed)
    
    os.makedirs(output_dir, exist_ok=True)
    
    num_files = (num_videos + records_per_file - 1) // records_per_file
    
    print(f"Generating {num_videos} videos across {num_files} ArrayRecord files...")
    print(f"Each file will contain up to {records_per_file} videos")
    print(f"Video dimensions: {seq_len} frames × {height}×{width}×{channels}")
    
    total_records = 0
    
    for file_idx in range(num_files):
        start_idx = file_idx * records_per_file
        end_idx = min((file_idx + 1) * records_per_file, num_videos)
        videos_in_file = end_idx - start_idx
        
        file_path = os.path.join(output_dir, f"videos_{file_idx:04d}.array_record")
        writer = array_record_module.ArrayRecordWriter(file_path, "group_size:1")
        
        print(f"Creating file {file_path} with {videos_in_file} videos...")
        
        for video_idx in range(videos_in_file):
            video = generate_synthetic_video(seq_len, height, width, channels)
            
            # Convert to uint8 format as expected by the dataloader
            video_uint8 = (video * 255).astype(np.uint8)
            
            # Create record in the format expected by ProcessEpisodeAndSlice
            record = {
                "raw_video": video_uint8.tobytes(),
                "sequence_length": seq_len
            }
            
            writer.write(pickle.dumps(record))
            total_records += 1
            
            if (video_idx + 1) % 100 == 0:
                print(f"  Generated {video_idx + 1}/{videos_in_file} videos")
        
        writer.close()
        print(f"Completed file {file_path}")
    
    print(f"
Dataset generation complete!")
    print(f"Total videos generated: {total_records}")
    print(f"Files created in: {output_dir}")
    print(f"Files: {len([f for f in os.listdir(output_dir) if f.endswith('.array_record')])}")


def verify_dataset(output_dir: str, num_samples: int = 5):
    """
    Verify the generated dataset using Grain's ArrayRecordDataSource.
    
    Args:
        output_dir: Directory containing ArrayRecord files
        num_samples: Number of samples to verify
    """
    print(f"
Verifying dataset in {output_dir}...")
    
    # Find all ArrayRecord files
    array_record_files = [
        os.path.join(output_dir, f) 
        for f in os.listdir(output_dir) 
        if f.endswith('.array_record')
    ]
    
    if not array_record_files:
        print("No ArrayRecord files found!")
        return
    
    print(f"Found {len(array_record_files)} ArrayRecord files")
    
    # Use Grain's ArrayRecordDataSource as shown in the documentation
    try:
        data_source = grain.sources.ArrayRecordDataSource(array_record_files[0])
        print(f"Number of records in first file: {len(data_source)}")
        
        # Load and verify a few samples
        for i in range(min(num_samples, len(data_source))):
            record_bytes = data_source[i]
            record = pickle.loads(record_bytes)
            
            # Reconstruct video from raw_video bytes
            video_shape = (record["sequence_length"], 90, 160, 3)  # Hardcoded for now
            video = np.frombuffer(record["raw_video"], dtype=np.uint8).reshape(video_shape)
            
            print(f"  Record {i}: video shape = {video.shape}, dtype = {video.dtype}")
            print(f"    Value range: [{video.min()}, {video.max()}]")
            print(f"    Mean: {video.mean():.1f}")
            print(f"    Sequence length: {record['sequence_length']}")
            
    except Exception as e:
        print(f"Error reading ArrayRecord file: {e}")
        print("This might indicate a file format issue.")


def main():
    parser = argparse.ArgumentParser(description="Generate ArrayRecord dataset for tokenizer training")
    parser.add_argument("--output_dir", type=str, default="data_arrayrecord/coinrun",
                       help="Output directory for ArrayRecord files")
    parser.add_argument("--num_videos", type=int, default=1000,
                       help="Total number of videos to generate")
    parser.add_argument("--seq_len", type=int, default=16,
                       help="Number of frames per video")
    parser.add_argument("--height", type=int, default=90,
                       help="Frame height")
    parser.add_argument("--width", type=int, default=160,
                       help="Frame width")
    parser.add_argument("--channels", type=int, default=3,
                       help="Number of color channels")
    parser.add_argument("--records_per_file", type=int, default=100,
                       help="Number of records per ArrayRecord file")
    parser.add_argument("--seed", type=int, default=42,
                       help="Random seed for reproducibility")
    parser.add_argument("--verify", action="store_true",
                       help="Verify the generated dataset")
    
    args = parser.parse_args()
    
    # Generate the dataset
    create_arrayrecord_dataset(
        output_dir=args.output_dir,
        num_videos=args.num_videos,
        seq_len=args.seq_len,
        height=args.height,
        width=args.width,
        channels=args.channels,
        records_per_file=args.records_per_file,
        seed=args.seed
    )
    
    # Verify if requested
    if args.verify:
        verify_dataset(args.output_dir)
    
if __name__ == "__main__":
    main()