code: def load(
      split: tfds.Split,
      *,
      batch_size: int,
      sequence_length: int,
  ) -> Iterator[Batch]:
      """Creates the Tiny Shakespeare dataset as a character modelling task."""
  
      def preprocess_fn(
          x: Mapping[str, tf.Tensor]
      ) -> Mapping[str, tf.Tensor]:
          x = x["text"]
          x = tf.strings.unicode_split(x, "UTF-8")
          x = tf.squeeze(tf.io.decode_raw(x, tf.uint8), axis=-1)
          x = tf.cast(x, tf.int32)
          return {"input": x[:-1], "target": x[1:]}
  
      ds = tfds.load(name="tiny_shakespeare", split=split)
      ds = ds.map(preprocess_fn)
      ds = ds.unbatch()
      ds = ds.batch(sequence_length, drop_remainder=True)
      ds = ds.shuffle(100)
      ds = ds.repeat()
      ds = ds.batch(batch_size)
      ds = ds.map(
          lambda b: tf.nest.map_structure(tf.transpose, b)
      )  # Time major.
  
      return iter(tfds.as_numpy(ds))
  
imports: from collections.abc import Iterator, Mapping
  import numpy
  import tensorflow.compat.v2
  import tensorflow_datasets
label_desc: Use image node_id 1 for calling a global function with example usage: load(split) and returns: iter, dict