code: def __init__(
      self, index_fname, batch_size, restore_state=None, text_tokens=256
  ):
      self.tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
      self.tokenizer.pad_token = "<|endoftext|>"
      self.tokenizer.add_special_tokens(
          {"sep_token": "<|sep|>", "pad_token": "<|pad|>"}
      )
  
      def map_fn(example):
          tokenizer = self.tokenizer
  
          def decode(x):
              return tokenizer(
                  ["<|endoftext|>" + i.decode() for i in x]
              )["input_ids"]
  
          texts = [
              decode(example["context_page_description"]),
              decode(example["context_section_description"]),
              decode(example["caption_reference_description"]),
              decode(example["caption_alt_text_description"]),
              decode(example["caption_attribution_description"]),
          ]
  
          output = []
  
          for text, dalle in zip(zip(*texts), example["dalle"]):
              all_text = list(itertools.chain(*text))[
                  -text_tokens + 1 :
              ]
  
              all_text += [tokenizer.pad_token_id] * (
                  (text_tokens - 1) - len(all_text)
              )
  
              assert len(all_text) == text_tokens - 1
  
              all_tokens = (
                  all_text
                  + [tokenizer.sep_token_id]
                  + list(dalle + tokenizer.vocab_size + 1)
              )
              output.append(all_tokens)
  
          return np.array(output)
  
      def tf_parse(example_proto):
          features = {
              "page_title": tf.io.FixedLenFeature([], tf.string),
              "section_title": tf.io.FixedLenFeature([], tf.string),
              "hierarchical_section_title": tf.io.FixedLenFeature(
                  [], tf.string
              ),
              "caption_reference_description": tf.io.FixedLenFeature(
                  [], tf.string
              ),
              "caption_attribution_description": tf.io.FixedLenFeature(
                  [], tf.string
              ),
              "caption_alt_text_description": tf.io.FixedLenFeature(
                  [], tf.string
              ),
              "mime_type": tf.io.FixedLenFeature([], tf.string),
              "context_page_description": tf.io.FixedLenFeature(
                  [], tf.string
              ),
              "context_section_description": tf.io.FixedLenFeature(
                  [], tf.string
              ),
              "dalle": tf.io.FixedLenFeature([1024], tf.int64),
          }
  
          parsed_features = tf.io.parse_single_example(
              example_proto, features
          )
  
          return parsed_features
  
      super().__init__(
          index_fname,
          batch_size,
          tf_parse,
          map_fn,
          restore_state=restore_state,
      )
  
file_path: mesh-transformer-jax/tfrecord_loader.py
imports: import jax
  import tensorflow
  import numpy
  from transformers import GPT2TokenizerFast
  import itertools
label_desc: Use image node_id 1 to create a new TFRecordWIT object from inherited base classes: TFRecordLoader with example: obj = TFRecordWIT(index_fname, batch_size, restore_state, text_tokens)