利用hugging face的Transformers实现多标签文本分类
多标签分类
文本分类是自然语言处理的基本任务之一。大部分文本分类是多分类，即数据有多个标签。实际工作或者项目中会遇到多标签的文本。笔者利用hugging face的Transformers实现多标签文本分类。笔者的tensorflow版本为2.4.0，transformers的版本为4.2.0
数据处理
利用transformers中的BertTokenizer对数据进行Tokenizer。代码如下：
def get_model_data(data, labels, max_seq_len=128):
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese", do_lower_case=True)
dataset_dict = {
"input_ids": [],
"attention_mask": [],
"label": []
}
assert len(data) == len(labels)
for i in range(len(data)):
sentence = data[i]
input_ids = tokenizer.encode(
sentence, # Sentence to encode.
add_special_tokens=True, # Add '[CLS]' and '[SEP]'
max_length=max_seq_len, # Truncate all sentences.
)
sentence_length = len(input_ids)
input_ids = pad_sequences([input_ids],
maxlen=max_seq_len,
dtype="long",
value=0,
truncating="post",
padding="post")
input_ids = input_ids.tolist()[0]
attention_mask = [1] * sentence_length + [0] * (max_seq_len - sentence_length)
dataset_dict["input_ids"].append(input_ids)
# dataset_dict["token_type_ids"].append(token_type_ids)
dataset_dict["attention_mask"].append(attention_mask)
dataset_dict["label"].append(labels[i])
for key in dataset_dict:
dataset_dict[key] = np.array(dataset_dict[key])
x = [
dataset_dict["input_ids"],
dataset_dict["attention_mask"],
]
y = dataset_dict["label"]
return x, y
多标签分类模型
利用Transformers搭建多标签分类模型。多标签分类在模型的最后一层全连接层的激活函数是sigmoid，而多分类的激活函数为softmax。多标签分类的损失函数为BinaryCrossentropy。代码如下：
class BertMultiClassifier(object):
def __init__(self, bert_model_name, label_num):
self.label_num = label_num
self.bert_model_name = bert_model_name
def get_model(self):
bert = TFBertModel.from_pretrained(self.bert_model_name)
input_ids = Input(shape=(None,), dtype=tf.int32, name="input_ids")
attention_mask = Input(shape=(None,), dtype=tf.int32, name="attention_mask")
outputs = bert(input_ids, attention_mask=attention_mask)[1]
cla_outputs = Dense(self.label_num, activation='sigmoid')(outputs)
model = Model(
inputs=[input_ids, attention_mask],
outputs=[cla_outputs])
return model
def create_model(bert_model_name, label_nums):
model = BertMultiClassifier(bert_model_name, label_nums).get_model()
optimizer = tf.keras.optimizers.Adam(lr=1e-5)
loss_object = tf.keras.losses.BinaryCrossentropy(from_logits=False)
model.compile(optimizer=optimizer, loss=loss_object,
metrics=['accuracy', tf.keras.metrics.Precision(),
tf.keras.metrics.Recall(),
tf.keras.metrics.AUC()]) # metrics=['accuracy']
return model
模型训练
利用tensorflow中的高阶API keras训练模型，模型保存为h5，同时将模型保存为pb模型。训练代码如下：
model.fit(train_x, train_y, epochs=args["epoch"], verbose=1,
batch_size=args["batch_size"],
callbacks=callbacks,
validation_data=(val_x, val_y),
validation_batch_size=args["batch_size"])
model_path = os.path.join("./output/model/", "mulclassifition.h5")
model.save_weights(model_path)
tf.keras.models.save_model(model, args["pbmodel_path"], save_format="tf", overwrite=True)
加载模型预测
一般训练模型可以直接加载做预测也可以用Tensorflow serving部署提供http服务，笔者分别介绍这两种方式。直接加载模型做预测的代码如下：
def predict(test_data, args, label_num):
# test_steps_per_epoch = len(test_data) // args["batch_size"]
tokenizer = BertTokenizer.from_pretrained("bert-base-chinese", do_lower_case=True)
testdata = get_model_data(test_data, tokenizer, args["max_length"])
print("testdata: ", testdata)
model = create_model(args['bert_model_name'], label_num)
model.load_weights("./output/model/mulclassifition.h5")
pred_logits = model.predict(testdata, batch_size=args["batch_size"])
pred = np.where(pred_logits >= 0.5, 1, 0).tolist()
return pred
HTTP服务
利用Tensorflow serving和Flask提供HTTP服务。代码如下：
@app.route("/multiclassfier", methods=['POST'])
def multiclassifier_pred():
data_para = json.loads(request.get_data(), encoding="utf-8")
sentence = data_para["sent"]
print("sentence: ", sentence)
# get model input
test_x = get_model_data(sentence, tokenizer, 256)
input_ids = test_x[0]
attention_mask = test_x[1]
data = json.dumps({"signature_name": "serving_default",
"inputs": {"input_ids": input_ids,
"attention_mask": attention_mask}})
headers = {"content-type": "application/json"}
result = requests.post("http://ip:port/v1/models/multiclass:predict", data=data, headers=headers)
if result.status_code == 200:
result = json.loads(result.text)
pred_logits = np.array(result["outputs"])
pred = np.where(pred_logits >= 0.5, 1, 0).tolist()
pred_encoder = label_encoder(pred, label)
return_result = {"code": 200, "sent": sentence, "label": pred_encoder[0]}
return jsonify(return_result)
else:
return jsonify({"code": result.status_code,
"message": traceback.format_exc()})
代码中的http://ip:port/v1/models/multiclass:predict是tensorflow serving的加载模型做预测的服务。采用docker 部署tensorflow serving部署服务。