TensorFlow2学习：RNN生成古诗词

来源 | CSDN博客作者 | 蒋含竹责编 | 徐威龙
利用循环神经网络RNN可以做各种连续性数据的预测，其中生成古诗词是一件非常有趣的事，特此分享我的学习经验。

宁静致远
宁随古峰一里乡，静在门林满树通。致有旧人身自住，远花不似水花中。

风起云涌
风山一夕月，起落鸟纷纷。云散生何处，涌深千尺村。

春夏秋冬
春来空树柳微时，夏火遥愁独寂寥。秋上北陵村未苦，冬来寒向入楼僧。

import math
import re
import numpy as np
import tensorflow as tf
from collections import Counter

过老子庙:仙居怀圣德，灵庙肃神心。草合人踪断，尘浓鸟迹深。流沙丹灶没，关路紫烟沉。独伤千载后，空馀松柏林。
途次陕州:境出三秦外，途分二陕中。山川入虞虢，风俗限西东。树古棠阴在，耕余让畔空。鸣笳从此去，行见洛阳宫。
野次喜雪:拂曙辟行宫，寒皋野望通。每云低远岫，飞雪舞长空。赋象恒依物，萦回屡逐风。为知勤恤意，先此示年丰。
送贺知章归四明:遗荣期入道，辞老竟抽簪。岂不惜贤达，其如高尚心。寰中得秘要，方外散幽襟。独有青门饯，群僚怅别深。
轩游宫十五夜:行迈离秦国，巡方赴洛师。路逢三五夜，春色暗中期。关外长河转，宫中淑气迟。歌钟对明月，不减旧游时。

# 数据路径
DATA_PATH = './datasets/poetry.txt'
# 单行诗最大长度
MAX_LEN = 64
# 禁用的字符，拥有以下符号的诗将被忽略
DISALLOWED_WORDS = ['（', '）', '(', ')', '__', '《', '》', '【', '】', '[', ']']

# 一首诗（一行）对应一个列表的元素
poetry = []

# 按行读取数据 poetry.txt
with open(DATA_PATH, 'r', encoding='utf-8') as f:
    lines = f.readlines()
# 遍历处理每一条数据    
for line in lines:
    # 利用正则表达式拆分标题和内容
    fields = re.split(r"[:：]", line)
    # 跳过异常数据
    if len(fields) != 2:
        continue
    # 得到诗词内容（后面不需要标题）
    content = fields[1]
    # 跳过内容过长的诗词
    if len(content) > MAX_LEN - 2:
        continue
    # 跳过存在禁用符的诗词
    if any(word in content for word in DISALLOWED_WORDS):
        continue

    poetry.append(content.replace('\n', '')) # 最后要记得删除换行符

for i in range(0, 5):
    print(poetry[i])

系马宫槐老，持杯店菊黄。故交今不见，流恨满川光。
世间何事不潸然，得失人情命不延。适向蔡家厅上饮，回头已见一千年。
只领千馀骑，长驱碛邑间。云州多警急，雪夜度关山。石响铃声远，天寒弓力悭。秦楼休怅望，不日凯歌还。
今日花前饮，甘心醉数杯。但愁花有语，不为老人开。
秋来吟更苦，半咽半随风。禅客心应乱，愁人耳愿聋。雨晴烟树里，日晚古城中。远思应难尽，谁当与我同。

# 最小词频
MIN_WORD_FREQUENCY = 8

# 统计词频，利用Counter可以直接按单个字符进行统计词频
counter = Counter()
for line in poetry:
    counter.update(line)
# 过滤掉低词频的词
tokens = [token for token, count in counter.items() if count >= MIN_WORD_FREQUENCY]

i = 0
for token, count in counter.items():
    if i >= 5:
        break;
    print(token, "->",count)
    i += 1

寒 -> 2627
随 -> 1039
穷 -> 487
律 -> 119
变 -> 286

# 补上特殊词标记：填充字符标记、未知词标记、开始标记、结束标记
tokens = ["[PAD]", "[NONE]", "[START]", "[END]"] + tokens

# 映射: 词 -> 编号
word_idx = {}
# 映射: 编号 -> 词
idx_word = {}
for idx, word in enumerate(tokens):
    word_idx[word] = idx
    idx_word[idx] = word

