[content] 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.array_record_module import ArrayRecordWriter
import argparse
import multiprocessing as mp
from functools import partial
import time


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)
    
    return video


def create_single_arrayrecord_file(
    file_info: tuple,
    output_dir: str,
    seq_len: int,
    height: int,
    width: int,
    channels: int,
    records_per_file: int,
    seed: int
) -> tuple:
    """
    Create a single ArrayRecord file with synthetic video data.
    
    Args:
        file_info: Tuple of (file_idx, start_idx, end_idx)
        output_dir: Directory to save ArrayRecord files
        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
        
    Returns:
        Tuple of (file_path, num_videos_created, success)
    """
    file_idx, start_idx, end_idx = file_info
    videos_in_file = end_idx - start_idx
    
    # Set seed for this process (add file_idx to make each file different)
    np.random.seed(seed + file_idx)
    
    file_path = os.path.join(output_dir, f"videos_{file_idx:04d}.array_record")
    
    try:
        writer = ArrayRecordWriter(file_path, "group_size:1")
        
        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))
        
        writer.close()
        return (file_path, videos_in_file, True)
        
    except Exception as e:
        print(f"Error creating file {file_path}: {e}")
        return (file_path, 0, False)


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,
    num_processes: int | None = None
):
    """
    Create ArrayRecord dataset with synthetic video data using multiprocessing.
    
    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
        num_processes: Number of processes to use (None for auto-detect)
    """
    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}")
    
    # Prepare file information for each worker
    file_infos = []
    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)
        file_infos.append((file_idx, start_idx, end_idx))
    
    # Set number of processes (use CPU count if not specified)
    if num_processes is None:
        num_processes = min(mp.cpu_count(), num_files)
    
    print(f"Using {num_processes} processes for parallel generation...")
    
    start_time = time.time()
    
    # Create partial function with fixed arguments
    worker_func = partial(
        create_single_arrayrecord_file,
        output_dir=output_dir,
        seq_len=seq_len,
        height=height,
        width=width,
        channels=channels,
        records_per_file=records_per_file,
        seed=seed
    )
    
    # Use multiprocessing to create files in parallel
    with mp.Pool(processes=num_processes) as pool:
        results = pool.map(worker_func, file_infos)
    
    end_time = time.time()
    
    # Process results
    total_records = 0
    successful_files = 0
    
    for file_path, num_videos_created, success in results:
        if success:
            print(f"✓ Created {file_path} with {num_videos_created} videos")
            total_records += num_videos_created
            successful_files += 1
        else:
            print(f"✗ Failed to create {file_path}")
    
    print(f"
Dataset generation complete!")
    print(f"Total videos generated: {total_records}")
    print(f"Successful files: {successful_files}/{num_files}")
    print(f"Files created in: {output_dir}")
    print(f"Generation time: {end_time - start_time:.2f} seconds")


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/dummy",
                       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("--num_processes", type=int, default=None,
                       help="Number of processes to use (default: auto-detect)")
    parser.add_argument("--verify", action="store_true",
                       help="Verify the generated dataset")
    
    args = parser.parse_args()
    
    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,
        num_processes=args.num_processes
    )
    
    # Verify if requested
    if args.verify:
        verify_dataset(args.output_dir)
    
if __name__ == "__main__":
    main()