64 lines
2.2 KiB
Python
64 lines
2.2 KiB
Python
|
#! -*- coding: utf-8 -*-
|
|||
|
|
# NEZHA模型做闲聊任务
|
|||
|
|
# 测试脚本
|
|||
|
|
# 测试环境：tensorflow 2.5.3 + keras 2.3.1 + bert4keras 0.11
|
|||
|
|
import os
|
|||
|
|
os.environ['TF_KERAS'] = "1"
|
|||
|
|
|
|||
|
|
import numpy as np
|
|||
|
|
from bert4keras.backend import keras, K
|
|||
|
|
from bert4keras.models import build_transformer_model
|
|||
|
|
from bert4keras.tokenizers import Tokenizer
|
|||
|
|
from bert4keras.snippets import AutoRegressiveDecoder
|
|||
|
|
|
|||
|
|
|
|||
|
|
class ChatBot(AutoRegressiveDecoder):
|
|||
|
|
"""基于随机采样对话机器人
|
|||
|
|
"""
|
|||
|
|
def __init__(self, start_id, end_id, maxlen, model, tokenizer):
|
|||
|
|
super().__init__(start_id, end_id, maxlen)
|
|||
|
|
self.model = model
|
|||
|
|
self.tokenizer = tokenizer
|
|||
|
|
|
|||
|
|
@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)
|
|||
|
|
curr_segment_ids = np.ones_like(output_ids) - segment_ids[0, -1]
|
|||
|
|
segment_ids = np.concatenate([segment_ids, curr_segment_ids], 1)
|
|||
|
|
return self.model.predict([token_ids, segment_ids])[:, -1]
|
|||
|
|
|
|||
|
|
def response(self, texts, topk=5):
|
|||
|
|
token_ids, segment_ids = [self.tokenizer._token_start_id], [0]
|
|||
|
|
for i, text in enumerate(texts):
|
|||
|
|
ids = self.tokenizer.encode(text)[0][1:]
|
|||
|
|
token_ids.extend(ids)
|
|||
|
|
segment_ids.extend([i % 2] * len(ids))
|
|||
|
|
results = self.random_sample([token_ids, segment_ids], 1, topk)
|
|||
|
|
return self.tokenizer.decode(results[0])
|
|||
|
|
|
|||
|
|
|
|||
|
|
def build_chat_model(model_path, tokenizer):
|
|||
|
|
# nezha配置
|
|||
|
|
config_path = f'{model_path}config.json'
|
|||
|
|
checkpoint_path = f'{model_path}model.ckpt'
|
|||
|
|
|
|||
|
|
# 建立并加载模型
|
|||
|
|
model = build_transformer_model(
|
|||
|
|
config_path,
|
|||
|
|
checkpoint_path,
|
|||
|
|
model='nezha',
|
|||
|
|
application='lm',
|
|||
|
|
)
|
|||
|
|
chatbot = ChatBot(start_id=None, end_id=tokenizer._token_end_id, maxlen=32, model=model, tokenizer=tokenizer)
|
|||
|
|
return chatbot
|
|||
|
|
|
|||
|
|
if __name__ == '__main__':
|
|||
|
|
tokenizer = Tokenizer("../models/nezha_gpt_dialog/vocab.txt", do_lower_case=True)
|
|||
|
|
chatbot = build_chat_model('../models/nezha_gpt_dialog', tokenizer)
|
|||
|
|
text_list = ['绿遍山原白满川', '子规声里雨如烟']
|
|||
|
|
print(chatbot.response(text_list))
|
|||
|
|
|
|||
|
|
|
|||
|
|
|