使用BERT对句子进行向量化(TensorFlow版和Pytorch版)

背景

用BERT对句子进行向量化

实施

TensorFlow版直接用肖涵博士的bert-as-service。使用方法真的很小白，简单概括为2点：server和client安装。

pip install bert-serving-server  # server
pip install bert-serving-client  # client, independent of `bert-serving-server`

在server安装完后，启动服务，比如：bert-serving-start -model_dir /home/pretained_models/chinese_wwm_ext_L-12_H-768_A-12 -num_worker=4
通过model_dir参数可以自行指定不同类型的BERT的模型路径，我这里使用的是哈工大发布的WWM-EXT版。在client上的测试代码：

def test_bert_tf(string):
    from bert_serving.client import BertClient
    bc = BertClient()
    s_encode = bc.encode([string])
    print(s_encode[0])

上述方案虽然简单易于上手，但是个人还是觉自己动手更香，比如基于huggingface的transformers。如何验证呢？就以bert-as-service编码得到的句向量作为标准值。将相同的文本输入到transformers试图得到与bert-as-service方案相同的句向量。

由于bert-as-service默认的句向量构造方案是取倒数第二层的隐状态值在token上的均值，即选用的层数是倒数第2层，池化策略是REDUCE_MEAN。

import torch
import pdb
from transformers import AutoConfig
from transformers import BertTokenizer, BertModel, BertConfig

UNCASE = "/home/pretained_models/chinese_wwm_ext_pytorch"
VOCAB = "vocab.txt"
tokenizer = BertTokenizer.from_pretrained(UNCASE + "/" + VOCAB)

model = BertModel.from_pretrained(UNCASE, output_hidden_states = True) # 如果想要获取到各个隐层值需要如此设置
model.eval()
string = '写代码不香吗'
string1 = "[CLS]" + string + "[SEP]"

# Convert token to vocabulary indices
tokenized_string = tokenizer.tokenize(string1)
tokens_ids = tokenizer.convert_tokens_to_ids(tokenized_string)
# Convert inputs to PyTorch tensors
tokens_tensor = torch.tensor([tokens_ids])
outputs = model(tokens_tensor) # encoded_layers, pooled_output

if model.config.output_hidden_states:
    hidden_states = outputs[2]
    # last_layer = outputs[-1]
    second_to_last_layer = hidden_states[-2]
    # 由于只要一个句子，所以尺寸为[1, 10, 768]
    token_vecs = second_to_last_layer[0]
    print(token_vecs.shape)
    # Calculate the average of all input token vectors.
    sentence_embedding = torch.mean(token_vecs, dim=0)
    print(sentence_embedding.shape)
    print(sentence_embedding[0:10])

print("tf version-----")
from bert_serving.client import BertClient
bc = BertClient()
s_encode = bc.encode([string])
print(s_encode[0].shape)
# pdb.set_trace()
print(s_encode[0][0:10])

结果如下：

从向量的前10维可以看出，两者向量是相同的。那么进一步计算二者的余弦相似度的结果：

tf_tensor = torch.tensor(s_encode[0])
similarity = torch.cosine_similarity(sentence_embedding, tf_tensor, dim=0)
print(similarity)

余弦相似度为1，所以两个向量相同。