p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] utils/dataloader.py (python)
    dataloader = grain.DataLoader(
        data_source=source,
        sampler=sampler,
        operations=operations,
        worker_count=num_workers,
        worker_buffer_size=1,
        read_options=read_options,
[tab] requirements.txt (pip-requirements)
dm_pix>=0.4.3
einops>=0.8.0
flax>=0.8.5
jax[cuda12]>=0.4.30
optax>=0.2.3
procgen>=0.10.7
tyro>=0.8.5
wandb>=0.17.4
<<<<<<< HEAD
grain>=0.2.10
pre-commit>=4.2.0
array-record>=0.7.2
=======
tensorflow>=2.1
pre-commit>=4.2.0
tqdm>=4.67.1
>>>>>>> main
[tab] tests/data/generate_dummy_data.py (python)
import tyro
import tensorflow as tf
import numpy as np
from pathlib import Path
from dataclasses import dataclass

@dataclass
class Args:
    data_dir: str = "data_tfrecords/dummy"
    num_episodes: int = 5
    episode_length: int = 16



def _bytes_feature(value):
    """Returns a bytes_list from a string / byte."""
    if isinstance(value, type(tf.constant(0))):
        value = value.numpy()  # BytesList won't unpack a string from an EagerTensor.
    return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value]))


def _int64_feature(value):
    """Returns an int64_list from a bool / enum / int / uint."""
    return tf.train.Feature(int64_list=tf.train.Int64List(value=[value]))


def create_tfrecord_example(episode_numpy_array):
    """Creates a TFRecord example from a numpy array video."""
    feature = {
        "height": _int64_feature(episode_numpy_array.shape[1]),
        "width": _int64_feature(episode_numpy_array.shape[2]),
        "channels": _int64_feature(episode_numpy_array.shape[3]),
        "sequence_length": _int64_feature(episode_numpy_array.shape[0]),
        "raw_video": _bytes_feature(episode_numpy_array.tobytes()),
    }
    return tf.train.Example(features=tf.train.Features(feature=feature))


def generate_dummy_tfrecord(
    output_path, num_episodes=5, episode_length=16, height=90, width=160, channels=3
):
    """Generates a dummy TFRecord file with synthetic video data."""
    print(f"Generating dummy TFRecord file at {output_path}")
    with tf.io.TFRecordWriter(str(output_path)) as writer:
        for i in range(num_episodes):
            np.random.seed(i)  # Seed per episode for some variation, but deterministic
            dummy_video = np.random.randint(
                0, 256, size=(episode_length, height, width, channels), dtype=np.uint8
            )
            tf_example = create_tfrecord_example(dummy_video)
            writer.write(tf_example.SerializeToString())
    print("Dummy TFRecord generation complete.")


if __name__ == "__main__":
    args = tyro.cli(Args)
    temp_dir = Path(args.data_dir)
    temp_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = temp_dir / "dummy_test_shard.tfrecord"
    generate_dummy_tfrecord(dummy_file, num_episodes=args.num_episodes, episode_length=args.episode_length)
    print(f"Generated dummy file: {dummy_file}")
[tab] tests/test_dataloader.py (python)
import unittest
import numpy as np
import tensorflow as tf
import tempfile
from pathlib import Path

from utils.dataloader import get_dataloader
from tests.data.generate_dummy_tfrecord import generate_dummy_tfrecord


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)
        self.dummy_tfrecord_path = self.test_data_dir / "dummy_test_shard.tfrecord"

        self.num_episodes = 5
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        generate_dummy_tfrecord(
            self.dummy_tfrecord_path,
            num_episodes=self.num_episodes,
            episode_length=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
        )
        self.tfrecord_files = [str(self.dummy_tfrecord_path)]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.tfrecord_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.tfrecord_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")  # type: ignore


if __name__ == "__main__":
    unittest.main()
[tab] reprocess_knoms_to_optimized_arrayrecord_layout.py (python)
from utils.dataloader import reprocess_arrayrecords_to_chunks

reprocess_arrayrecords_to_chunks(
    arrayrecord_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards",
    output_folder="/hkfs/work/workspace/scratch/tum_ind3695-jafa_ws_shared/data/knoms_arrayrecords_500_shards_optimized_layout",
)
[content] tests/test_dataloader.py (python)
import unittest
import numpy as np
import tensorflow as tf
import tempfile
from pathlib import Path

from utils.dataloader import get_dataloader
from tests.data.generate_dummy_tfrecord import generate_dummy_tfrecord


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)
        self.dummy_tfrecord_path = self.test_data_dir / "dummy_test_shard.tfrecord"

        self.num_episodes = 5
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        generate_dummy_tfrecord(
            self.dummy_tfrecord_path,
            num_episodes=self.num_episodes,
            episode_length=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
        )
        self.tfrecord_files = [str(self.dummy_tfrecord_path)]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.tfrecord_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.tfrecord_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")  # type: ignore


