ai_platform_nlu/nlp/text_classification.py

#! -*- coding:utf-8 -*-
import os

# os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"  #（保证程序cuda序号与实际cuda序号对应）
# os.environ['CUDA_VISIBLE_DEVICES'] = "0, 1"  #（代表仅使用第0，1号GPU）
os.environ['TF_KERAS'] = '1'

import numpy as np
import pandas as pd
import tensorflow as tf
from logzero import logger
from bert4keras.backend import keras, set_gelu
from bert4keras.tokenizers import Tokenizer
from bert4keras.models import build_transformer_model
from bert4keras.optimizers import Adam, extend_with_piecewise_linear_lr
from bert4keras.snippets import sequence_padding, DataGenerator
from bert4keras.snippets import open
from keras.layers import Lambda, Dense

config = tf.compat.v1.ConfigProto()
config.gpu_options.allow_growth = True  # 按需分配显存
tf_session = tf.compat.v1.Session(graph=tf.compat.v1.get_default_graph(), config=config)
tf.compat.v1.keras.backend.set_session(tf_session)

from sklearn.model_selection import train_test_split

set_gelu('tanh')  # 切换gelu版本

MAX_LEN = 128
BATCH_SIZE = 32
AdamLR = extend_with_piecewise_linear_lr(Adam, name='AdamLR')


# config_path = '/home/zhaojh/pretrain_models/roberta_base/bert_config.json'
# checkpoint_path = '/home/zhaojh/pretrain_models/roberta_base/bert_model.ckpt'
# dict_path = '/home/zhaojh/pretrain_models/roberta_base/vocab.txt'
# 建立分词器
# tokenizer = Tokenizer(dict_path, do_lower_case=True)

# def label_to_digit(data_df: pd.DataFrame, labels):
#     data = list()
#     for i in range(data_df.shape[0]):
#         data.append((data_df.iloc[i]['text'], int(labels.index(data_df.iloc[i]['label']))))
#     return data

# def load_data(all_data:pd.DataFrame):
#     """加载数据
#     单条格式：(文本, 标签id)
#     """
#     logger.info(f"数据集大小：{all_data.shape}")
#     assert ('text' in all_data.columns and 'label' in all_data.columns)
#     use_data = all_data[['text', 'label']].copy()
#     labels = use_data.label.unique().tolist()
#     train, valid = train_test_split(use_data, test_size=0.3, random_state=42, shuffle=True)
#     valid, test = train_test_split(valid, test_size=0.7, random_state=42, shuffle=True)
#     # 转换数据集
#     train_data = label_to_digit(train, labels)
#     valid_data = label_to_digit(valid, labels)
#     test_data = label_to_digit(test, labels)
#     train_generator = MyDataGenerator(train_data, BATCH_SIZE)
#     valid_generator = MyDataGenerator(valid_data, BATCH_SIZE)
#     test_generator = MyDataGenerator(test_data, BATCH_SIZE)
#     return train_generator, valid_generator, test_generator, labels


# class MyDataGenerator(DataGenerator):
#     """数据生成器
#     """
#     def __iter__(self, random=False):
#         batch_token_ids, batch_segment_ids, batch_labels = [], [], []
#         for is_end, (text, label) in self.sample(random):
#             token_ids, segment_ids = tokenizer.encode(text, maxlen=MAX_LEN)
#             batch_token_ids.append(token_ids)
#             batch_segment_ids.append(segment_ids)
#             batch_labels.append([label])
#             if len(batch_token_ids) == self.batch_size or is_end:
#                 batch_token_ids = sequence_padding(batch_token_ids)
#                 batch_segment_ids = sequence_padding(batch_segment_ids)
#                 batch_labels = sequence_padding(batch_labels)
#                 yield [batch_token_ids, batch_segment_ids], batch_labels
#                 batch_token_ids, batch_segment_ids, batch_labels = [], [], []


