code: def __init__(
      self,
      index_fname,
      batch_size,
      parse_fn,
      map_fn=None,
      restore_state=None,
  ):
      if restore_state is not None:
          self.file_idx = restore_state["file_idx"]
          self.file_idx_init = False
          self.used = restore_state["used"]
      else:
          self.file_idx = 0
          self.file_idx_init = True
          self.used = []
  
      self.index = open(index_fname).read().splitlines()
      self.clean_index = list(
          filter(lambda x: x not in self.used, self.index)
      )
      self.bs = batch_size
      # self.seq = sample_size
      self.parse_fn = parse_fn
  
      if map_fn:
          self.map_fn = map_fn
      else:
          self.map_fn = lambda x: x
  
      self.sample_fn = self.sample_once()
  
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 TFRecordLoader object  with example: obj = TFRecordLoader(index_fname, batch_size, parse_fn, map_fn, restore_state)