def update_bert_input_mask(shard_token, shard_quantity, tokenizer,
                           bert_input_mask):
  """Utility function to update the BERT input mask for an incoming shard."""
  if shard_token == "[EMPTY]":
    bert_input_mask.extend([1 for _ in range(shard_quantity)])

  elif shard_token in tokenizer.vocab and shard_token == "[PAD]":
    bert_input_mask.extend([0 for _ in range(shard_quantity)])

  elif shard_token in tokenizer.vocab and shard_token != "[PAD]":
    bert_input_mask.extend([1 for _ in range(shard_quantity)])

  else:
    shard_token_pieces = tokenizer.tokenize(shard_token)
    for _ in range(shard_quantity):
      bert_input_mask.extend([1 for _ in shard_token_pieces])
Utility function to update the BERT input mask for an incoming shard.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py