p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] train_tokenizer.py (python)
            if step % args.time_measurement_interval == 0:
                jax.block_until_ready(train_state)
                elapsed_time = (time.time() - start_time)
                avg_step_time = elapsed_time / args.time_measurement_interval
                print(f"Step {step}, loss: {loss}, avg step time: {avg_step_time:.2f}s")
                if args.log and jax.process_index() == 0:
                    wandb.log({"avg_step_time_s": avg_step_time, "step": step})
                start_time = time.time()
            else:
                print(f"Step {step}, loss: {loss}")
[tab] experiments/dummy_time_measurement.sh (shellscript)
module unload mpi/openmpi/5.0
module unload devel/cuda/12.4
source .venv/bin/activate

srun python train_tokenizer.py \
    --data_dir /hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/dummy \
    --batch_size 192 \
    --time_measurement_interval 10
[tab] experiments/tokenizer_knoms_time_measurement.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_tfrecords_500_shards'

# 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 768 \
    --num_steps 300000 \
    --warmup_steps 10000 \
    --seed 0 \
    --min_lr=0.0000866 \
    --max_lr=0.0000866 \
    --data_dir $data_dir
[content] experiments/tokenizer_knoms_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_tfrecords_500_shards'

# 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
[tab] convert_knoms_to_arrayrecords.py (python)
from utils.dataloader import convert_tfrecords_to_arrayrecords

convert_tfrecords_to_arrayrecords(
    tfrecord_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_tfrecords_500_shards",
    output_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards",
)
[content] generate_arrayrecord_dataset.py (python)
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)
            video = record["videos"]
            
            print(f"  Record {i}: video shape = {video.shape}, dtype = {video.dtype}")
            print(f"    Value range: [{video.min():.3f}, {video.max():.3f}]")
            print(f"    Mean: {video.mean():.3f}")
            
    except Exception as e:
        print(f"Error reading ArrayRecord file: {e}")
        print("This might indicate a file format issue.")
[content] generate_arrayrecord_dataset.py (python)
    # 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)
[tab] test_arrayrecord_simple.py (python)
#!/usr/bin/env python3
"""
Simple test script following the ArrayRecord documentation exactly.
"""

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


def main():
    print("Testing ArrayRecord format following documentation...")
    
    # Create a simple test file
    example_file_path = "./test.array_record"
    
    # Write the dataset into a test array_record file
    writer = array_record_module.ArrayRecordWriter(example_file_path, "group_size:1")
    record_count = 0
    
    # Create simple test data (similar to what we need for training)
    for i in range(10):
        # Create a simple video-like record
        video = np.random.rand(16, 90, 160, 3).astype(np.float32)
        record = {"videos": video}
        
        writer.write(pickle.dumps(record))
        record_count += 1
    
    writer.close()
    
    print(f"Number of records written to array_record file {example_file_path}: {record_count}")
    
    # Test reading with Grain's ArrayRecordDataSource
    try:
        example_array_record_data_source = grain.sources.ArrayRecordDataSource(example_file_path)
        print(f"Number of records: {len(example_array_record_data_source)}")
        
        # Test reading first record
        first_record = example_array_record_data_source[0]
        parsed_record = pickle.loads(first_record)
        video = parsed_record["videos"]
        
        print(f"First record video shape: {video.shape}")
        print(f"First record video dtype: {video.dtype}")
        print(f"Value range: [{video.min():.3f}, {video.max():.3f}]")
        
        print("✅ ArrayRecord test successful!")
        
    except Exception as e:
        print(f"❌ Error reading ArrayRecord file: {e}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    main()