class Tokenizer:
    """
    分词器
    """

    def __init__(self, tokens):
        # 词汇表大小
        self.dict_size = len(tokens)
        # 生成映射关系
        self.token_id = {} # 映射: 词 -> 编号
        self.id_token = {} # 映射: 编号 -> 词
        for idx, word in enumerate(tokens):
            self.token_id[word] = idx
            self.id_token[idx] = word

        # 各个特殊标记的编号id，方便其他地方使用
        self.start_id = self.token_id["[START]"]
        self.end_id = self.token_id["[END]"]
        self.none_id = self.token_id["[NONE]"]
        self.pad_id = self.token_id["[PAD]"]

    def id_to_token(self, token_id):
        """
        编号 -> 词
        """
        return self.id_token.get(token_id)

    def token_to_id(self, token):
        """
        词 -> 编号
        """
        return self.token_id.get(token, self.none_id)

    def encode(self, tokens):
        """
        词列表 -> [START]编号 + 编号列表 + [END]编号
        """
        token_ids = [self.start_id, ] # 起始标记
        # 遍历，词转编号
        for token in tokens:
            token_ids.append(self.token_to_id(token))
        token_ids.append(self.end_id) # 结束标记
        return token_ids

    def decode(self, token_ids):
        """
        编号列表 -> 词列表(去掉起始、结束标记)
        """
        # 起始、结束标记
        flag_tokens = {"[START]", "[END]"}

        tokens = []
        for idx in token_ids:
            token = self.id_to_token(idx)
            # 跳过起始、结束标记
            if token not in flag_tokens:
                tokens.append(token)
        return tokens

tokenizer = Tokenizer(tokens)

class PoetryDataSet:
    """
    古诗数据集生成器
    """

    def __init__(self, data, tokenizer, batch_size):
        # 数据集
        self.data = data
        self.total_size = len(self.data)
        # 分词器，用于词转编号
        self.tokenizer = tokenizer
        # 每批数据量
        self.batch_size = BATCH_SIZE
        # 每个epoch迭代的步数
        self.steps = int(math.floor(len(self.data) / self.batch_size))

    def pad_line(self, line, length, padding=None):
        """
        对齐单行数据
        """
        if padding is None:
            padding = self.tokenizer.pad_id

        padding_length = length - len(line)
        if padding_length > 0:
            return line + [padding] * padding_length
        else:
            return line[:length]

    def __len__(self):
        return self.steps

    def __iter__(self):
        # 打乱数据
        np.random.shuffle(self.data)
        # 迭代一个epoch，每次yield一个batch
        for start in range(0, self.total_size, self.batch_size):
            end = min(start + self.batch_size, self.total_size)
            data = self.data[start:end]

            max_length = max(map(len, data)) 

            batch_data = []
            for str_line in data:
                # 对每一行诗词进行编码、并补齐padding
                encode_line = self.tokenizer.encode(str_line)
                pad_encode_line = self.pad_line(encode_line, max_length + 2) # 加2是因为tokenizer.encode会添加START和END
                batch_data.append(pad_encode_line)

            batch_data = np.array(batch_data)
            # yield 特征、标签
            yield batch_data[:, :-1], batch_data[:, 1:]

    def generator(self):
        while True:
            yield from self.__iter__()

特征：[START]我有辞乡剑，玉锋堪截云。襄阳走马客，意气自生春。朝嫌剑花净，暮嫌剑光冷。能持剑向人，不解持照身。[END][PAD][PAD][PAD]
标签：我有辞乡剑，玉锋堪截云。襄阳走马客，意气自生春。朝嫌剑花净，暮嫌剑光冷。能持剑向人，不解持照身。[END][PAD][PAD][PAD][PAD]

dataset = PoetryDataSet(poetry, tokenizer, BATCH_SIZE)

model = tf.keras.Sequential([
    # 词嵌入层
    tf.keras.layers.Embedding(input_dim=tokenizer.dict_size, output_dim=150),
    # 第一个LSTM层
    tf.keras.layers.LSTM(150, dropout=0.5, return_sequences=True),
    # 第二个LSTM层
    tf.keras.layers.LSTM(150, dropout=0.5, return_sequences=True),
    # 利用TimeDistributed对每个时间步的输出都做Dense操作(softmax激活)
    tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(tokenizer.dict_size, activation='softmax')),
])

model.summary()

