从 BERT 获取嵌入查找结果

Getting embedding lookup result from BERT

在通过 BERT 传递我的令牌之前,我想对它们的嵌入执行一些处理(嵌入查找层的结果)。 HuggingFace BERT TensorFlow implementation 允许我们使用以下方式访问嵌入查找的输出:

import tensorflow as tf
from transformers import BertConfig, BertTokenizer, TFBertModel

bert_tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')

input_ids = tf.constant(bert_tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]
attention_mask = tf.stack([tf.ones(shape=(len(sent),)) for sent in input_ids])
token_type_ids = tf.stack([tf.ones(shape=(len(sent),)) for sent in input_ids])

config = BertConfig.from_pretrained('bert-base-uncased', output_hidden_states=True)
bert_model = TFBertModel.from_pretrained('bert-base-uncased', config=config)

result = bert_model(inputs={'input_ids': input_ids, 
                            'attention_mask': attention_mask, 
                            'token_type_ids': token_type_ids})
inputs_embeds = result[-1][0]  # output of embedding lookup

随后,可以处理 inputs_embeds,然后使用以下方法将其作为输入发送到同一模型:

inputs_embeds = process(inputs_embeds)  # some processing on inputs_embeds done here (dimensions kept the same)
result = bert_model(inputs={'inputs_embeds': inputs_embeds, 
                            'attention_mask': attention_mask, 
                            'token_type_ids': token_type_ids})
output = result[0]

其中 output 现在包含修改后输入的 BERT 输出。但是,这需要两次完全通过 BERT。而不是 运行 BERT 一直执行嵌入查找,我只想获得嵌入查找层的输出。 这可能吗?如果可能,怎么做?

将第一个输出 result[-1][0] 视为嵌入查找的结果实际上是不正确的。原始嵌入查找由以下公式给出:

embeddings = bert_model.bert.get_input_embeddings()
word_embeddings = embeddings.word_embeddings
inputs_embeds = tf.gather(word_embeddings, input_ids)

result[-1][0] 给出了嵌入查找 plus 位置嵌入和标记类型嵌入。上面的代码不需要完全通过 BERT,结果可以在输入到 BERT 的其余层之前进行处理。

编辑:要获得对任意 inputs_embeds 添加位置和标记类型嵌入的结果,可以使用:

full_embeddings = embeddings(inputs=[None, None, token_type_ids, inputs_embeds])

此处,embeddings 对象的 call 方法接受一个列表,该列表被馈送到 _embeddings 方法。第一个值为 input_ids,第二个值为 position_ids,第三个值为 token_type_ids,第四个值为 inputs_embeds。 (请参阅 here 了解更多详情。)如果您在一个输入中有多个句子,则可能需要设置 position_ids.