bojone / cdial-gpt-tf Goto Github PK

#! -*- coding: utf-8 -*-
# bert4keras加载CDial-GPT
# https://github.com/bojone/CDial-GPT-tf

import numpy as np
from bert4keras.models import build_transformer_model
from bert4keras.tokenizers import Tokenizer
from bert4keras.snippets import AutoRegressiveDecoder
from bert4keras.snippets import uniout
from tensorflow import keras
import json
import requests


config_path = './GPT_LCCC-large-tf/config.json'
checkpoint_path = './GPT_LCCC-large-tf/gpt_model.ckpt'
dict_path = './GPT_LCCC-large-tf/vocab.txt'

tokenizer = Tokenizer(dict_path, do_lower_case=True)  # 建立分词器
speakers = [
    tokenizer.token_to_id('[speaker1]'),
    tokenizer.token_to_id('[speaker2]')
]

model = build_transformer_model(
    config_path=config_path,
    checkpoint_path=checkpoint_path,
    model='gpt_openai'
)  # 建立模型，加载权重


print(type(model))
model.summary(line_length=150)
# model.save('./b4k_model')
# k_model = keras.models.load_model('./b4k_model')

mls_uri = 'http://cdial-gpt.t9k-sample.n.ksvc.tensorstack.net/v1/models/cdial-gpt:predict'

class ChatBot(AutoRegressiveDecoder):
    """基于随机采样对话机器人
    """
    @AutoRegressiveDecoder.wraps(default_rtype='probas')
    def predict(self, inputs, output_ids, states):
        token_ids, segment_ids = inputs
        curr_segment_ids = np.zeros_like(output_ids) + token_ids[0, -1]
        token_ids = np.concatenate([token_ids, output_ids], 1)
        segment_ids = np.concatenate([segment_ids, curr_segment_ids], 1)
        
        serving_url = mls_uri
        instances = [{'Input-Segment': segment_ids.tolist(), 'Input-Token': token_ids.tolist()}]
        data = {'signature_name': 'serving_default', 'instances': instances}
        data = json.dumps(data)
        headers = {"content-type": "application/json"}

        json_response = requests.post(serving_url, data=data, headers=headers)
        try:
            pred_res = json.loads(json_response.text)
        except json.decoder.JSONDecodeError:
            raise Exception("Invalid response: {}".format(json_response))
        print(type(pred_res))
        print(pred_res)
        pred = pred_res['predictions']
        return pred[:, -1]

    def response(self, texts, topk=5):
        token_ids = [tokenizer._token_start_id, speakers[0]]
        segment_ids = [tokenizer._token_start_id, speakers[0]]
        for i, text in enumerate(texts):
            ids = tokenizer.encode(text)[0][1:-1] + [speakers[(i + 1) % 2]]
            token_ids.extend(ids)
            segment_ids.extend([speakers[i % 2]] * len(ids))
            segment_ids[-1] = speakers[(i + 1) % 2]
        results = self.random_sample([token_ids, segment_ids], 1, topk)
        return tokenizer.decode(results[0])


chatbot = ChatBot(start_id=None, end_id=tokenizer._token_end_id, maxlen=32)
while True:
    raw_text = input(">>> ")
    print(chatbot.response([raw_text]))
"""
回复是随机的，例如：你还有我 | 那就不要爱我 | 你是不是傻 | 等等。
"""