def update_token_type_ids(shard_token, shard_quantity, tokenizer,
                          token_type_ids, current_token_type):
  """Utility function to update token_type_ids list for an incoming shard."""
  if shard_token == "[EMPTY]":
    token_type_ids.extend([current_token_type for _ in range(shard_quantity)])

  elif shard_token in tokenizer.vocab and shard_token == "[SEP]":
    token_type_ids.extend([current_token_type for _ in range(shard_quantity)])
    current_token_type = 1 - current_token_type

  elif shard_token in tokenizer.vocab and shard_token != "[SEP]":
    token_type_ids.extend([current_token_type for _ in range(shard_quantity)])

  else:
    # BERT tokenizer will not retain [SEP] special tokens so current_token_type
    # will not swap during this shard.
    shard_token_pieces = tokenizer.tokenize(shard_token)
    for _ in range(shard_quantity):
      token_type_ids.extend([current_token_type for _ in shard_token_pieces])

  return current_token_type
Utility function to update token_type_ids list for an incoming shard.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py