循环神经网络：用训练好的model写诗歌

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)))