| from __future__ import annotations |
|
|
| from pathlib import Path |
|
|
| import tensorflow as tf |
|
|
|
|
| class GoDataset: |
| """Factory for train/val/test tf.data.Dataset objects.""" |
|
|
| def __init__( |
| self, |
| processed_dir: str | Path, |
| seq_len: int = 2048, |
| batch_size: int = 32, |
| shuffle_buffer: int = 50_000, |
| seed: int = 42, |
| ): |
| self.processed_dir = Path(processed_dir) |
| self.seq_len = seq_len |
| self.batch_size = batch_size |
| self.shuffle_buffer = shuffle_buffer |
| self.seed = seed |
|
|
| def build(self, split: str) -> tf.data.Dataset: |
| split_dir = self.processed_dir / split |
| tfrecords = sorted(split_dir.glob("*.tfrecord")) |
| if not tfrecords: |
| raise FileNotFoundError(f"No TFRecords found in {split_dir}") |
|
|
| fileset = tf.data.Dataset.from_tensor_slices([str(p) for p in tfrecords]) |
| if split == "train": |
| fileset = fileset.shuffle(len(tfrecords), seed=self.seed, reshuffle_each_iteration=True) |
|
|
| raw = fileset.interleave( |
| lambda p: tf.data.TFRecordDataset(p, compression_type=""), |
| cycle_length=16, |
| num_parallel_calls=tf.data.AUTOTUNE, |
| deterministic=(split != "train"), |
| ) |
|
|
| feature_spec = { |
| "input_ids": tf.io.FixedLenSequenceFeature([], tf.int64, allow_missing=True) |
| } |
|
|
| def _parse(serialised: tf.Tensor): |
| parsed = tf.io.parse_single_example(serialised, feature_spec) |
| ids = tf.cast(parsed["input_ids"], tf.int32) |
| return ids[:-1], ids[1:] |
|
|
| ds = raw.map(_parse, num_parallel_calls=tf.data.AUTOTUNE) |
|
|
| if split == "train": |
| ds = ds.shuffle(self.shuffle_buffer, seed=self.seed) |
|
|
| ds = ( |
| ds.batch(self.batch_size, drop_remainder=True) |
| .prefetch(tf.data.AUTOTUNE) |
| ) |
| return ds |
|
|
| def train(self) -> tf.data.Dataset: |
| return self.build("train") |
|
|
| def val(self) -> tf.data.Dataset: |
| return self.build("val") |
|
|
| def test(self) -> tf.data.Dataset: |
| return self.build("test") |
|
|
| def steps_per_epoch(self, split: str = "train") -> int: |
| """Estimate steps; requires at least one pass through the dataset.""" |
| count = sum( |
| 1 for _ in self.build(split).unbatch().batch(self.batch_size) |
| ) |
| return count |
|
|
