def run_bert_embeddings(input_ids, config):
  """Extract only the word embeddings of the original BERT model."""
  with tf.variable_scope("bert", reuse=tf.compat.v1.AUTO_REUSE):
    with tf.variable_scope("embeddings"):
      # Perform embedding lookup on the word ids.
      embedding_output, embedding_var = modeling.embedding_lookup(
          input_ids=input_ids,
          vocab_size=config.vocab_size,
          embedding_size=config.hidden_size,
          initializer_range=config.initializer_range,
          word_embedding_name="word_embeddings",
          use_one_hot_embeddings=False)
      return embedding_output, embedding_var
Extract only the word embeddings of the original BERT model.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py