87 lines
3.3 KiB
Python
87 lines
3.3 KiB
Python
|
|
|||
|
|
import numpy as np
|
|||
|
|
import settings
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_random_poetry(tokenizer, model, s=''):
|
|||
|
|
"""
|
|||
|
|
随机生成一首诗
|
|||
|
|
:param tokenizer: 分词器
|
|||
|
|
:param model: 用于生成古诗的模型
|
|||
|
|
:param s: 用于生成古诗的起始字符串，默认为空串
|
|||
|
|
:return: 一个字符串，表示一首古诗
|
|||
|
|
"""
|
|||
|
|
# 将初始字符串转成token
|
|||
|
|
token_ids = tokenizer.encode(s)
|
|||
|
|
# 去掉结束标记[SEP]
|
|||
|
|
token_ids = token_ids[:-1]
|
|||
|
|
while len(token_ids) < settings.MAX_LEN:
|
|||
|
|
# 进行预测，只保留第一个样例（我们输入的样例数只有1）的、最后一个token的预测的、不包含[PAD][UNK][CLS]的概率分布
|
|||
|
|
output = model(np.array([token_ids, ], dtype=np.int32))
|
|||
|
|
_probas = output.numpy()[0, -1, 3:]
|
|||
|
|
del output
|
|||
|
|
# print(_probas)
|
|||
|
|
# 按照出现概率，对所有token倒序排列
|
|||
|
|
p_args = _probas.argsort()[::-1][:100]
|
|||
|
|
# 排列后的概率顺序
|
|||
|
|
p = _probas[p_args]
|
|||
|
|
# 先对概率归一
|
|||
|
|
p = p / sum(p)
|
|||
|
|
# 再按照预测出的概率，随机选择一个词作为预测结果
|
|||
|
|
target_index = np.random.choice(len(p), p=p)
|
|||
|
|
target = p_args[target_index] + 3
|
|||
|
|
# 保存
|
|||
|
|
token_ids.append(target)
|
|||
|
|
if target == 3:
|
|||
|
|
break
|
|||
|
|
return tokenizer.decode(token_ids)
|
|||
|
|
|
|||
|
|
|
|||
|
|
def generate_acrostic(tokenizer, model, head):
|
|||
|
|
"""
|
|||
|
|
随机生成一首藏头诗
|
|||
|
|
:param tokenizer: 分词器
|
|||
|
|
:param model: 用于生成古诗的模型
|
|||
|
|
:param head: 藏头诗的头
|
|||
|
|
:return: 一个字符串，表示一首古诗
|
|||
|
|
"""
|
|||
|
|
# 使用空串初始化token_ids，加入[CLS]
|
|||
|
|
token_ids = tokenizer.encode('')
|
|||
|
|
token_ids = token_ids[:-1]
|
|||
|
|
# 标点符号，这里简单的只把逗号和句号作为标点
|
|||
|
|
punctuations = ['，', '。']
|
|||
|
|
punctuation_ids = {tokenizer.token_to_id(token) for token in punctuations}
|
|||
|
|
# 缓存生成的诗的list
|
|||
|
|
poetry = []
|
|||
|
|
# 对于藏头诗中的每一个字，都生成一个短句
|
|||
|
|
for ch in head:
|
|||
|
|
# 先记录下这个字
|
|||
|
|
poetry.append(ch)
|
|||
|
|
# 将藏头诗的字符转成token id
|
|||
|
|
token_id = tokenizer.token_to_id(ch)
|
|||
|
|
# 加入到列表中去
|
|||
|
|
token_ids.append(token_id)
|
|||
|
|
# 开始生成一个短句
|
|||
|
|
while True:
|
|||
|
|
# 进行预测，只保留第一个样例（我们输入的样例数只有1）的、最后一个token的预测的、不包含[PAD][UNK][CLS]的概率分布
|
|||
|
|
output = model(np.array([token_ids, ], dtype=np.int32))
|
|||
|
|
_probas = output.numpy()[0, -1, 3:]
|
|||
|
|
del output
|
|||
|
|
# 按照出现概率，对所有token倒序排列
|
|||
|
|
p_args = _probas.argsort()[::-1][:100]
|
|||
|
|
# 排列后的概率顺序
|
|||
|
|
p = _probas[p_args]
|
|||
|
|
# 先对概率归一
|
|||
|
|
p = p / sum(p)
|
|||
|
|
# 再按照预测出的概率，随机选择一个词作为预测结果
|
|||
|
|
target_index = np.random.choice(len(p), p=p)
|
|||
|
|
target = p_args[target_index] + 3
|
|||
|
|
# 保存
|
|||
|
|
token_ids.append(target)
|
|||
|
|
# 只有不是特殊字符时，才保存到poetry里面去
|
|||
|
|
if target > 3:
|
|||
|
|
poetry.append(tokenizer.id_to_token(target))
|
|||
|
|
if target in punctuation_ids:
|
|||
|
|
break
|
|||
|
|
return ''.join(poetry)
|