code: def load_ascii_dataset(
      corpus: str,
      *,
      batch_size: int,
      sequence_length: int,
      num_shuffle_batches: int = 10,
  ) -> Iterator[Batch]:
      """Loads a single-file ASCII dataset in memory."""
  
      if not corpus.isascii():
          raise ValueError("Loaded corpus is not ASCII.")
  
      if chr(PAD_TOKEN) in corpus:  # Reserve 0 codepoint for pad token.
          raise ValueError("Corpus must not contain the null byte.")
  
      # Naively tokenise by taking ASCII codepoints.
      corpus = np.array([ord(c) for c in corpus]).astype(np.int32)
      assert np.max(corpus) < VOCAB_SIZE
  
      crop_len = sequence_length + 1
      num_batches, remainder = divmod(
          corpus.size, batch_size * crop_len
      )
      if remainder:
          corpus = corpus[
              :-remainder
          ]  # Drop remainder (incomplete) batch.
      ds = corpus.reshape([-1, crop_len])
  
      if num_batches < num_shuffle_batches:
          raise ValueError(
              f"Only {num_batches} batches in the dataset; consider using a shorter "
              "sequence length or a smaller batch batch size.",
          )
  
      ds = repeat(ds)
      ds = shuffle(ds, buffer_size=batch_size * num_shuffle_batches)
      while True:
          batch = np.stack([next(ds) for _ in range(batch_size)])
          yield Batch(inputs=batch[:, :-1], targets=batch[:, 1:])
  
file_path: dm-haiku/examples/transformer/dataset.py
imports: from collections.abc import Iterable, Iterator
  import itertools
  import random
  from typing import NamedTuple, TypeVar
  import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: load_ascii_dataset(corpus) without return types