PyTorch示例——使用Transformer写古诗

蒋含竹

已于 2024-04-22 16:10:33 修改

阅读量2.9k

点赞数 23

分类专栏： MachineLearning # PyTorch 文章标签： pytorch transformer 人工智能古诗深度学习 python

于 2024-04-03 21:06:29 首次发布

本文链接：https://blog.csdn.net/alionsss/article/details/137357691

版权

文章目录

PyTorch示例——使用Transformer写古诗

PyTorch示例——使用Transformer写古诗

1. 前言

很早、很早以前，在TensorFlow2 学习——RNN生成古诗词_rnn古诗生成头词汇是 “ 日、红、山、夜、湖、海、月。-CSDN博客中已使用TensorFlow+RNN的方式实现过写古诗的功能，现在来个Pytorch+Transformer的示例😄
数据处理逻辑和前面博文中大致相似，本文中就不再赘述
Kaggle Notebook地址: PyTorch示例-使用Transformer写古诗x

2. 版本信息

PyTorch: 2.1.2
Python: 3.10.13

3. 导包

import math
import numpy as np
from collections import Counter
import torch
from torch import nn
from torch.utils.data import TensorDataset
from torch.utils.data import DataLoader
import tqdm
import random
import sys

print("Pytorch 版本：", torch.__version__)
print("Python  版本：", sys.version)

Pytorch 版本： 2.1.2
Python  版本： 3.10.13 | packaged by conda-forge | (main, Dec 23 2023, 15:36:39) [GCC 12.3.0]

4. 数据与预处理

数据下载

度盘： https://pan.baidu.com/s/1HIROi4mPMv0RBWHIHvUDVg，提取码：b2pp
Kaggle：https://www.kaggle.com/datasets/alionsss/poetry

先看一下原始数据

# 数据路径
DATA_PATH = '/kaggle/input/poetry/poetry.txt'

# 先看下原始数据，每一行格式为"诗的标题:诗的内容"
with open(DATA_PATH, 'r', encoding='utf-8') as f:
    lines = f.readlines()
    for i in range(0, 5):
        print(lines[i])
    print(f"origin_line_count = {
     len(lines)}")

首春:寒随穷律变，春逐鸟声开。初风飘带柳，晚雪间花梅。碧林青旧竹，绿沼翠新苔。芝田初雁去，绮树巧莺来。

初晴落景:晚霞聊自怡，初晴弥可喜。日晃百花色，风动千林翠。池鱼跃不同，园鸟声还异。寄言博通者，知予物外志。

初夏:一朝春夏改，隔夜鸟花迁。阴阳深浅叶，晓夕重轻烟。哢莺犹响殿，横丝正网天。珮高兰影接，绶细草纹连。碧鳞惊棹侧，玄燕舞檐前。何必汾阳处，始复有山泉。

度秋:夏律昨留灰，秋箭今移晷。峨嵋岫初出，洞庭波渐起。桂白发幽岩，菊黄开灞涘。运流方可叹，含毫属微理。

仪鸾殿早秋:寒惊蓟门叶，秋发小山枝。松阴背日转，竹影避风移。提壶菊花岸，高兴芙蓉池。欲知凉气早，巢空燕不窥。

origin_line_count = 43030

开始处理数据，过滤掉异常数据

# 单行诗最大长度
MAX_LEN = 64
MIN_LEN = 5
# 禁用的字符，拥有以下符号的诗将被忽略
DISALLOWED_WORDS = ['（', '）', '(', ')', '__', '《', '》', '【', '】', '[', ']', '？', '；']

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

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

寒随穷律变，春逐鸟声开。初风飘带柳，晚雪间花梅。碧林青旧竹，绿沼翠新苔。芝田初雁去，绮树巧莺来。
晚霞聊自怡，初晴弥可喜。日晃百花色，风动千林翠。池鱼跃不同，园鸟声还异。寄言博通者，知予物外志。
夏律昨留灰，秋箭今移晷。峨嵋岫初出，洞庭波渐起。桂白发幽岩，菊黄开灞涘。运流方可叹，含毫属微理。
寒惊蓟门叶，秋发小山枝。松阴背日转，竹影避风移。提壶菊花岸，高兴芙蓉池。欲知凉气早，巢空燕不窥。
山亭秋色满，岩牖凉风度。疏兰尚染烟，残菊犹承露。古石衣新苔，新巢封古树。历览情无极，咫尺轮光暮。
current_line_count = 24375

过滤掉出现频率较低的字符串，后面统一当作 UNKNOWN

# 最小词频
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]

# 打印一下出现次数前5的字
for i, (token, count) in enumerate(counter.items()):
    print(token, "->",count)
    if i >= 4:
        break;

寒 -> 2612
随 -> 1036
穷 -> 482
律 -> 118
变 -> 286

定义词典编码器 Tokenizer

class Tokenizer:
    """
    词典编码器
    """
    UNKNOWN = "<unknown>"
    PAD = "<pad>"
    BOS = "<bos>" 
    EOS = "<eos>" 

    def __init__(self, tokens):
        # 补上特殊词标记：未知词标记、填充字符标记、开始标记、结束标记
        tokens = [Tokenizer.UNKNOWN, Tokenizer.PAD, Tokenizer.BOS, Tokenizer.EOS] + tokens
        # 词汇表大小
        self.dict_size = len