import torch
import math
import time
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from torch.utils.data import Dataset, DataLoader
import sys
#1.构建数据集
datas = np.load('tang.npz')
data, word2ix, ix2word = datas['data'], datas['word2ix'].item(), datas['ix2word'].item()
#2.搭建神经网络
class PoetryModel(nn.Module):#vocab_size=input_size=num_class
def __init__(self, vocab_size, embedding_size, hidden_size):
super(PoetryModel, self).__init__()
self.hidden_size = hidden_size
self.embeddings = nn.Embedding(vocab_size, embedding_size) #(seq_len,batch_size)->(seq_len,batch_size,embedding_size)
self.lstm = nn.LSTM(embedding_size, self.hidden_size, num_layers=2) #(inputsize,hiddensize,num_layzers) 这里inputsize=embeddingsize
self.linear = nn.Linear(self.hidden_size, vocab_size) #input(seqlen,batchsize,inputsize),h0(num_layers,batchsize,hiddensize),c0
#out(seqlen,batchsize,hiddensize),hn(num_layers,batchsize,hiddensize),cn
def forward(self, input, hidden=None): #data.size([57580, 125]) 要做一个转置
# input=input.t() #之后的设置中有转置项
seq_len, batch_size = input.size() #(seq_len,batch_size)
if hidden is None:
# h_0 = 0.01*torch.Tensor(2, batch_size, self.hidden_size).normal_().cuda()
# c_0 = 0.01*torch.Tensor(2, batch_size, self.hidden_size).normal_().cuda()
h_0 = input.data.new(2, batch_size, self.hidden_size).fill_(0).float()
c_0 = input.data.new(2, batch_size, self.hidden_size).fill_(0).float()
else:
h_0, c_0 = hidden
# size: (seq_len,batch_size,embedding_size) #经过embedding层输出的维度
embeds = self.embeddings(input)
# output size: (seq_len,batch_size,hidden_size)
output, hidden = self.lstm(embeds, (h_0, c_0))
# size: (seq_len*batch_size,vocab_size) #(seq_len*batch_size,hidden_size)->(seq_len*batch_size,vocab_size)
output = self.linear(output.view(seq_len * batch_size, -1))
return output, hidden
#(inputsize,embeddingsize,hiddensize) 输出output(seq_len*batch_size,vocab_size) hidden(hn,cn)
model = PoetryModel(len(word2ix), 128, 256)
device=torch.device("cuda:0"if torch.cuda.is_available()else"cpu") #cuda:1 第二块显卡
model.to(device) #将模型迁移到GPU
model=torch.load('model.pt')
#5.AI写诗:
"""
给定几个词,根据这几个词接着生成一首完整的诗歌
start_words:u'春江潮水连海平'
比如start_words 为 春江潮水连海平,可以生成:
"""
start_words = '闲云潭影日悠悠' # 诗歌开始
prefix_words = '细雨鱼儿出,微风燕子斜。' # 不是诗歌的组成部分,用来控制生成诗歌的意境
max_gen_len = 200 # 生成诗歌最长长度
acrostic = False # 是否是藏头诗
def generate(model, start_words, ix2word, word2ix, prefix_words=None):
"""
给定几个词,根据这几个词接着生成一首完整的诗歌
start_words:u'春江潮水连海平'
比如start_words 为 春江潮水连海平,可以生成:
"""
results = list(start_words)
start_word_len = len(start_words)
# 手动设置第一个词为<START>
input = torch.Tensor([word2ix['<START>']]).view(1, 1).long()
input = input.cuda()
hidden = None
if prefix_words:
for word in prefix_words: #input output
output, hidden = model(input, hidden) #(seqlen,batch_size)->(seq_len*batch_size,vocab_size)
input = input.data.new([word2ix[word]]).view(1, 1) #这个循环影响了hidden,进而影响之后的output
for i in range(max_gen_len):
output, hidden = model(input, hidden)
if i < start_word_len:
w = results[i]
input = input.data.new([word2ix[w]]).view(1, 1)
else: #topk(dim=?) 默认对最后一个维度求最大值 ,1指对前一个数据求最大值
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] #del函数,删除results列表中的最后一个元素
break
return results
def gen_acrostic(model, start_words, ix2word, word2ix, prefix_words=None):
"""
生成藏头诗
start_words : u'深度学习'
生成:
深木通中岳,青苔半日脂。
度山分地险,逆浪到南巴。
学道兵犹毒,当时燕不移。
习根通古岸,开镜出清羸。
"""
results = []
start_word_len = len(start_words)
input = (torch.Tensor([word2ix['<START>']]).view(1, 1).long())
input = input.cuda()
hidden = None
index = 0 # 用来指示已经生成了多少句藏头诗
# 上一个词
pre_word = '<START>'
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(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 {u'。', u'!', '<START>'}):
# 如果遇到句号,藏头的词送进去生成
if index == start_word_len:
# 如果生成的诗歌已经包含全部藏头的词,则结束
break
else:
# 把藏头的词作为输入送入模型
w = start_words[index]
index += 1
input = (input.data.new([word2ix[w]])).view(1, 1)
else:
# 否则的话,把上一次预测是词作为下一个词输入
input = (input.data.new([word2ix[w]])).view(1, 1)
results.append(w)
pre_word = w
return results
def gen(**kwargs):
"""
提供命令行接口,用以生成相应的诗
"""
start_words = '闲云潭影日悠悠' # 诗歌开始
prefix_words = '细雨鱼儿出,微风燕子斜。' # 不是诗歌的组成部分,用来控制生成诗歌的意境
max_gen_len = 200 # 生成诗歌最长长度
acrostic = False # 是否是藏头诗
model = torch.load('model.pt')
model.cuda()
# python2和python3 字符串兼容
if sys.version_info.major == 3:
if start_words.isprintable():
start_words = start_words
prefix_words = prefix_words if prefix_words else None
else:
start_words = start_words.encode('ascii', 'surrogateescape').decode('utf8')
prefix_words = prefix_words.encode('ascii', 'surrogateescape').decode(
'utf8') if prefix_words else None
else:
start_words = start_words.decode('utf8')
prefix_words = prefix_words.decode('utf8') if prefix_words else None
start_words = start_words.replace(',', u',') \
.replace('.', u'。') \
.replace('?', u'?')
gen_poetry = gen_acrostic if acrostic else generate
result = gen_poetry(model, start_words, ix2word, word2ix, prefix_words)
print(''.join(result))
gen_poetries = []
# 分别以这几个字作为诗歌的第一个字,生成8首诗
for word in list(u'春江花月夜凉如水'):
gen_poetry = ''.join(generate(model, word, ix2word, word2ix))
gen_poetries.append(gen_poetry)
print(gen_poetries)
#以这几个字为诗歌的开头生成一首诗
word=u'春江'
print(''.join(generate(model, word, ix2word, word2ix)))
#书写藏头诗
start_words=u'电信十班'
print(''.join(gen_acrostic(model, start_words, ix2word, word2ix, prefix_words=None)))
循环神经网络:用训练好的model写诗歌
最新推荐文章于 2023-06-12 22:45:57 发布