Model: "sequential_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
embedding_2 (Embedding)      (None, None, 150)         515100    
_________________________________________________________________
lstm_4 (LSTM)                (None, None, 150)         180600    
_________________________________________________________________
lstm_5 (LSTM)                (None, None, 150)         180600    
_________________________________________________________________
time_distributed_2 (TimeDist (None, None, 3434)        518534    
=================================================================
Total params: 1,394,834
Trainable params: 1,394,834
Non-trainable params: 0
_________________________________________________________________

model.compile(
    optimizer=tf.keras.optimizers.Adam(), 
    loss=tf.keras.losses.sparse_categorical_crossentropy
)

model.fit(
    dataset.generator(), 
    steps_per_epoch=dataset.steps, 
    epochs=10
)

Train for 767 steps
Epoch 1/10
767/767 [==============================] - 34s 44ms/step - loss: 4.8892
Epoch 2/10
767/767 [==============================] - 31s 41ms/step - loss: 4.2494
Epoch 3/10
767/767 [==============================] - 31s 40ms/step - loss: 4.1113
Epoch 4/10
767/767 [==============================] - 31s 40ms/step - loss: 3.9864
Epoch 5/10
767/767 [==============================] - 31s 40ms/step - loss: 3.8660
Epoch 6/10
767/767 [==============================] - 31s 40ms/step - loss: 3.7879
Epoch 7/10
767/767 [==============================] - 31s 40ms/step - loss: 3.7339
Epoch 8/10
767/767 [==============================] - 31s 40ms/step - loss: 3.6826
Epoch 9/10
767/767 [==============================] - 31s 40ms/step - loss: 3.6275
Epoch 10/10
767/767 [==============================] - 31s 40ms/step - loss: 3.5999

# 需要先将词转为编号
token_ids = [tokenizer.token_to_id(word) for word in ["月", "光", "静", "谧"]]
# 进行预测
result = model.predict([token_ids ,])
print(result)

[[[2.0809230e-04 9.3881181e-03 5.5695949e-07 ... 5.6030808e-06
   8.5241054e-06 2.0507096e-06]
  [7.6916285e-06 6.1246334e-03 1.8850582e-08 ... 4.8418292e-06
   2.8483141e-06 5.3288642e-07]
  [5.0856406e-06 3.1365673e-03 1.9067786e-08 ... 4.5156207e-06
   1.0479171e-05 9.7814757e-07]
  [7.1793047e-06 2.2729969e-02 2.0391434e-08 ... 2.0609916e-06
   2.2420336e-06 2.1413473e-06]]]

def predict(model, token_ids):
    """
    在概率值为前100的词中选取一个词(按概率分布的方式)
    :return: 一个词的编号(不包含[PAD][NONE][START])
    """
    # 预测各个词的概率分布
    # -1 表示只要对最新的词的预测
    # 3: 表示不要前面几个标记符
    _probas = model.predict([token_ids, ])[0, -1, 3:]
    # 按概率降序，取前100
    p_args = _probas.argsort()[-100:][::-1] # 此时拿到的是索引
    p = _probas[p_args] # 根据索引找到具体的概率值
    p = p / sum(p) # 归一
    # 按概率抽取一个
    target_index = np.random.choice(len(p), p=p)
    # 前面预测时删除了前几个标记符，因此编号要补上3位，才是实际在tokenizer词典中的编号
    return p_args[target_index] + 3

token_ids = tokenizer.encode("清风明月")[:-1]
while len(token_ids) < 13:
    # 预测词的编号
    target = predict(model, token_ids)
    # 保存结果
    token_ids.append(target)
    # 到达END
    if target == tokenizer.end_id: 
        break

print("".join(tokenizer.decode(token_ids)))

清风明月夜，晚色北堂残。

def generate_random_poem(tokenizer, model, text=""):
    """
    随机生成一首诗
    :param tokenizer: 分词器
    :param model: 古诗模型
    :param text: 古诗的起始字符串，默认为空
    :return: 一首古诗的字符串
    """
    # 将初始字符串转成token_ids，并去掉结束标记[END]
    token_ids = tokenizer.encode(text)[:-1]
    while len(token_ids) < MAX_LEN:
        # 预测词的编号
        target = predict(model, token_ids)
        # 保存结果
        token_ids.append(target)
        # 到达END
        if target == tokenizer.end_id: 
            break

    return "".join(tokenizer.decode(token_ids))

for i in range(5):
    print(generate_random_poem(tokenizer, model))

江亭路断暮，归去见芳洲。惆怅门中去，心年少地深。夜期深木静，水落夕阳深。秋去人南雨，凄头望海中。
洛陌江阳宫下树，玉门宫夜似东云。今更已长逢醉士，一明先语似相春。
春山风半夜初归，万岁空声去去过。自惜秦生犹送酒，何人无计不安稀。
何处东陵路，无年已复还。晓莺逢半急，潮望月云稀。暗影通三度，烟沙水鸟深。当年相忆望，何处问渔家。
清夜向阳阁，一风看北宫。雨分红蕊草，红杏药茶行。野石翻山远，猿晴不独天。谁知一山下，飞首却悠悠。

print(generate_random_poem(tokenizer, model, "春眠不觉晓，"))
print(generate_random_poem(tokenizer, model, "白日依山尽，"))
print(generate_random_poem(tokenizer, model, "秦时明月汉时关，"))

春眠不觉晓，坐住树深空。风月飘犹晓，春多出水流。
白日依山尽，相逢独水声。唯疑见心意，一老泪鸣归。落晚南游客，吟猿见柳寒。何堪看暮望，还见有军情。
秦时明月汉时关，欲望时恩不道心。莫忆旧乡僧雁在，始堪曾在牡苓流。

def generate_acrostic_poem(tokenizer, model, heads):
    """
    生成一首藏头诗
    :param tokenizer: 分词器
    :param model: 古诗模型
    :param heads: 藏头诗的头
    :return: 一首古诗的字符串
    """
    # token_ids，只包含[START]编号
    token_ids = [tokenizer.start_id, ]
    # 逗号和句号标记编号
    punctuation_ids = {tokenizer.token_to_id("，"), tokenizer.token_to_id("。")}
    content = []
    # 为每一个head生成一句诗
    for head in heads:
        content.append(head)
        # head转为编号id，放入列表，用于预测
        token_ids.append(tokenizer.token_to_id(head))
        # 开始生成一句诗
        target = -1;
        while target not in punctuation_ids: # 遇到逗号、句号，说明本句结束，开始下一句
            # 预测词的编号
            target = predict(model, token_ids)
            # 因为可能预测到END，所以加个判断
            if target > 3:
                # 保存结果到token_ids中，下一次预测还要用
                token_ids.append(target)
                content.append(tokenizer.id_to_token(target))

    return "".join(content)

print(generate_acrostic_poem(tokenizer, model, heads="上善若水"))
print(generate_acrostic_poem(tokenizer, model, heads="明月清风"))
print(generate_acrostic_poem(tokenizer, model, heads="点个赞吧"))

上亭清色望，善地半烟霞。若辨从秋日，水花清上清。
明夕远多尽，月生开雨明。清山看楚雪，风色水堂钟。
点阁风空雪，个枝时未开。赞君初合泪，吧石似春风。

class ShowSaveCallback(tf.keras.callbacks.Callback):

    def __init__(self):
        super().__init__()
        # 给一个初始最大值
        self.loss = float("inf")

    def on_epoch_end(self, epoch, logs=None):
        # 保留损失最低的模型
        if logs['loss'] <= self.loss:
            self.loss = logs['loss']
            model.save("./rnn_model.h5")
        # 查看一下本次训练的效果
        print()
        for i in range(5):
            print(generate_random_poem(tokenizer, model))

# 开始训练
model.fit(
    dataset.generator(), 
    steps_per_epoch=dataset.steps, 
    epochs=10,
    callbacks=[ShowSaveCallback()]
)

model = tf.keras.models.load_model("./rnn_model.h5")

# 后面就可以继续进行预测了

【end】

◆有奖征文◆
推荐阅读
2020年，5种将死的编程语言检测、量化、追踪新冠病毒，基于深度学习的自动CT图像分析有多靠谱？GitHub 接连封杀开源项目惹众怒，CEO 亲自道歉！智能合约编写之 Solidity 的设计模式低学历、文科出身，我如何从月薪不到 3000 逆袭为大厂高薪程序员？从提取层、处理层、基础结构入手，带你了解Spark和Kafka！你点的每个“在看”，我都认真当成了AI