def run_one_hot_embeddings(one_hot_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"):
      # branched from modeling.embedding_lookup
      embedding_table = tf.get_variable(
          name="word_embeddings",
          shape=[config.vocab_size, config.hidden_size],
          initializer=modeling.create_initializer(config.initializer_range))

      flat_input_ids = tf.reshape(one_hot_input_ids, [-1, config.vocab_size])
      output = tf.matmul(flat_input_ids, embedding_table)

      input_shape = modeling.get_shape_list(one_hot_input_ids)

      output = tf.reshape(output, input_shape[0:-1] + [config.hidden_size])

      return (output, embedding_table)
Extract only the word embeddings of the original BERT model.
language/bert_extraction/steal_bert_classifier/embedding_perturbations/embedding_util.py