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)
Input pipeline for fine-tuning BERT on GLUE tasks.
trax/supervised/pretrain_finetune.py