p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/dataloader.py (python)
    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:
        # Default feature description for video data
        feature_description = {
            'raw_video': tf.io.FixedLenFeature([], tf.string),
            'sequence_length': tf.io.FixedLenFeature([], tf.int64),
        }
    
    # Create output directory if it doesn't exist
    os.makedirs(output_folder, exist_ok=True)
    
    # Find all TFRecord files
    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")
    
    # Create dataset from all 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())
        }
    
    # Process records and write to ArrayRecord files
    arrayrecord_files = []
    record_count = 0
    file_index = 0
    
    current_writer = None
    
    for record in dataset:
        parsed_record = parse_tfrecord(record)
        
        # Create new writer if needed
        if current_writer is None:
            output_file = os.path.join(output_folder, f"data_{file_index:04d}.arrayrecord")
            current_writer = grain.ArrayRecordWriter(output_file)
            arrayrecord_files.append(output_file)
        
        # Write record
        current_writer.write(parsed_record)
        record_count += 1
        
        # Close current file and start new one if we've reached the limit
        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")
    
    # Close final writer if it exists
    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