if __name__ == "__main__":
    unittest.main()
[content] tests/test_dataloader.py (python)
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")

    def test_dataloader_batch_shape_and_dtype(self):
        """Test that dataloader yields batches with correct shape and dtype."""
        seq_len = 8
        batch_size = 2

        dataloader = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        
        batch = next(dataloader)
        
        # Check shape: (batch_size, seq_len, height, width, channels)
        expected_shape = (batch_size, seq_len, self.image_height, self.image_width, self.image_channels)
        self.assertEqual(batch.shape, expected_shape)
        
        # Check dtype: should be float32 and normalized to [0, 1]
        self.assertEqual(batch.dtype, np.float32)
        self.assertTrue(np.all(batch >= 0.0) and np.all(batch <= 1.0))

    def test_dataloader_handles_different_sequence_lengths(self):
        """Test that dataloader works with different sequence lengths."""
        batch_size = 2
        
        for seq_len in [4, 8, 12]:
            dataloader = get_dataloader(
                self.array_record_files,
                seq_len,
                batch_size,
                self.image_height,
                self.image_width,
                self.image_channels,
                seed=self.fixed_seed,
            )
            
            batch = next(dataloader)
            expected_shape = (batch_size, seq_len, self.image_height, self.image_width, self.image_channels)
            self.assertEqual(batch.shape, expected_shape)
[content] tests/test_dataloader.py (python)
        
        # Get list of ArrayRecord files
        self.array_record_files = [
            str(self.test_data_dir / f) 
            for f in os.listdir(self.test_data_dir) 
            if f.endswith('.array_record')
        ]
[content] tests/test_dataloader.py (python)
        
        # Generate ArrayRecord dataset for testing
        dummy_file = self.test_data_dir / "dummy_test_shard.array_record"
        generate_dummy_arrayrecord(
            dummy_file,
            num_videos=self.num_videos,
[tab] tests/data/generate_dummy_arrayrecord.py (python)
import os
import pickle
import numpy as np
from pathlib import Path
from array_record.python.array_record_module import ArrayRecordWriter


def generate_dummy_arrayrecord(
    output_path: Path,
    num_videos: int = 5,
    episode_length: int = 16,
    height: int = 64,
    width: int = 64,
    channels: int = 3,
    seed: int = 42
):
    """Generates a dummy ArrayRecord file with synthetic video data for testing."""
    print(f"Generating dummy ArrayRecord file at {output_path}")
    
    # Ensure output directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    writer = ArrayRecordWriter(str(output_path), "group_size:1")
    try:
        for i in range(num_videos):
            np.random.seed(seed + i)  # Seed per video for some variation, but deterministic
            dummy_video = np.random.randint(
                0, 256, size=(episode_length, height, width, channels), dtype=np.uint8
            )
            
            # Create record in the format expected by ProcessEpisodeAndSlice
            record = {
                "raw_video": dummy_video.tobytes(),
                "sequence_length": episode_length
            }
            
            writer.write(pickle.dumps(record))
    finally:
        writer.close()
    
    print("Dummy ArrayRecord generation complete.")


if __name__ == "__main__":
    # Example usage
    test_dir = Path("tests/data/dummy_arrayrecord")
    test_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = test_dir / "dummy_test_shard.array_record"
    
    generate_dummy_arrayrecord(
        dummy_file,
        num_videos=5,
        episode_length=16,
        height=64,
        width=64,
        channels=3
    )
    print(f"Generated dummy file: {dummy_file}")
[content] tests/data/generate_dummy_data.py (python)
import os
import pickle
import numpy as np
from pathlib import Path
from array_record.python.array_record_module import ArrayRecordWriter


def generate_dummy_arrayrecord(
    output_path: Path,
    num_videos: int = 5,
    episode_length: int = 16,
    height: int = 64,
    width: int = 64,
    channels: int = 3,
    seed: int = 42
):
    """Generates a dummy ArrayRecord file with synthetic video data for testing."""
    print(f"Generating dummy ArrayRecord file at {output_path}")
    
    # Ensure output directory exists
    output_path.parent.mkdir(parents=True, exist_ok=True)
    
    writer = ArrayRecordWriter(str(output_path), "group_size:1")
    try:
        for i in range(num_videos):
            np.random.seed(seed + i)  # Seed per video for some variation, but deterministic
            dummy_video = np.random.randint(
                0, 256, size=(episode_length, height, width, channels), dtype=np.uint8
            )
            
            # Create record in the format expected by ProcessEpisodeAndSlice
            record = {
                "raw_video": dummy_video.tobytes(),
                "sequence_length": episode_length
            }
            
            writer.write(pickle.dumps(record))
    finally:
        writer.close()
    
    print("Dummy ArrayRecord generation complete.")


if __name__ == "__main__":
    # Example usage
    test_dir = Path("tests/data/dummy_arrayrecord")
    test_dir.mkdir(parents=True, exist_ok=True)
    dummy_file = test_dir / "dummy_test_shard.array_record"
    
    generate_dummy_arrayrecord(
        dummy_file,
        num_videos=5,
        episode_length=16,
        height=64,
        width=64,
        channels=3
    )
    print(f"Generated dummy file: {dummy_file}")
[tab] tests/test_dataloader.py (python)
import unittest
import numpy as np
import tempfile
import os
from pathlib import Path

from utils.dataloader import get_dataloader
from tests.data.generate_dummy_data import generate_dummy_arrayrecord


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        self.num_videos = 10
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        
        dummy_file = self.test_data_dir / "dummy_test_shard.array_record"
        generate_dummy_arrayrecord(
            dummy_file,
            num_videos=self.num_videos,
            episode_length=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
            seed=42
        )
        
        self.array_record_files = [
            str(self.test_data_dir / f) 
            for f in os.listdir(self.test_data_dir) 
            if f.endswith('.array_record')
        ]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")


if __name__ == "__main__":
    unittest.main()
[content] tests/test_dataloader.py (python)
import tensorflow as tf
import tempfile
from pathlib import Path

from utils.dataloader import get_dataloader
from tests.data.generate_dummy_tfrecord import generate_dummy_tfrecord


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)
        self.dummy_tfrecord_path = self.test_data_dir / "dummy_test_shard.tfrecord"

        self.num_episodes = 5
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        generate_dummy_tfrecord(
            self.dummy_tfrecord_path,
            num_episodes=self.num_episodes,
            episode_length=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
        )
        self.tfrecord_files = [str(self.dummy_tfrecord_path)]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.tfrecord_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.tfrecord_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")  # type: ignore
