p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] experiments/tokenizer_time_measurement_grain.sh (shellscript)
#!/usr/bin/env bash
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate


data_dir='/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards_group_size_1'

# FIXME: the `time_measurement_interval` line is only for debugging and should be removed
srun python train_tokenizer.py \
    --time_measurement_interval 10 \
    --batch_size 48 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] generate_arrayrecord_dataset.py (python)
        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}")
[content] generate_arrayrecord_dataset.py (python)
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)
[content] generate_arrayrecord_dataset.py (python)
        videos_in_file = end_idx - start_idx
        
        file_path = os.path.join(output_dir, f"videos_{file_idx:04d}.array_record")
        writer = 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}")