he-stackデータセットから10MBだけ読み込んで、先頭を表示する

the_stack_load.py

# the-stackデータセットから10MBだけ読み込んで、先頭を表示する

import sys
from datasets import load_dataset

dataset = load_dataset("bigcode/the-stack", split="train", streaming=True)

data_subset = []
total_size = 0

for sample in dataset:
    sample_size = sys.getsizeof(sample)
    if total_size + sample_size > 10 * 1024 * 1024:  # 10MB
        break
    data_subset.append(sample)
    total_size += sample_size

print(len(data_subset))
print(total_size)
print(data_subset[0])