code: def load_shakespeare_corpus(token_type="char"):
      """Load the tiny_shakespeare TFDS and return its corpus and vodabulary."""
      data = tfds.as_numpy(
          tfds.load(name="tiny_shakespeare", batch_size=-1)
      )
      train_data = data["train"]
      input_string = train_data["text"][
          0
      ].decode()  # decode binary string
      re.sub("[^A-Za-z]+", " ", input_string.strip().lower())
      lines = input_string.splitlines()
  
      token_list = tokenize(lines, token_type)
  
      vocab = Vocabulary(token_list)
      corpus = [vocab[tk] for tokens in token_list for tk in tokens]
      return corpus, vocab
  
imports: import collections
  import random
  import re
  import jax
  import jax.numpy
  import tensorflow_datasets
  import objax
  from objax.functional import one_hot
  from objax.zoo.rnn import RNN
label_desc: Use image node_id 4 for calling a global function with example usage: load_shakespeare_corpus(token_type) and returns: corpus,  vocab