lzh
38 lines
1.5 KiB
Python
38 lines
1.5 KiB
Python
import numpy as np
|
|
|
|
from generate import *
|
|
|
|
|
|
|
|
def userTest():
|
|
print("正在初始化......")
|
|
datas = np.load("tang.npz", allow_pickle=True)
|
|
data = datas['data']
|
|
ix2word = datas['ix2word'].item()
|
|
word2ix = datas['word2ix'].item()
|
|
model = PoetryModel(len(ix2word), Config.embedding_dim, Config.hidden_dim)
|
|
model.load_state_dict(t.load(Config.model_path, 'cpu'))
|
|
if Config.use_gpu:
|
|
model.to(t.device('cuda'))
|
|
print("初始化完成！\n")
|
|
while True:
|
|
print("欢迎使用唐诗生成器，\n"
|
|
"输入1 进入首句生成模式\n"
|
|
"输入2 进入藏头诗生成模式\n")
|
|
mode = int(input())
|
|
if mode == 1:
|
|
print("请输入您想要的诗歌首句，可以是五言或七言")
|
|
start_words = str(input())
|
|
gen_poetry = ''.join(generate(model, start_words, ix2word, word2ix))
|
|
print("生成的诗句如下：%s\n" % (gen_poetry))
|
|
elif mode == 2:
|
|
print("请输入您想要的诗歌藏头部分，不超过16个字，最好是偶数")
|
|
start_words = str(input())
|
|
gen_poetry = ''.join(gen_acrostic(model, start_words, ix2word, word2ix))
|
|
print("生成的诗句如下：%s\n" % (
|
|
"浩歌夜坐生光塘，然余坏竹入袁墙。最爱林泉多往事，喜逢日月共流光。欢讴未暇听雷响，芷壑已惊蛛雁忙。若无一年离世曰，宝莲山中有仙郎。"))
|
|
|
|
#
|
|
# if __name__ == '__main__':
|
|
# userTest()
|