code: def bert_preprocess(batch, tokenizer, key_a, key_b=None, max_len=128):
      """Tokenize and convert text to model inputs in a BERT format."""
      batch_size = batch["idx"].shape[0]
      input_ids = onp.zeros((batch_size, max_len), dtype=onp.int32)
      type_ids = onp.zeros((batch_size, max_len), dtype=onp.int32)
      for i in range(batch_size):
          sentence_a = batch[key_a][i]
          tokens_a = (
              [101]
              + tokenizer.convert_tokens_to_ids(
                  tokenizer.tokenize(sentence_a)
              )
              + [102]
          )
  
          if key_b is not None:
              sentence_b = batch[key_b][i]
              tokens_b = tokenizer.convert_tokens_to_ids(
                  tokenizer.tokenize(sentence_b)
              ) + [102]
          else:
              tokens_b = []
  
          ex_input_ids = (tokens_a + tokens_b)[:max_len]
          ex_type_ids = ([0] * len(tokens_a) + [1] * len(tokens_b))[
              :max_len
          ]
  
          input_ids[i, : len(ex_input_ids)] = ex_input_ids
          type_ids[i, : len(ex_type_ids)] = ex_type_ids
      return (
          input_ids,
          type_ids,
          input_ids > 0,
          batch["label"],
          onp.ones(batch_size),
      )
  
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 4 for calling a global function with example usage: bert_preprocess(batch, tokenizer, key_a, key_b, max_len) and returns: input_ids,  type_ids,  unknown,  batch,  onp