| import json |
| import time |
| import random |
| from typing import Literal |
|
|
| import requests |
| import zstandard as zstd |
| from torch.utils.data import IterableDataset, get_worker_info |
|
|
|
|
| Subset = Literal["train", "val", "test"] |
| URLs = { |
| "val": [ |
| "https://the-eye.eu/public/AI/pile/val.jsonl.zst", |
| ], |
| "test": [ |
| "https://the-eye.eu/public/AI/pile/test.jsonl.zst", |
| ], |
| "train": [ |
| "https://the-eye.eu/public/AI/pile/train/00.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/01.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/02.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/03.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/04.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/05.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/06.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/07.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/08.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/09.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/10.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/11.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/12.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/13.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/14.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/15.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/16.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/17.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/18.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/19.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/20.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/21.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/22.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/23.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/24.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/25.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/26.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/27.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/28.jsonl.zst", |
| "https://the-eye.eu/public/AI/pile/train/29.jsonl.zst", |
| ], |
| } |
|
|
|
|
| def _read_line_from_stream(reader, initial_line="", buffer_size=4096): |
| line = initial_line |
| while True: |
| c = reader.read(buffer_size) |
| if not c: |
| raise StopIteration |
| line += c.decode("utf-8") |
| if "\n" in line: |
| break |
| return line.split("\n", 1) |
|
|
|
|
| def _line_streamer(reader, buffer_size=4096): |
| rest = "" |
| while True: |
| try: |
| line, rest = _read_line_from_stream( |
| reader, |
| rest, |
| buffer_size, |
| ) |
| yield line |
| except StopIteration: |
| break |
|
|
|
|
| class ThePile(IterableDataset): |
| TEXT_BUFFER_SIZE = 4096 |
|
|
| def __init__(self, subset: Subset): |
| self.subset = subset |
|
|
| def __iter__(self): |
| urls = URLs[self.subset].copy() |
| while True: |
| wi = get_worker_info() |
| seed = wi.id if wi is not None else None |
| rnd = random.Random(seed) |
| rnd.shuffle(urls) |
| for url in urls: |
| r = requests.get(url, stream=True) |
| with zstd.ZstdDecompressor().stream_reader(r.raw) as reader: |
| for line in _line_streamer(reader, self.TEXT_BUFFER_SIZE): |
| data = json.loads(line) |
| yield data |
|
|
|
|
| if __name__ == "__main__": |
| from tqdm import tqdm |
|
|
| dataset = ThePile("train") |
| for data in tqdm(dataset, smoothing=0.01): |
| pass |
| |
