【AI达人创造营第二期】基于LSTM的现代诗生成器

# 导入相关依赖from paddle.io import Datasetimport paddle.fluid as fluidimport numpy as npimport paddleimport paddle.nnfrom paddlenlp.embeddings import TokenEmbeddingfrom paddlenlp.data import JiebaTokenizer,Vocabimport visualdl

# 定义超参数class Config(object):
    # version = 'models/version1-modern/version1.pdparams' # 现代诗风格
    version = 'models/version2-ancient/version2.pdparams' # 古体诗风格
    maxl = 120
    filepath = "vocab/poems_without_title.txt"
    filepath2 = "vocab/poems_zh.txt"
    embedding_dim = 300
    hidden_dim = 512
    num_layers = 3

    max_gen_len = 150
    prefix = "爱你一生一世" # 前置风格，可以调整生成文本的风格
    beginning = "陪伴是最长情的告白" # 需要给出诗篇的开头，模型进行续写config = Config()

# 加载词汇表vocabdic = {'[PAD]':0,'<start>':1,'<end>':2,'[UNK]':3}
cnt=4with open (config.filepath) as fp:    for line in fp:        for char in line:            if char not in dic:
                dic[char] = cnt
                cnt+=1with open (config.filepath2) as fp:    for line in fp:        for char in line:            if char not in dic:
                dic[char] = cnt
                cnt+=1vocab = Vocab.from_dict(dic,unk_token='[UNK]')

# 加载模型class Poetry(paddle.nn.Layer):
    def __init__(self,vocab_size,embedding_dim,hidden_dim):
        super().__init__()
        self.embeddings = paddle.nn.Embedding(vocab_size,embedding_dim)
        self.lstm = paddle.nn.LSTM(
            input_size=embedding_dim,
            hidden_size=hidden_dim,
            num_layers=config.num_layers,
        )
        self.linear = paddle.nn.Linear(in_features=hidden_dim,out_features=vocab_size)    def forward(self,input,hidden=None):
        batch_size, seq_len = paddle.shape(input)
        embeds = self.embeddings(input)        if hidden is None:
            output,hidden = self.lstm(embeds)        else:
            output,hidden = self.lstm(embeds,hidden)
        output = paddle.reshape(output,[seq_len*batch_size,Config.hidden_dim])
        output = self.linear(output)        return output,hidden

poetry = Poetry(len(vocab),config.embedding_dim,config.hidden_dim)
poetry.set_state_dict(paddle.load(config.version))

results = [i for i in config.beginning]
start_words_len = len(results)input = (paddle.to_tensor(vocab("<start>"))).reshape([1,1])
hidden = Noneif config.prefix:
    words = [i for i in config.prefix]    for word in words:
        _, hidden = poetry(input, hidden)        input = (paddle.to_tensor(vocab(word))).reshape([1,1])for i in range(config.max_gen_len):
    output, hidden = poetry(input, hidden)    if i < start_words_len:
        word = results[i]        input = (paddle.to_tensor(vocab(word))).reshape([1,1])    else:
        _,top_index = paddle.fluid.layers.topk(output[0],k=1)
        top_index = top_index.item()
        word = vocab.to_tokens(top_index)
        results.append(word)        input = paddle.to_tensor([top_index])        input = paddle.reshape(input,[1,1])    if word == '<end>':        del results[-1]        breakresults = ''.join(results)print(results)