code: def glue_inputs(
      dataset_name=gin.REQUIRED,
      batch_size=16,
      eval_batch_size=None,
      data_dir=None,
      max_len=128,
      tokenizer=bert_tokenizer,
  ):
      """Input pipeline for fine-tuning BERT on GLUE tasks."""
      if callable(
          tokenizer
      ):  # If we pass a function, e.g., through gin, call it.
          tokenizer = bert_tokenizer()
  
      eval_split = tfds.Split.VALIDATION
      if dataset_name == "glue/mnli":
          eval_split = "validation_matched"
          # TODO(kitaev): Support diagnostic dataset (AX)
  
      keys_lookup = {
          "glue/cola": ("sentence", None),
          "glue/sst2": ("sentence", None),
          "glue/mrpc": ("sentence1", "sentence2"),
          "glue/qqp": ("question1", "question2"),
          "glue/stsb": ("sentence1", "sentence2"),
          "glue/mnli": (
              "premise",
              "hypothesis",
          ),  # TODO(kitaev): swap the two?
          "glue/qnli": (
              "question",
              "sentence",
          ),  # TODO(kitaev) swap the two?
          "glue/rte": ("sentence1", "sentence2"),
          "glue/wnli": ("sentence1", "sentence2"),
      }
  
      key_a, key_b = keys_lookup[dataset_name]
  
      preprocess_fn = functools.partial(
          bert_preprocess,
          tokenizer=tokenizer,
          key_a=key_a,
          key_b=key_b,
          max_len=max_len,
      )
      return tfds_inputs(  # TODO(piotrekp1): use data_streams instead
          dataset_name=dataset_name,
          preprocess_fun=preprocess_fn,
          batch_size=batch_size,
          eval_batch_size=eval_batch_size,
          data_dir=data_dir,
          train_split=tfds.Split.TRAIN,
          eval_split=eval_split,
      )
  
file_path: trax/trax/supervised/pretrain_finetune.py
imports: import functools
  import gin
  import numpy
  import tensorflow_datasets
  from trax.data.inputs import Inputs
label_desc: Use image node_id 5 for calling a global function with example usage: glue_inputs(dataset_name, batch_size, eval_batch_size, data_dir, max_len, tokenizer) and returns: tfds_inputs