p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    # Ensure we are operating on a single ArrayRecord file and not a folder
    if not os.path.isfile(args.input):
        raise AssertionError(
            "--input must be a single ArrayRecord file path, not a directory"
        )
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    idx = rng.randrange(num_records)

    raw: Optional[bytes] = None
    # Prefer Python indexing if implemented
    try:
        raw = reader[idx]
    except Exception:
        # Fall back to common explicit methods exposed by some builds
        for method_name in ("read", "read_at", "get", "record_at"):
            method = getattr(reader, method_name, None)
            if method is None:
                continue
            try:
                candidate = method(int(idx))
                if candidate is not None:
                    raw = candidate
                    break
            except Exception:
                continue

    if raw is None:
        raise RuntimeError("Failed to read a record by index without iteration.")

    return pickle.loads(raw)
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    try:
        num_records = len(reader)
    except Exception as e:
        raise AssertionError(
            "ArrayRecordReader must support len() for random access; please update array-record."
        ) from e
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    # Reservoir sample one record uniformly at random without loading all into memory
    chosen: Optional[bytes] = None
    count = 0
    for raw in reader:
        count += 1
        if rng.random() < 1.0 / count:
            chosen = raw
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    rng = random.Random(seed)
    # Reservoir sample one record uniformly at random without loading all into memory
    chosen: Optional[bytes] = None
    count = 0
    for raw in reader:
        count += 1
        if rng.random() < 1.0 / count:
            chosen = raw
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    record = load_one_sequence_via_grain(
        data_dir=args.data_dir,
        seq_len=int(args.seq_len),
        image_h=int(args.height),
        image_w=int(args.width),
        image_c=int(args.channels),
        seed=int(args.seed),
    )
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    # Obtain frames either directly (from Grain) or by decoding raw bytes
    if "videos" in record:
        frames = np.asarray(record["videos"], dtype=np.uint8)
        seq_len = int(frames.shape[0])
        h, w = int(frames.shape[1]), int(frames.shape[2])
        assert frames.shape[3] == channels, f"Expected {channels} channels, got {frames.shape[3]}"
    else:
        seq_len = int(record["sequence_length"])  # type: ignore[index]
        raw_video: bytes = record["raw_video"]  # type: ignore[index]
        arr = np.frombuffer(raw_video, dtype=np.uint8)
        total_elements = arr.size
        if height is not None and width is not None:
            h, w = int(height), int(width)
            expected = seq_len * h * w * channels
            assert expected == total_elements, f"Expected {expected} elements, got {total_elements}"
        else:
            h, w = infer_hw_from_bytes(total_elements, seq_len, channels)
        frames = arr.reshape(seq_len, h, w, channels)
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    state = dataloader._create_initial_state()
    it = grain.DataLoaderIterator(dataloader, state)
    batch = next(it)
    videos = np.asarray(batch["videos"][0])  # (T,H,W,C) uint8
    actions = None
    if "actions" in batch:
        actions = np.asarray(batch["actions"][0]).reshape(-1)
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    dataloader = get_dataloader(
        array_record_files,
        seq_len,
        1,  # global batch size 1 for visualization
        image_h,
        image_w,
        image_c,
        num_workers=1,
        prefetch_buffer_size=1,
        seed=0 if seed is None else int(seed),
    )
[content] data/jasmine_data/atari/visualize_array_record.py (python)
def load_one_sequence_via_grain(
    data_dir: str,
    seq_len: int,
    image_h: int,
    image_w: int,
    image_c: int,
    seed: Optional[int] = None,
) -> dict:
    array_record_files = [
        os.path.join(data_dir, x)
        for x in os.listdir(data_dir)
        if x.endswith(".array_record")
    ]
    if not array_record_files:
        raise RuntimeError(f"No .array_record files found in {data_dir}")
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    # Decode frames
    arr = np.frombuffer(raw_video, dtype=np.uint8)
    total_elements = arr.size

    if height is not None and width is not None:
        h, w = int(height), int(width)
        expected = seq_len * h * w * channels
        assert expected == total_elements, f"Expected {expected} elements, got {total_elements}"
    else:
        h, w = infer_hw_from_bytes(total_elements, seq_len, channels)

    frames = arr.reshape(seq_len, h, w, channels)
[content] data/jasmine_data/atari/visualize_array_record.py (python)
def load_random_record(input_path: str, seed: Optional[int] = None) -> dict:
    assert ArrayRecordReader is not None, "array_record is not installed. pip install array-record"
    reader = ArrayRecordReader(str(input_path))
