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)
Tokenize and convert text to model inputs in a BERT format.
trax/supervised/pretrain_finetune.py