import torch
import torch as t
import numpy as np
import random
from torch.utils.data import DataLoader
from torch import optim
from torch import nn
from torchnet import meter
import tqdm
class Config(object):
num_layers = 3 # LSTM层数
data_path = 'data/' # 诗歌的文本文件存放路径
pickle_path = 'tang.npz' # 预处理好的二进制文件
author = None # 只学习某位作者的诗歌
constrain = None # 长度限制
category = 'poet.tang' # 类别，唐诗还是宋诗歌(poet.song)
lr = 1e-3
weight_decay = 1e-4
epoch = 5000
batch_size = 16
maxlen = 125 # 超过这个长度的之后字被丢弃，小于这个长度的在前面补空格
plot_every = 1000 # 每20个batch 可视化一次
use_gpu = True
# use_env = True # 是否使用visodm
env = 'poetry' # visdom env
max_gen_len = 200 # 生成诗歌最长长度
debug_file = '/tmp/debugp'
model_path = "./cps_new/tang_0.pth" # 预训练模型路径
prefix_words = '床前明月光，疑是地上霜。' # 不是诗歌的组成部分，用来控制生成诗歌的意境
start_words = '窗前明月光' # 诗歌开始
acrostic = False # 是否是藏头诗
model_prefix = 'cps_new/tang' # 模型保存路径
embedding_dim = 256
hidden_dim = 512
class PoetryModel(nn.Module):
def __init__(self, vocab_size, embedding_dim, hidden_dim):
super(PoetryModel, self).__init__()
self.hidden_dim = hidden_dim
# 词向量层，词表大小 * 向量维度
self.embeddings = nn.Embedding(vocab_size, embedding_dim)
# 网络主要结构
self.lstm = nn.LSTM(embedding_dim, self.hidden_dim, num_layers=Config.num_layers, dropout = 0.1)
# 进行分类
self.linear = nn.Linear(self.hidden_dim, vocab_size)
def forward(self, input, hidden=None):
seq_len, batch_size = input.size()
if hidden is None:
h_0 = input.data.new(Config.num_layers, batch_size, self.hidden_dim).fill_(0).float()
c_0 = input.data.new(Config.num_layers, batch_size, self.hidden_dim).fill_(0).float()
else:
h_0, c_0 = hidden
# 输入 序列长度 * batch(每个汉字是一个数字下标)，
# 输出 序列长度 * batch * 向量维度
embeds = self.embeddings(input)
# 输出hidden的大小： 序列长度 * batch * hidden_dim
output, hidden = self.lstm(embeds, (h_0, c_0))
output = self.linear(output.view(seq_len * batch_size, -1))
return output, hiddendef isChinese(word):
for ch in word:
if 'u4e00' <= ch <= 'u9fff':
return True
return False
def Load():
print("Begin Loading...")
datas = np.loadtxt("./data_new.txt",
dtype = str, delimiter = '
', encoding = "utf-8")
len_datas = len(datas)
_char = []
for i in range(len_datas):
# if len(datas[i]) == 24:
for j in range(len(datas[i])):
_char.append(datas[i][j])
print("Done")
_char = list(set(_char))
return datas, _char
if True:
print("正在初始化......")
datas, char = Load()
word2ix = {}
ix2word = {}
char.append('</s>')
char.append('<START>')
char.append('<EOP>')
index = 0
for item in char:
ix2word[index] = item
index += 1
for key in ix2word:
word2ix[ix2word[key]]=key
len_datas = len(datas)
cnt = 0;
for i in range(len_datas):
if len(datas[i]) == 24:
cnt += 1
data = np.empty([cnt, 26], dtype = int)
index = 0
for i in range(len_datas):
if len(datas[i]) != 24:
continue
tmp = []
tmp.append(word2ix['<START>'])
for j in range(len(datas[i])):
tmp.append(word2ix[datas[i][j]])
tmp.append(word2ix['<EOP>'])
data[index] = tmp
index += 1
data = t.from_numpy(data)
dataloader = DataLoader(data,
batch_size=Config.batch_size,
shuffle=True,
num_workers=0)
def generate(model, start_words, ix2word, word2ix, prefix_words=None):
results = list(start_words)
start_words_len = len(start_words)
# 第一个词语是<START>
input = t.Tensor([word2ix['<START>']]).view(1, 1).long()
if Config.use_gpu:
input = input.cuda()
hidden = None
# 若有风格前缀，则先用风格前缀生成hidden
if prefix_words:
# 第一个input是<START>，后面就是prefix中的汉字
# 第一个hidden是None，后面就是前面生成的hidden
for word in prefix_words:
output, hidden = model(input, hidden)
input = input.data.new([word2ix[word]]).view(1, 1)
# 开始真正生成诗句，如果没有使用风格前缀，则hidden = None，input = <START>
# 否则，input就是风格前缀的最后一个词语，hidden也是生成出来的
for i in range(24):
output, hidden = model(input, hidden)
# 如果还在诗句内部，输入就是诗句的字，不取出结果，只为了得到
# 最后的hidden
if i < start_words_len:
w = results[i]
input = input.data.new([word2ix[w]]).view(1, 1)
# 否则将output作为下一个input进行
else:
top_index = output.data[0].topk(1)[1][0].item()
w = ix2word[top_index]
results.append(w)
input = input.data.new([top_index]).view(1, 1)
# if w == '<EOP>':
# del results[-1]
# break
return results
# 生成藏头诗
def gen_acrostic(model, start_words, ix2word, word2ix, prefix_words=None):
result = []
start_words_len = len(start_words)
input = (t.Tensor([word2ix['<START>']]).view(1, 1).long())
if Config.use_gpu:
input = input.cuda()
# 指示已经生成了几句藏头诗
index = 0
pre_word = '<START>'
hidden = None
# 存在风格前缀，则生成hidden
if prefix_words:
for word in prefix_words:
output, hidden = model(input, hidden)
input = (input.data.new([word2ix[word]])).view(1, 1)
# 开始生成诗句
for i in range(Config.max_gen_len):
output, hidden = model(input, hidden)
top_index = output.data[0].topk(1)[1][0].item()
w = ix2word[top_index]
# 说明上个字是句末
if pre_word in {'。', '，', '?', '！', '<START>'}:
if index == start_words_len:
break
else:
w = start_words[index]
index += 1
input = (input.data.new([word2ix[w]])).view(1, 1)
else:
input = (input.data.new([top_index])).view(1, 1)
result.append(w)
pre_word = w
return result
# 根据 topic 生成
def gen_topic(model, start_words, ix2word, word2ix, prefix_words=None):
results = list(start_words)
start_words_len = len(start_words)
# 第一个词语是<START>
input = t.Tensor([word2ix['<START>']]).view(1, 1).long()
if Config.use_gpu:
input = input.cuda()
hidden = None
# 若有风格前缀，则先用风格前缀生成hidden
if prefix_words:
# 第一个input是<START>，后面就是prefix中的汉字
# 第一个hidden是None，后面就是前面生成的hidden
prefix_words_len = len(prefix_words)
for word in prefix_words:
output, hidden = model(input, hidden)
input = input.data.new([word2ix[word]]).view(1, 1)
for i in range(24 - prefix_words_len):
output, hidden = model(input, hidden)
# 如果还在诗句内部，输入就是诗句的字，不取出结果，只为了得到
# 最后的hidden
if i < start_words_len:
w = results[i]
input = input.data.new([word2ix[w]]).view(1, 1)
# 否则将output作为下一个input进行
else:
top_index = output.data[0].topk(1)[1][0].item()
w = ix2word[top_index]
input = input.data.new([top_index]).view(1, 1)
if w == '<EOP>':
del results[-1]
break
# 开始真正生成诗句，如果没有使用风格前缀，则 hidden = None，input = <START>
# 否则，input就是风格前缀的最后一个词语，hidden也是生成出来的
output, hidden = model(input, hidden)
top_index = output.data[0].topk(1)[1][0].item()
w = ix2word[top_index]
input = input.data.new([word2ix[w]]).view(1, 1)
for i in range(24):
output, hidden = model(input, hidden)
# 如果还在诗句内部，输入就是诗句的字，不取出结果，只为了得到
# 最后的hidden
if i < start_words_len:
w = results[i]
input = input.data.new([word2ix[w]]).view(1, 1)
# 否则将output作为下一个input进行
else:
top_index = output.data[0].topk(1)[1][0].item()
w = ix2word[top_index]
results.append(w)
input = input.data.new([top_index]).view(1, 1)
if w == '<EOP>':
del results[-1]
break
return results
def train():
if Config.use_gpu:
Config.device = t.device("cuda")
else:
Config.device = t.device("cpu")
device = Config.device
# 定义模型
model = PoetryModel(len(word2ix),
embedding_dim=Config.embedding_dim,
hidden_dim = Config.hidden_dim)
Configimizer = optim.Adam(model.parameters(),lr=Config.lr)
criterion = nn.CrossEntropyLoss()
model.load_state_dict(t.load(Config.model_path))
# 转移到相应计算设备上
model.to(device)
loss_meter = meter.AverageValueMeter()
# 进行训练
f = open('result.txt','w')
for epoch in range(Config.epoch):
loss_meter.reset()
for li,data_ in tqdm.tqdm(enumerate(dataloader)):
data_ = data_.long().transpose(1,0).contiguous()
# 注意这里，也转移到了计算设备上
data_ = data_.to(device)
Configimizer.zero_grad()
# n个句子，前n-1句作为输入，后n-1句作为输出，二者一一对应
input_,target = data_[:-1,:],data_[1:,:]
output,_ = model(input_)
# 这里为什么view(-1)
# if (1+li)%2436==0:
# print(target.shape,target.view(-1).shape)
loss = criterion(output,target.view(-1))
loss.backward()
Configimizer.step()
loss_meter.add(loss.item())
# 进行可视化
if (1+li)%1000 == 0:
print("训练损失为%s"%(str(loss_meter.mean)))
word = "明月"
gen_poetry = ''.join(gen_topic(model,'',ix2word,word2ix,word))
print(gen_poetry)
t.save(model.state_dict(),'%s_%s.pth'%(Config.model_prefix,epoch % 2))
if __name__ == '__main__':
train()
def userTest():
while True:
print("欢迎使用唐诗生成器，
"
"输入1 进入首句生成模式
"
"输入2 进入藏头诗生成模式
"
"输入3 进入主题诗生成模式
")
mode = int(input())
model = PoetryModel(len(ix2word), Config.embedding_dim, Config.hidden_dim)
model.load_state_dict(t.load("./cps_new/tang_0.pth"))
if Config.use_gpu:
model.to(t.device('cuda'))
if mode == 1:
print("请输入您想要的诗歌首句，必须是五言")
# start_words = str(input())
start_words = ''
gen_poetry = ''.join(generate(model, start_words, ix2word, word2ix))
print("生成的诗句如下：%s
" % (gen_poetry))
elif mode == 2:
print("请输入您想要的诗歌藏头部分，必须是 4 个字")
start_words = str(input())
gen_poetry = ''.join(gen_acrostic(model, start_words, ix2word, word2ix))
print("生成的诗句如下：%s
" % gen_poetry)
elif mode == 3:
print("请输入您想要的诗歌topic，必须是 2 个字")
# prefix_words = str(input())
gen_poetry = ''.join(gen_topic(model, Config.start_words, ix2word, word2ix, Config.prefix_words))
print("生成的诗句如下：%s
" % gen_poetry)
if __name__ == '__main__':
userTest()