[content] data/jasmine_data/atari/visualize_array_record.py (python)
    array_record_files = [
        os.path.join(data_dir, x)
        for x in os.listdir(data_dir)
        if x.endswith(".array_record")
    ]
    if not array_record_files:
        raise RuntimeError(f"No .array_record files found in {data_dir}")

    dataloader = get_dataloader(
        array_record_files,
        seq_len,
        1,  # global batch size 1 for visualization
        image_h,
        image_w,
        image_c,
        num_workers=1,
        prefetch_buffer_size=1,
        seed=0 if seed is None else int(seed),
    )
[content] data/jasmine_data/atari/visualize_array_record.py (python)
def load_one_sequence_via_grain(
    data_dir: str,
    seq_len: int,
    image_h: int,
    image_w: int,
    image_c: int,
    seed: Optional[int] = None,
) -> dict:
    import grain  # lazy import to avoid linter issues when grain is unavailable
[content] data/jasmine_data/atari/visualize_array_record.py (python)
def load_one_sequence_via_grain(
    data_dir: str,
    seq_len: int,
    image_h: int,
    image_w: int,
    image_c: int,
    seed: Optional[int] = None,
) -> dict:
    grain = importlib.import_module("grain")  # dynamic import to avoid linter issues
[tab] data/pyproject.toml (plaintext)
[project]
name = "jasmine-data"
version = "0.1.0"
authors = [
  {name = "Alfred Nguyen", email = "alfrednguyen02@gmail.com"},
  {name = "Mihir Mahajan", email = "mihir@pdoom.org"},
  {name = "Franz Srambical", email = "franz@pdoom.org"},
]
requires-python = "==3.10.*"
dependencies = [
  "procgen>=0.10.7",
  "gym3>=0.3.3",
  "array-record>=0.7.2",
  "numpy",
  "hf-transfer==0.1.9",
  "huggingface-hub[cli]>=0.34.3",
  "ffmpeg-python==0.2.0",
  "pillow>=11.3.0",
  "tqdm>=4.67.1",
  "tyro>=0.8.5",
  "gymnasium[atari, accept-rom-license]==0.29.1",
  "torch>=2.8.0",
  "tensorboard>=2.20.0",
  "opencv-python>=4.6.0.66,<5",
  "einops>=0.8.0",
]

[build-system]
requires = ["uv_build>=0.8.22,<0.9.0"]
build-backend = "uv_build"

[tool.uv.build-backend]
module-root = ""
[content] data/jasmine_data/atari/visualize_array_record.py (python)


def get_dataloader(
    array_record_paths: list[str],
    seq_len: int,
    global_batch_size: int,
    image_h: int,
    image_w: int,
    image_c: int,
    num_workers: int = 1,
    prefetch_buffer_size: int = 1,
    seed: int = 42,
):
    """
    Creates a data loading pipeline using Grain.
    """
    if not array_record_paths:
        raise ValueError("array_record_paths list cannot be empty.")

    num_processes = jax.process_count()

    if global_batch_size % num_processes != 0:
        raise ValueError(
            f"Global batch size {global_batch_size} must be divisible by "
            f"the number of JAX processes {num_processes} for proper sharding."
        )
    per_process_batch_size = global_batch_size // num_processes

    source = grain.sources.ArrayRecordDataSource(array_record_paths)

    sampler = grain.samplers.IndexSampler(
        num_records=len(source),
        shard_options=grain.sharding.ShardByJaxProcess(drop_remainder=True),
        shuffle=True,
        num_epochs=None,
        seed=seed,
    )

    operations = [
        EpisodeLengthFilter(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        ProcessEpisodeAndSlice(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        grain.transforms.Batch(batch_size=per_process_batch_size, drop_remainder=True),
    ]

    read_options = grain.ReadOptions(
        prefetch_buffer_size=prefetch_buffer_size,
        num_threads=1,
    )
    dataloader = grain.DataLoader(
        data_source=source,
        sampler=sampler,
        operations=operations,
        worker_count=num_workers,
        worker_buffer_size=1,
        read_options=read_options,
    )

    return dataloader
[content] data/jasmine_data/atari/visualize_array_record.py (python)
EpisodeLengthFilter(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),
        ProcessEpisodeAndSlice(
            seq_len=seq_len, image_h=image_h, image_w=image_w, image_c=image_c
        ),