如何利用gradio部署基于bert4keras的小说随机生成模型

from __future__ import print_function

import os
os.environ['TF_KERAS'] = '1'  # 必须使用tf.keras
import glob
import numpy as np
from bert4keras.backend import keras, K
from keras.backend import set_session

from bert4keras.layers import Loss
from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer, load_vocab
from bert4keras.optimizers import Adam
from bert4keras.snippets import sequence_padding, open
from bert4keras.snippets import DataGenerator, AutoRegressiveDecoder
from keras.models import Model
import tensorflow as tf  # 导入tf,备用
graph = tf.get_default_graph()
sess = tf.Session(graph=graph)
# graph = tf.Graph() 
# 基本参数
maxlen = 256+128
batch_size = 64
steps_per_epoch = 1000
epochs = 10000
# bert配置 autodl-tmp/chinese_roformer-v2-char_L-24_H-1024_A-16
config_path = '/root/autodl-tmp/chinese_roformer-v2-char_L-12_H-768_A-12/bert_config.json'
checkpoint_path = '/root/autodl-tmp/chinese_roformer-v2-char_L-12_H-768_A-12/bert_model.ckpt'
dict_path = '/root/autodl-tmp/chinese_roformer-v2-char_L-12_H-768_A-12/vocab.txt'

# 加载数据集
# autodl-nas/all_txt.txt
# 加载并精简词表,建立分词器
# token_dict, keep_tokens = load_vocab(
#     dict_path=dict_path,
#     simplified=True,
#     startswith=['[PAD]', '[UNK]', '[CLS]', '[SEP]'],
# )
# tokenizer = Tokenizer(token_dict, do_lower_case=True)
tokenizer = Tokenizer(dict_path, do_lower_case=True)


class CrossEntropy(Loss):
    """交叉熵作为loss,并mask掉输入部分
    """
    def compute_loss(self, inputs, mask=None):
        y_true, y_mask, y_pred = inputs
        y_true = y_true[:, 1:]  # 目标token_ids
        y_mask = y_mask[:, 1:]  # segment_ids,刚好指示了要预测的部分
        y_pred = y_pred[:, :-1]  # 预测序列,错开一位
        loss = K.sparse_categorical_crossentropy(y_true, y_pred)
        loss = K.sum(loss * y_mask) / K.sum(y_mask)
        return loss

class AutoTitle(AutoRegressiveDecoder):
    """seq2seq解码器
    """
    def __init__(self, start_id, end_id, maxlen, model):
        super().__init__(start_id, end_id, maxlen)
        self.model = model
        self.models = {}
    @AutoRegressiveDecoder.wraps(default_rtype='probas')
    def predict(self, inputs, output_ids, states):
        token_ids, segment_ids = inputs
        token_ids = np.concatenate([token_ids, output_ids], 1)
        segment_ids = np.concatenate([segment_ids, np.ones_like(output_ids)], 1)
        return self.last_token(self.model).predict([token_ids, segment_ids])

    def generate(self, text, topk=1):
        max_c_len = maxlen - self.maxlen
        token_ids, segment_ids = tokenizer.encode(text, maxlen=max_c_len)
        output_ids = self.beam_search([token_ids, segment_ids],
                                      topk=topk)  # 基于beam search
        return tokenizer.decode(output_ids)





# 在model加载前添加set_session

# graph = tf.get_default_graph()
# sess = tf.Session(graph=graph) 
# strategy = tf.distribute.MirroredStrategy()  # 建立单机多卡策略

# with strategy.scope():  # 调用该策略
set_session(sess)

bert = build_transformer_model(
                        config_path,
                        checkpoint_path=None,
                        model='roformer_v2',
                        # with_mlm='linear',
                        application='unilm',
                        return_keras_model=False
                    )

model = bert.model  # 这个才是keras模型
output = CrossEntropy(2)(model.inputs + model.outputs)

model = Model(model.inputs, output)
# model.compile(optimizer=Adam(1e-5))
# model.summary()
bert.load_weights_from_checkpoint(checkpoint_path)  # 必须最后才加载预训练权重

model.load_weights('/root/autodl-tmp/best_model.weights')
autotitle = AutoTitle(start_id=None, end_id=tokenizer._token_end_id, maxlen=128+64,model=model)
def model_predict(text):
    # 每次使用有关TensorFlow的请求时
    # in each request (i.e. in each thread):
    global sess
    global graph
    # graph = tf.get_default_graph()
    # sess = tf.Session(graph=graph)

    # set_session(sess)


    
    with graph.as_default():
        set_session(sess)
        summary_data = autotitle.generate(text)
    # K.clear_session()
    return summary_data

import gradio as gr

for i in range(2):
    print(model_predict("电荷交换"))
demo = gr.Interface(fn=model_predict, inputs="text", outputs="text")
gr.close_all()
demo.launch(server_port=6006)

你可能感兴趣的:(深度学习框架,深度学习,tensorflow,python)