[content] tests/test_dataloader.py (python)
import tempfile
import os
from pathlib import Path

from utils.dataloader import get_dataloader
from tests.data.generate_dummy_data import generate_dummy_arrayrecord


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        self.num_videos = 10
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        
        dummy_file = self.test_data_dir / "dummy_test_shard.array_record"
        generate_dummy_arrayrecord(
            dummy_file,
            num_videos=self.num_videos,
            episode_length=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
            seed=42
        )
        
        self.array_record_files = [
            str(self.test_data_dir / f) 
            for f in os.listdir(self.test_data_dir) 
            if f.endswith('.array_record')
        ]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")
[tab] tests/test_dataloader.py (python)
import unittest
import numpy as np
import tempfile
import os
from pathlib import Path

from utils.dataloader import get_dataloader
from tests.data.generate_dummy_data import generate_dummy_arrayrecord


class DataloaderReproducibilityTest(unittest.TestCase):

    def setUp(self):
        super().setUp()
        self._temp_dir_manager = tempfile.TemporaryDirectory()
        self.test_data_dir = Path(self._temp_dir_manager.name)
        self.addCleanup(self._temp_dir_manager.cleanup)

        self.num_videos = 10
        self.episode_length = 16
        self.image_height = 64
        self.image_width = 64
        self.image_channels = 3
        
        dummy_file = self.test_data_dir / "dummy_test_shard.array_record"
        generate_dummy_arrayrecord(
            dummy_file,
            num_videos=self.num_videos,
            episode_length=self.episode_length,
            height=self.image_height,
            width=self.image_width,
            channels=self.image_channels,
            seed=42
        )
        
        self.array_record_files = [
            str(self.test_data_dir / f) 
            for f in os.listdir(self.test_data_dir) 
            if f.endswith('.array_record')
        ]

        self.fixed_seed = 42

    def test_dataloader_yields_reproducible_batches(self):
        seq_len = 8
        batch_size = 2

        dataloader1 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=self.fixed_seed,
        )
        batches1 = [next(dataloader1) for _ in range(3)]

        dataloader2 = get_dataloader(
            self.array_record_files,
            seq_len,
            batch_size,
            self.image_height,
            self.image_width,
            self.image_channels,
            num_workers=8,
            prefetch_buffer_size=1,
            seed=self.fixed_seed,
        )
        batches2 = [next(dataloader2) for _ in range(3)]

        for i, (b1, b2) in enumerate(zip(batches1, batches2)):
            np.testing.assert_array_equal(b1, b2, err_msg=f"Batch {i} is not reproducible")


if __name__ == "__main__":
    unittest.main()