def build_model(config_path, ckpt_path, num_classes, lr):
    # 加载预训练模型
    bert = build_transformer_model(
        config_path=config_path,
        checkpoint_path=ckpt_path,
        model='bert',
        return_keras_model=False,
    )

    output = Lambda(lambda x: x[:, 0], name='CLS-token')(bert.model.output)
    if num_classes == 2:
        activitor = 'sigmoid'
    else:
        activitor = 'softmax'
    output = Dense(
        units=num_classes,
        activation=activitor,
        kernel_initializer=bert.initializer
    )(output)

    model = keras.models.Model(bert.model.input, output)

    # 派生为带分段线性学习率的优化器。
    AdamLR = extend_with_piecewise_linear_lr(Adam, name='AdamLR')

    model.compile(
        loss='sparse_categorical_crossentropy',
        optimizer=AdamLR(learning_rate=lr, lr_schedule={
            1000: 1,
            2000: 0.1
        }),
        metrics=['accuracy'],
    )
    return model


def evaluate(data, model):
    total, right = 0., 0.
    for x_true, y_true in data:
        y_pred = model.predict(x_true).argmax(axis=1)
        y_true = y_true[:, 0]
        total += len(y_true)
        right += (y_true == y_pred).sum()
    return right / total


# class Evaluator(keras.callbacks.Callback):
#     """评估与保存
#     """
#     def __init__(self, model, valid_generator, test_generator):
#         super().__init__()
#         self.best_val_acc = 0.
#         self.model = model
#         self.valid = valid_generator
#         self.test = test_generator

#     def on_epoch_end(self, epoch, logs=None):
#         val_acc = evaluate(self.valid, self.model)
#         if val_acc > self.best_val_acc:
#             self.best_val_acc = val_acc
#             self.model.save_weights('../models/text_classifier/best_model.weights')
#         test_acc = evaluate(self.test, self.model)
#         logger.info(f"val_acc: {val_acc}, best_val_acc:{self.best_val_acc}, test_acc: {test_acc}")


# def train_model(train_generator, valid_generator, test_generator, num_classes, epochs, lr):
#     model = build_model(num_classes, lr)
#     evaluator = Evaluator(model, valid_generator, test_generator)

#     model.fit(
#         train_generator.forfit(),
#         steps_per_epoch=len(train_generator),
#         epochs=epochs,
#         callbacks=[evaluator]
#     )
#     model.load_weights('../models/text_classifier/best_model.weights')
#     return model


def run_cls(model, tokenizer, test_text: str, labels=["negative", "positive"]):
    token_ids, segment_ids = tokenizer.encode(test_text, maxlen=MAX_LEN)
    tok_ids = sequence_padding([token_ids])
    seg_ids = sequence_padding([segment_ids])
    predicted_data = model.predict([tok_ids, seg_ids])[0]
    return labels[np.argmax(predicted_data)]


# def train(df:pd.DataFrame, epochs:int=20, lr:float=1e-3):
#     train, valid, test, labels = load_data(df)
#     model = train_model(train, valid, test, len(labels), epochs, lr)
#     # score = evaluate(test.forfit(), model)
#     model.save('../models/text_classifier/best_model.h5')
#     return model, labels

def load_cls_model(config_path: str, ckpt_path: str, weight_path: str):
    """load local classification model

    Args:
        config_path (str): config.json
        ckpt_path (str): ckpt
        weight_path (str): best weight

    Returns:
        keras.Model: pretrained model
    """
    model = build_model(config_path, ckpt_path, 2, 0.001)
    model.load_weights(weight_path)
    return model


if __name__ == '__main__':
    from keras.models import load_model

    dict_path = '../models/tokenizer/vocab.txt'
    tokenizer = Tokenizer(dict_path, do_lower_case=True)
    model = load_model('../models/text_classifier/best_model.h5')
    rst = run_cls(model, tokenizer, "这部电影太棒了")
    print(rst)