[BERT]tensorflow获取中文句子向量(基于estimator和基于bertServer两种方法)

目的：将数据集输入BERT，直接获取句子向量做后续训练

数据集格式：一句话一行（已过滤掉各种符号）

目录

一、利用BertServer

二、保存bert得到的句子向量-tensorflow基于estimator版

一、利用BertServer

环境：python3.6 + tensorflow1.14

参考：使用Bert得到句向量简单总结_介样不勾的博客-CSDN博客_bert获取句子向量

我的数据保存：

1. 安装BertServer：

pip install bert-serving-client
pip install bert-serving-server

输出：

2. 启动服务端

bert-serving-start -model_dir ./chinese_L-12_H-768_A-12/ -num_worker 2 -max_seq_len 64

 max_seq_len的取值是经过统计数据集之后得到的长度，num_worker=2是因为服务器上有两张卡是空闲的可以被我拿来用~

 3. 测试一下

import numpy as np
from bert_serving.client import BertClient
bc = BertClient()
sentence = bc.encode(['左侧 肾上腺 增粗 PET 上 异常 放射性 浓聚 考虑 为 肾上腺 转移'])
print(sentence.shape)
print(sentence)

python bert2senvec.py

 输出：

(1, 768)
[[ 2.54443586e-01  3.98949206e-01 -2.38077953e-01 -2.89931267e-01
   3.15846503e-01 -7.56896377e-01 -1.96446180e-01 -8.28849748e-02
  ...
  -2.79337615e-01  2.05526114e-01 -3.87653321e-01 -2.14681253e-01
  -7.32321620e-01  6.60689592e-01 -8.09417069e-01  6.66649314e-04]]

目前看来还比较成功

4. 批量处理一下

参考：使用bert-service获取句向量和相似度计算_马苏比拉米G的博客-CSDN博客_bert获取句向量

from bert_serving.client import BertClient
import numpy as np
import os


def get_sen_vec(file, new_name):
    bc = BertClient()
    sen_vec = []
    len = 0
    with open(file, 'r', encoding='utf-8') as f:
        for line in f.readlines():
            len += 1
            if line.strip('\n').strip() != '':
                # 如果不是空
                sen = bc.encode([line])
                sen_vec.append(sen)               
    print(f"{file}一共有{len}句话")
    np.save(new_name, np.concatenate(sen_vec, axis=0))


if __name__ == '__main__':
    dir = "file.txt"
    output_dir = "file_vec"
    get_sen_vec(file_name, new_name)
    newfile = np.load(new_name+'.npy')
    print(f"{new_name}的维度是{newfile.shape}")

启动后运行

nohup python bert2senvec.py > run.log 2>&1 &

nohup：后台自动运行

所有提示信息输入到run.log中，输出消息如下：

nohup: ignoring input
file.txt一共有5515句话
file_vec的维度是(5515, 768)

目前看来非常成功，不知道能否顺利拿来应用，如果有后续持续更新

尝试了一下超出max_seq_len的句子如何处理，结果是输出维度依旧是(1,768)

bc = BertClient()
line = '左侧 肾上腺 增粗 PET 上 异常 放射性 浓聚 考虑 为 肾上腺 转移左侧 肾上腺 增粗 PET 上 异常 放射性 浓聚 考虑 为 肾上腺 转移左侧 肾上腺 增粗 PET 上 异常 放射性 浓聚 考虑 为 肾上腺 转移'
sentence = bc.encode([line])
print(sentence.shape)
print(sentence)

$ python bert2senvec.py 
(1, 768)
...

二、保存bert得到的句子向量-tensorflow基于estimator版

环境：tensorflow1.15 + tensorflow-bert + python3.7

1. 模型部分

def create_model(bert_config, is_training, input_ids, input_mask, segment_ids,
                 labels, use_one_hot_embeddings, use_tpu):
    """Creates a classification model."""
    model = modeling.BertModel(
        config=bert_config,
        is_training=is_training,
        input_ids=input_ids,
        input_mask=input_mask,
        token_type_ids=segment_ids,
        use_one_hot_embeddings=use_one_hot_embeddings)

    output_final_layer = model.get_sequence_output()
    embedding = tf.squeeze(output_final_layer[:, 0:1, :], axis=1)
    # embedding即为需要保存的对象


    return embedding

2. 返回到model_fn部分

def model_fn_builder(bert_config, init_checkpoint, learning_rate,
                     num_train_steps, num_warmup_steps, use_tpu,
                     use_one_hot_embeddings):
  """Returns `model_fn` closure for TPUEstimator."""

  def model_fn(features, labels, mode, params):
    """The `model_fn` for TPUEstimator."""
    del params  # unused
    del labels  # unused
    is_training = (mode == tf.estimator.ModeKeys.TRAIN)


    features = pad_feature(features)
    
    input_ids = features["input_ids"]
    input_mask = features["input_mask"]
    segment_ids = features["segment_ids"]
    label_ids = features["label_ids"]
    is_real_example = None

    if "is_real_example" in features:
      is_real_example = tf.cast(features["is_real_example"], dtype=tf.float32)
    else:
      is_real_example = tf.ones(tf.shape(label_ids), dtype=tf.float32)

    sentence_embeddings = create_model(bert_config, is_training, input_ids, input_mask, segment_ids, label_ids, use_one_hot_embeddings, use_tpu)

    predictions = {
      'sentence_embeddings': sentence_embeddings, # 即为所求 
    }
 
    output_spec = tf.estimator.EstimatorSpec(
      mode=mode,
      predictions=predictions
    )

    return output_spec

  return model_fn

3. main函数中保存

ts = experiment_utils.run_experiment(
      model_fn=model_fn,
      train_input_fn=train_input_fn,
      eval_input_fn=train_input_fn,
      params=params)
# 此处对estimator代码进行过修改，直接将estimator返回了

tensor_pred = ts.predict(input_fn=train_input_fn)

for pred in tensor_pred:
    vec.append(pred['sentence_embeddings'])
np.save(f"./{FLAGS.continual_learning}_embeddings.npy", vec)
# 此npy文件即为保存的向量