def template_to_input_tensor(template, flex_input, config, tokenizer,
                             max_seq_length):
  """This function converts a template to a corresponding input tensor."""
  template_shards = template.split(" <piece> ")
  embedding_input_ids = []
  index_array = []
  flex_input_mask = []
  # bert_input_mask will be applied during self-attention, it is mainly applied
  # to PADDING tokens.
  bert_input_mask = []
  # token_type_ids are needed for pairwise classification tasks like MNLI.
  token_type_ids = []
  current_index = 0
  current_token_type = 0

  batch_size = flex_input.shape[0]

  for shard in template_shards:
    # Precautionary measure to prevent errors for templates > max_seq_length
    if current_index >= max_seq_length:
      tf.logging.warning("Template is longer than max_seq_length = %d",
                         max_seq_length)
      break

    shard_token, shard_quantity, _ = parse_shard(
        shard=shard, current_index=current_index, max_seq_length=max_seq_length)

    update_bert_input_mask(
        shard_token=shard_token,
        shard_quantity=shard_quantity,
        tokenizer=tokenizer,
        bert_input_mask=bert_input_mask)

    current_token_type = update_token_type_ids(
        shard_token=shard_token,
        shard_quantity=shard_quantity,
        tokenizer=tokenizer,
        token_type_ids=token_type_ids,
        current_token_type=current_token_type)

    if shard_token == "[EMPTY]":
      flex_input_mask.extend([1 for _ in range(shard_quantity)])
      # [EMPTY] shard tokens indicate using the flex_input matrix
      index_array.extend(
          list(range(current_index, current_index + shard_quantity)))
      # Move the index forward by the number of tokens done
      current_index += shard_quantity
    else:
      if shard_token in tokenizer.vocab:
        token_ids = [tokenizer.vocab[shard_token]]
      else:
        token_ids = tokenizer.convert_tokens_to_ids(
            tokenizer.tokenize(shard_token))
      # repeat sequence of tokens according to number of shards
      for _ in range(shard_quantity):
        # These tokens are fixed, so we use 0 in their positions in the mask
        flex_input_mask.extend([0 for _ in token_ids])
        index_array.extend([
            max_seq_length + len(embedding_input_ids) + token_num
            for token_num in range(len(token_ids))
        ])
        # Move the index forward by the number of tokens done
        current_index += len(token_ids)
        embedding_input_ids.extend(token_ids)

  if embedding_input_ids:
    # We store the embedding table to carry out nearest neighbour computes
    embeddings, embed_var = run_bert_embeddings(
        input_ids=tf.constant([embedding_input_ids]), config=config)
    batched_embeddings = tf.tile(embeddings, [batch_size, 1, 1])
    all_vectors = tf.concat([flex_input, batched_embeddings], axis=1)
  else:
    # Just a dummy run to get the final embedding table
    _, embed_var = run_bert_embeddings(
        input_ids=tf.constant([[0]]), config=config)
    all_vectors = flex_input

  # Build final masks by replication in the batch axis and converting to boolean
  flex_input_mask = tf.cast(
      tf.tile(tf.constant([flex_input_mask], dtype=tf.int32), [batch_size, 1]),
      tf.bool)

  bert_input_mask = tf.tile(
      tf.constant([bert_input_mask], dtype=tf.int32), [batch_size, 1])

  token_type_ids = tf.tile(
      tf.constant([token_type_ids], dtype=tf.int32), [batch_size, 1])

  final_input = tf.gather(all_vectors, tf.constant(index_array), axis=1)
  return (final_input, embed_var, flex_input_mask, bert_input_mask,
          token_type_ids)
This function converts a template to a corresponding input tensor.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py