| """FineWeb-Edu text data preparation and streaming. |
| |
| Dataset: https://huggingface.co/datasets/HuggingFaceFW/fineweb-edu |
| Subset: sample-10BT (~10B tokens) or sample-100BT (~100B tokens) |
| Storage: ~6 GB (10BT) or ~60 GB (100BT) |
| License: ODC-BY |
| """ |
| import math, torch, numpy as np |
| from dataclasses import dataclass |
| from typing import Optional |
| import sys, os |
| sys.path.insert(0, os.path.join(os.path.dirname(__file__), "..", "..")) |
| from arbitor.config import SPECIAL_VOCAB |
|
|
|
|
| @dataclass |
| class FineWebConfig: |
| subset: str = "sample-10BT" |
| ctx: int = 2048 |
| batch_size: int = 8 |
| shuffle_buffer: int = 10000 |
| split: str = "train" |
|
|
|
|
| class FineWebStream: |
| """Streaming text batches from FineWeb-Edu via HF datasets iterable.""" |
|
|
| def __init__(self, cfg: FineWebConfig): |
| self.cfg = cfg |
| self._ds = None |
| self._bos = SPECIAL_VOCAB['BOS'] |
| self._eos = SPECIAL_VOCAB['EOS'] |
| self._pad = SPECIAL_VOCAB['PAD'] |
|
|
| def _lazy_init(self): |
| if self._ds is not None: |
| return |
| from datasets import load_dataset |
| ds = load_dataset("HuggingFaceFW/fineweb-edu", self.cfg.subset, |
| split=self.cfg.split, streaming=True) |
| self._ds = ds.shuffle(buffer_size=self.cfg.shuffle_buffer, seed=42) |
|
|
| def _tokenize(self, text: str) -> torch.LongTensor: |
| raw = text.encode("utf-8") |
| tokens = [self._bos] + list(raw) + [self._eos] |
| return torch.tensor(tokens[:self.cfg.ctx + 1], dtype=torch.long) |
|
|
| def batches(self): |
| self._lazy_init() |
| buf = [] |
| for example in self._ds: |
| tokens = self._tokenize(example["text"]) |
| if tokens.numel() < 5: |
| continue |
| buf.append(tokens) |
| if len(buf) >= self.cfg.batch_size: |
| batch = buf[:self.cfg.batch_size] |
| buf = buf[self.cfg.batch_size:] |
| max_len = max(t.numel() for t in batch) |
| padded = torch.stack([ |
| torch.cat([t, t.new_full((max_len - t.numel(),), self._pad)]) |
| for t in batch |
| ]) |
| yield padded[:, :-1].contiguous(), padded[:, 1:].contiguous() |
|
|
| def num_samples(self) -> int: |
| if self.cfg.subset == "sample-10BT": |
| return 10_000_000_000 // self.cfg.ctx |
| return 100_000_000_000 // self.cfg.ctx |
|
|
