code: def get_training_data(
      seq_len: int = 32,
      dataset_name: str = "tiny_shakespeare",
      truncate: int = -1,
  ):
      """
      ### Load Dataset
  
      :param seq_len: is the sequence length of a single training sample
      :param dataset_name: is the name of the dataset
      :return: the dataset
      """
  
      ds = DATASETS[dataset_name]
      # Load the content
      text = load_text(lab.get_data_path() / ds["file"], ds["url"])
      # Tokenize
      tokenizer = get_tokenizer()
      tokens = tokenizer.encode_batch([text])[0]
  
      if truncate > 0:
          token_ids = tokens.ids[: truncate * seq_len]
      else:
          token_ids = tokens.ids
  
      #
      return NeoXDataset(token_ids, seq_len)
  
file_path: annotated_deep_learning_paper_implementations/labml_nn/neox/utils/text_dataset.py
imports: from pathlib import PurePath, Path
  from typing import Optional, List
  import torch
  import torch.utils.data
  from labml import lab
  from labml import monit
  from labml.logger import inspect
  from labml.utils.download import download_file
  from labml_nn.neox.tokenizer import get_tokenizer
label_desc: Use image node_id 2 for calling a global function with example usage: get_training_data(seq_len, dataset_name, truncate) and returns: NeoXDataset