1
0

refactor: merge multiple project into one and create new project

This commit is contained in:
2026-04-07 08:30:41 +08:00
parent 7aa7ae3335
commit 6cb1a89751
49 changed files with 2932 additions and 4 deletions

3
dl-exp/exp3/datasets/.gitignore vendored Normal file
View File

@@ -0,0 +1,3 @@
# Ignore datasets and processed datasets
*.txt
*.pickle

2
dl-exp/exp3/models/.gitignore vendored Normal file
View File

@@ -0,0 +1,2 @@
# Ignore every saved model files
*.pth

View File

@@ -0,0 +1,271 @@
from pathlib import Path
import typing
import pickle
from collections import Counter
import numpy
import torch
from torch.utils.data import DataLoader, Dataset
import torch.nn.functional as F
import settings
TOKEN_PAD: str = '[PAD]'
"""使用古诗词数据时的特殊字符RNN填充时使用的填充字符。"""
TOKEN_UNK: str = '[UNK]'
"""使用古诗词数据时的特殊字符,词频不足的生僻字。"""
TOKEN_CLS: str = '[CLS]'
"""使用古诗词数据时的特殊字符,标记古诗词开始。"""
TOKEN_SEP: str = '[SEP]'
"""使用古诗词数据时的特殊字符,标记古诗词结束。"""
class Tokenizer:
"""分词器"""
token_dict: dict[str, int]
"""词->编号的映射"""
token_dict_rev: dict[int, str]
"""编号->词的映射"""
vocab_size: int
"""词汇表大小"""
def __init__(self, token_dict: dict[str, int]):
self.token_dict = token_dict
self.token_dict_rev = {value: key for key, value in self.token_dict.items()}
self.vocab_size = len(self.token_dict)
def id_to_token(self, token_id: int) -> str:
"""
给定一个编号,查找词汇表中对应的词。
:param token_id: 带查找词的编号
:return: 编号对应的词
"""
return self.token_dict_rev[token_id]
def token_to_id(self, token: str):
"""
给定一个词,查找它在词汇表中的编号。
未找到则返回低频词[UNK]的编号。
:param token: 带查找编号的词
:return: 词的编号
"""
return self.token_dict.get(token, self.token_dict['[UNK]'])
def encode(self, tokens: str) -> list[int]:
"""
给定一个字符串s在头尾分别加上标记开始和结束的特殊字符并将它转成对应的编号序列
:param tokens: 待编码字符串
:return: 编号序列
"""
# 加上开始标记
token_ids: list[int] = [self.token_to_id(TOKEN_CLS), ]
# 加入字符串编号序列
for token in tokens:
token_ids.append(self.token_to_id(token))
# 加上结束标记
token_ids.append(self.token_to_id(TOKEN_SEP))
return token_ids
def decode(self, token_ids: typing.Iterable[int]) -> str:
"""
给定一个编号序列,将它解码成字符串
:param token_ids: 待解码的编号序列
:return: 解码出的字符串
"""
# 起止标记字符特殊处理
spec_tokens = {TOKEN_CLS, TOKEN_SEP}
# 保存解码出的字符的list
tokens: list[str] = []
for token_id in token_ids:
token = self.id_to_token(token_id)
if token in spec_tokens:
continue
tokens.append(token)
# 拼接字符串
return ''.join(tokens)
class PoetryPreprocessor:
"""
古诗词数据集的预处理器。
该类负责古诗词数据的读取,清洗和数据持久化。
"""
tokenizer: Tokenizer
"""分词器"""
poetry: list[str]
"""古诗词数据集,每一项是一首诗"""
def __init__(self, force_reclean: bool=False):
# 加载古诗词数据集
if force_reclean or (not settings.CLEAN_DATASET_PATH.is_file()):
(self.poetry, self.tokenizer) = self.__load_from_dirty()
else:
(self.poetry, self.tokenizer) = self.__load_from_clean()
def __load_from_clean(self) -> tuple[list[str], Tokenizer]:
"""直接读取清洗后的数据"""
with open(settings.CLEAN_DATASET_PATH, 'rb') as f:
return pickle.load(f)
def __load_from_dirty(self) -> tuple[list[str], Tokenizer]:
"""从原始数据加载,清洗数据后,写入缓存文件,并返回清洗后的数据"""
# 加载脏的古诗数据
with open(settings.DIRTY_DATASET_PATH, 'r', encoding='utf-8') as f:
lines = f.readlines()
# 清洗古诗数据
poetry = self.__wash_dirty_poetry(lines)
# 构建分词器
tokenizer = self.__build_tokenizer(poetry)
# 数据清理完毕
# 写入干净数据
with open(settings.CLEAN_DATASET_PATH, 'wb') as f:
pickle.dump((poetry, tokenizer), f)
# 返回结果
return poetry, tokenizer
def __wash_dirty_poetry(self, poetry: list[str]) -> list[str]:
"""
清洗给定的古诗数据。
:param poetry: 要清洗的古诗数据,每一行是一首古诗。
古诗开头是标题,然后是一个冒号(全角或半角),然后是古诗主体。
:return: 清洗完毕的古诗。
"""
# 禁用词列表,包含如下字符的诗歌将被忽略
BAD_WORDS = ['', '', '(', ')', '__', '', '', '', '', '[', ']']
# 数据集列表
clean_poetry: list[str] = []
# 逐行处理读取到的数据
for line in poetry:
# 删除空白字符
line = line.strip()
# 将全角冒号替换为半角的
line = line.replace('', ':')
# 有且只能有一个冒号用来分割标题
if line.count(':') != 1: continue
# 获取后半部分(删除标题)
_, last_part = line.split(':')
# 长度不能超过最大长度减去2是因为古诗首尾要加特殊符号
if len(last_part) > settings.POETRY_MAX_LEN - 2:
continue
# 不能包含禁止词
for bad_word in BAD_WORDS:
if bad_word in last_part:
break
else:
# 如果循环正常结束就表明没有bad words推入数据列表
clean_poetry.append(last_part)
# 返回清洗完毕的结果
return clean_poetry
def __build_tokenizer(self, poetry: list[str]) -> Tokenizer:
"""
根据给定古诗统计词频,并构建分词器。
:param poetry: 清洗完毕后的古诗,每一行是一句诗。
:return: 构建完毕的分词器。
"""
# 统计词频
counter: Counter[str] = Counter()
for line in poetry:
counter.update(line)
# 过滤掉低频词
tokens = ((token, count) for token, count in counter.items() if count >= settings.POETRY_MIN_WORD_FREQ)
# 按词频排序
tokens = sorted(tokens, key=lambda x: -x[1])
# 去掉词频,只保留词列表
tokens = list(token for token, _ in tokens)
# 将特殊词和数据集中的词拼接起来
tokens = ['[PAD]', '[UNK]', '[CLS]', '[SEP]'] + tokens
# 创建词典 token->id映射关系
token_id_dict = dict(zip(tokens, range(len(tokens))))
# 使用新词典重新建立分词器
tokenizer = Tokenizer(token_id_dict)
# 直接返回,此处无需混洗数据
return tokenizer
class PoetryDataset(Dataset):
"""适配PyTorch的古诗词Dataset"""
preprocessor: PoetryPreprocessor
def __init__(self, poetry: PoetryPreprocessor):
self.preprocessor = poetry
def __getitem__(self, index):
# 获取古诗词并编码
poetry = self.preprocessor.poetry[index]
encoded_poetry = self.preprocessor.tokenizer.encode(poetry)
# 直接返回编码后的古诗词数据数据的padding和输入输出构成由DataLoader来做。
return encoded_poetry
def __len__(self):
return len(self.preprocessor.poetry)
class PoetryDataLoader:
"""适配PyTorch的古诗词数据Loader"""
preprocessor: PoetryPreprocessor
dataset: PoetryDataset
loader: DataLoader
def __init__(self, batch_size: int, force_reclean: bool=False):
self.preprocessor = PoetryPreprocessor(force_reclean)
self.dataset = PoetryDataset(self.preprocessor)
self.loader = DataLoader(dataset=self.dataset,
batch_size=batch_size,
# 对古诗词做padding后返回
collate_fn=lambda batch: self.__collect_fn(batch),
# 混洗数据以防止过拟合
shuffle=True)
def get_vocab_size(self) -> int:
"""一个便捷的获取vocab_size的函数避免层层调用"""
return self.preprocessor.tokenizer.vocab_size
def get_tokenizer(self) -> Tokenizer:
"""一个便捷的获取Tokenizer的函数避免层层调用"""
return self.preprocessor.tokenizer
def __collect_fn(self, batch: list[list[int]]) -> tuple[torch.Tensor, torch.Tensor]:
"""
适用于DataLoader的样本收集器。
用于将上传的古诗词样本做padding后打包返回。
"""
# 计算填充长度
length = max(map(len, batch))
# 获取填充数据
padding = self.preprocessor.tokenizer.token_to_id(TOKEN_PAD)
# 开始填充
padded_batch: list[list[int]] = []
for entry in batch:
padding_length = length - len(entry)
if padding_length > 0:
# 不足就进行填充
padded_batch.append(numpy.concatenate([entry, [padding] * padding_length]))
else:
# 超过就进行截断
padded_batch.append(entry[:length])
numpy_batch = numpy.array(padded_batch)
# 生成输入和输出。
# 输入是去除最后一个字符的部分,输出是去除第一个字符的部分。
# 这么做是为了让RNN从输入推到输出下一个字符
# 此外输出要做onehot编码
input = torch.tensor(numpy_batch[:, :-1], dtype=torch.long)
output = torch.tensor(numpy_batch[:, 1:], dtype=torch.long)
# 返回结果
return input, output

View File

@@ -0,0 +1,41 @@
import torch
import torch.nn.functional as F
class TimeDistributed(torch.nn.Module):
"""模拟tensorflow中的TimeDistributed包装层因为pytorch似乎不提供这个。"""
layer: torch.nn.Module
"""内部节点"""
def __init__(self, layer: torch.nn.Module):
super(TimeDistributed, self).__init__()
self.layer = layer
def forward(self, x: torch.Tensor):
# 获取批次大小,时间步个数,特征个数
batch_size, time_steps, features = x.size()
# 把时间步维度合并到批次维度中然后运算,这样在其他层看来这就是不同的批次而已。
x = x.reshape(-1, features)
outputs: torch.Tensor = self.layer(x)
# 再把时间步维度还原出来
outputs = outputs.reshape(batch_size, time_steps, -1)
return outputs
class Rnn(torch.nn.Module):
"""循环神经网络"""
def __init__(self, vocab_size: int):
super(Rnn, self).__init__()
self.embedding = torch.nn.Embedding(vocab_size, 128)
self.lstm1 = torch.nn.LSTM(128, 128, batch_first=True, dropout=0.5)
self.lstm2 = torch.nn.LSTM(128, 128, batch_first=True, dropout=0.5)
self.timedfc = TimeDistributed(torch.nn.Linear(128, vocab_size))
def forward(self, x):
x = self.embedding(x)
x, _ = self.lstm1(x)
x, _ = self.lstm2(x)
x = self.timedfc(x)
return x

View File

@@ -0,0 +1,147 @@
from pathlib import Path
import sys
import numpy
import torch
import torch.nn.functional as F
import settings
from dataset import Tokenizer, PoetryDataLoader
from model import Rnn
sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
import gpu_utils
def generate_random_poetry(tokenizer: Tokenizer, model: Rnn, device: torch.device, s: str='') -> str:
"""
随机生成一首诗
:param tokenizer: 分词器
:param model: 用于生成古诗的模型
:param s: 用于生成古诗的起始字符串,默认为空串
:return: 一个字符串,表示一首古诗
"""
# 将初始字符串转成token
token_ids = tokenizer.encode(s)
# 去掉结束标记[SEP]
token_ids = token_ids[:-1]
while len(token_ids) < settings.POETRY_MAX_LEN:
# 进行预测其中batch_size=1
input = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0)
output: torch.Tensor = model(input.to(device))
# 计算最后一个字符的概率分布。
# 由于后续预测概率时,需要批次维度,所以方括号里第一项写:保留批次维度。
# 然后因为只有最后一个字符是预测的,其他字符都是辅助推断的,所以方括号第二项-1表示取这个最后一个字符。
# 最后,它的概率分布中不包含[PAD][UNK][CLS]的概率分布所以方括号第三项3:把这些东西删掉这些编号是Tokenizer在编译时写死的详细查看对应模块
possibilities = F.softmax(output[:, -1, 3:], dim=-1)
# 按照预测出的概率,随机选择一个词作为预测结果。
# 如果需要贪心则用argmax替代。
target_index = torch.multinomial(possibilities, num_samples=1)
# 记得把之前删除的维度加回来才是token id
target_id = target_index.item() + 3
# 把target_id加入序列
token_ids.append(target_id)
# 如果target_id是[SEP],表示输出结束,需要退出
if target_id == 3: break
# 解码并返回结果
return tokenizer.decode(token_ids)
def generate_acrostic(tokenizer: Tokenizer, model: Rnn, device: torch.device, head: str) -> str:
"""
随机生成一首藏头诗
:param tokenizer: 分词器
:param model: 用于生成古诗的模型
:param head: 藏头诗的头
:return: 一个字符串,表示一首古诗
"""
# 使用空串初始化token_ids
token_ids = tokenizer.encode('')
# 去掉结束标记[SEP],只保留[CLS]
token_ids = token_ids[:-1]
# 标点符号,这里简单的只把逗号和句号作为标点
punctuations = ['', '']
punctuation_ids = {tokenizer.token_to_id(token) for token in punctuations}
# 缓存生成的诗的list
poetry: list[str] = []
# 对于藏头诗中的每一个字,都生成一个短句
for ch in head:
# 先记录下这个字
poetry.append(ch)
# 将藏头诗的字符转成token id
token_id = tokenizer.token_to_id(ch)
# 加入到列表中去
token_ids.append(token_id)
# 开始生成一个短句
while True:
# 与generate_random_poetry函数相同的方式不断地生成诗句的下一个字。
input = torch.tensor(token_ids, dtype=torch.long).unsqueeze(0)
output: torch.Tensor = model(input.to(device))
possibilities = F.softmax(output[:, -1, 3:], dim=-1)
target_index = torch.multinomial(possibilities, num_samples=1)
target_id = target_index.item() + 3
# 把target_id加入序列
token_ids.append(target_id)
# 只有对应ID不是特殊符号的ID我们才把这个字符推入诗句中
if target_id > 3: poetry.append(tokenizer.id_to_token(target_id))
# 此外,与上面不同的是,当输出为标点符号时,我们退出当前循环,进而生成藏头诗的下一句。
if target_id in punctuation_ids: break
# 解码并返回结果
return ''.join(poetry)
class Predictor:
device: torch.device
data_loader: PoetryDataLoader
model: Rnn
def __init__(self):
self.device = gpu_utils.get_gpu_device()
self.data_loader = PoetryDataLoader(batch_size=settings.N_BATCH_SIZE)
self.model = Rnn(self.data_loader.get_vocab_size()).to(self.device)
# 加载保存好的模型参数
self.model.load_state_dict(torch.load(settings.SAVED_MODEL_PATH))
self.model.eval()
def generate_random_poetry(self, s: str = ''):
"""随机生成一首诗"""
with torch.no_grad():
print(generate_random_poetry(self.data_loader.get_tokenizer(),
self.model,
self.device,
s))
def generate_acrostic(self, s: str):
"""随机生成一首藏头诗"""
with torch.no_grad():
print(generate_acrostic(self.data_loader.get_tokenizer(),
self.model,
self.device,
s))
def main():
predictor = Predictor()
# 随机生成一首诗
predictor.generate_random_poetry()
# 给出部分信息的情况下,随机生成剩余部分
predictor.generate_random_poetry('床前明月光,')
# 生成藏头诗
predictor.generate_acrostic('好好学习天天向上')
if __name__ == "__main__":
gpu_utils.print_gpu_availability()
main()

View File

@@ -0,0 +1,19 @@
from pathlib import Path
POETRY_MAX_LEN: int = 64
"""古诗词句子最大允许长度(该长度包含首尾填充的特殊字符),超过该长度的诗句将被删除。"""
POETRY_MIN_WORD_FREQ: int = 8
"""古诗词最小允许词频,小于该词频的词将在编解码时被视为[UNK]生僻字。"""
DIRTY_DATASET_PATH: Path = Path(__file__).resolve().parent.parent / 'datasets' / 'poetry.txt'
"""脏的(未清洗的)古诗数据的路径"""
CLEAN_DATASET_PATH: Path = Path(__file__).resolve().parent.parent / 'datasets' / 'poetry.pickle'
"""干净的(已经清洗过的)古诗数据的路径"""
SAVED_MODEL_PATH: Path = Path(__file__).resolve().parent.parent / 'models' / 'rnn.pth'
"""训练完毕的模型进行保存的路径"""
N_EPOCH: int = 10
"""训练时的epoch"""
N_BATCH_SIZE: int = 50
"""训练时的batch size"""

View File

@@ -0,0 +1,79 @@
from pathlib import Path
import sys
import typing
import torch
import torchinfo
import ignite.engine
import ignite.metrics
from ignite.engine import Engine, Events
from ignite.handlers.tqdm_logger import ProgressBar
from dataset import PoetryDataLoader
from model import Rnn
from predict import generate_random_poetry
import settings
sys.path.append(str(Path(__file__).resolve().parent.parent.parent))
import gpu_utils
class Trainer:
"""核心训练器"""
device: torch.device
data_loader: PoetryDataLoader
model: Rnn
trainer: Engine
pbar: ProgressBar
def __init__(self):
# 创建训练设备,模型和数据加载器。
self.device = gpu_utils.get_gpu_device()
self.data_loader = PoetryDataLoader(batch_size=settings.N_BATCH_SIZE)
self.model = Rnn(self.data_loader.get_vocab_size()).to(self.device)
# 展示模型结构。批次为指定批次数量最大诗歌长度同时输入一定是int32
torchinfo.summary(self.model,
(settings.N_BATCH_SIZE, settings.POETRY_MAX_LEN),
dtypes=[torch.int32,])
# 优化器和损失函数
optimizer = torch.optim.Adam(self.model.parameters(), eps=1e-7)
criterion = torch.nn.CrossEntropyLoss()
# 创建训练器
self.trainer = ignite.engine.create_supervised_trainer(
self.model, optimizer, criterion, self.device,
# 由于PyTorch的交叉熵函数总是要求概率在dim=1所以要调换一下维度才能传入。
model_transform=lambda output: self.__adjust_for_loss(output))
# 将训练器关联到进度条
self.pbar = ProgressBar(persist=True)
self.pbar.attach(self.trainer, output_transform=lambda loss: {"loss": loss})
# 每次epoch后作诗一首看看结果
self.trainer.add_event_handler(
Events.EPOCH_COMPLETED,
lambda: print(generate_random_poetry(self.data_loader.get_tokenizer(), self.model, self.device))
)
def __adjust_for_loss(self, output: torch.Tensor) -> torch.Tensor:
return output.permute(0, 2, 1)
def train_model(self):
# 训练模型
self.trainer.run(self.data_loader.loader, max_epochs=settings.N_EPOCH)
def save_model(self):
# 确保保存模型的文件夹存在。
settings.SAVED_MODEL_PATH.parent.mkdir(parents=True, exist_ok=True)
# 仅保存模型参数
torch.save(self.model.state_dict(), settings.SAVED_MODEL_PATH)
print(f'Model was saved into: {settings.SAVED_MODEL_PATH}')
def main():
trainer = Trainer()
trainer.train_model()
trainer.save_model()
if __name__ == "__main__":
gpu_utils.print_gpu_availability()
main()

View File

@@ -0,0 +1,199 @@
#ANLI College of Artificial Intelligence
from collections import Counter
import math
import numpy as np
import tensorflow as tf
import settings
class Tokenizer:
"""
分词器
"""
def __init__(self, token_dict):
# 词->编号的映射
self.token_dict = token_dict
# 编号->词的映射
self.token_dict_rev = {value: key for key, value in self.token_dict.items()}
# 词汇表大小
self.vocab_size = len(self.token_dict)
def id_to_token(self, token_id):
"""
给定一个编号,查找词汇表中对应的词
:param token_id: 带查找词的编号
:return: 编号对应的词
"""
return self.token_dict_rev[token_id]
def token_to_id(self, token):
"""
给定一个词,查找它在词汇表中的编号
未找到则返回低频词[UNK]的编号
:param token: 带查找编号的词
:return: 词的编号
"""
return self.token_dict.get(token, self.token_dict['[UNK]'])
def encode(self, tokens):
"""
给定一个字符串s在头尾分别加上标记开始和结束的特殊字符并将它转成对应的编号序列
:param tokens: 待编码字符串
:return: 编号序列
"""
# 加上开始标记
token_ids = [self.token_to_id('[CLS]'), ]
# 加入字符串编号序列
for token in tokens:
token_ids.append(self.token_to_id(token))
# 加上结束标记
token_ids.append(self.token_to_id('[SEP]'))
return token_ids
def decode(self, token_ids):
"""
给定一个编号序列,将它解码成字符串
:param token_ids: 待解码的编号序列
:return: 解码出的字符串
"""
# 起止标记字符特殊处理
spec_tokens = {'[CLS]', '[SEP]'}
# 保存解码出的字符的list
tokens = []
for token_id in token_ids:
token = self.id_to_token(token_id)
if token in spec_tokens:
continue
tokens.append(token)
# 拼接字符串
return ''.join(tokens)
# 禁用词
disallowed_words = settings.DISALLOWED_WORDS
# 句子最大长度
max_len = settings.MAX_LEN
# 最小词频
min_word_frequency = settings.MIN_WORD_FREQUENCY
# mini batch 大小
batch_size = settings.BATCH_SIZE
# 加载数据集
with open(settings.DATASET_PATH, 'r', encoding='utf-8') as f:
lines = f.readlines()
# 将冒号统一成相同格式
lines = [line.replace('', ':') for line in lines]
# 数据集列表
poetry = []
# 逐行处理读取到的数据
for line in lines:
# 有且只能有一个冒号用来分割标题
if line.count(':') != 1:
continue
# 后半部分不能包含禁止词
__, last_part = line.split(':')
ignore_flag = False
for dis_word in disallowed_words:
if dis_word in last_part:
ignore_flag = True
break
if ignore_flag:
continue
# 长度不能超过最大长度
if len(last_part) > max_len - 2:
continue
poetry.append(last_part.replace('\n', ''))
# 统计词频
counter = Counter()
for line in poetry:
counter.update(line)
# 过滤掉低频词
_tokens = [(token, count) for token, count in counter.items() if count >= min_word_frequency]
# 按词频排序
_tokens = sorted(_tokens, key=lambda x: -x[1])
# 去掉词频,只保留词列表
_tokens = [token for token, count in _tokens]
# 将特殊词和数据集中的词拼接起来
_tokens = ['[PAD]', '[UNK]', '[CLS]', '[SEP]'] + _tokens
# 创建词典 token->id映射关系
token_id_dict = dict(zip(_tokens, range(len(_tokens))))
# 使用新词典重新建立分词器
tokenizer = Tokenizer(token_id_dict)
# 混洗数据
np.random.shuffle(poetry)
class PoetryDataGenerator:
"""
古诗数据集生成器
"""
def __init__(self, data, random=False):
# 数据集
self.data = data
# batch size
self.batch_size = batch_size
# 每个epoch迭代的步数
self.steps = int(math.floor(len(self.data) / self.batch_size))
# 每个epoch开始时是否随机混洗
self.random = random
def sequence_padding(self, data, length=None, padding=None):
"""
将给定数据填充到相同长度
:param data: 待填充数据
:param length: 填充后的长度不传递此参数则使用data中的最大长度
:param padding: 用于填充的数据,不传递此参数则使用[PAD]的对应编号
:return: 填充后的数据
"""
# 计算填充长度
if length is None:
length = max(map(len, data))
# 计算填充数据
if padding is None:
padding = tokenizer.token_to_id('[PAD]')
# 开始填充
outputs = []
for line in data:
padding_length = length - len(line)
# 不足就进行填充
if padding_length > 0:
outputs.append(np.concatenate([line, [padding] * padding_length]))
# 超过就进行截断
else:
outputs.append(line[:length])
return np.array(outputs)
def __len__(self):
return self.steps
def __iter__(self):
total = len(self.data)
# 是否随机混洗
if self.random:
np.random.shuffle(self.data)
# 迭代一个epoch每次yield一个batch
for start in range(0, total, self.batch_size):
end = min(start + self.batch_size, total)
batch_data = []
# 逐一对古诗进行编码
for single_data in self.data[start:end]:
batch_data.append(tokenizer.encode(single_data))
# 填充为相同长度
batch_data = self.sequence_padding(batch_data)
# yield x,y
yield batch_data[:, :-1], tf.one_hot(batch_data[:, 1:], tokenizer.vocab_size)
del batch_data
def for_fit(self):
"""
创建一个生成器,用于训练
"""
# 死循环当数据训练一个epoch之后重新迭代数据
while True:
# 委托生成器
yield from self.__iter__()

View File

@@ -0,0 +1,16 @@
#ANLI College of Artificial Intelligence
import tensorflow as tf
from dataset import tokenizer
import settings
import utils
# 加载训练好的模型
model = tf.keras.models.load_model(settings.BEST_MODEL_PATH)
# 随机生成一首诗
print(utils.generate_random_poetry(tokenizer, model))
# 给出部分信息的情况下,随机生成剩余部分
print(utils.generate_random_poetry(tokenizer, model, s='床前明月光,'))
# 生成藏头诗
print(utils.generate_acrostic(tokenizer, model, head='好好学习天天向上'))

View File

@@ -0,0 +1,19 @@
#ANLI College of Artificial Intelligence
# 禁用词,包含如下字符的唐诗将被忽略
DISALLOWED_WORDS = ['', '', '(', ')', '__', '', '', '', '', '[', ']']
# 句子最大长度
MAX_LEN = 64
# 最小词频
MIN_WORD_FREQUENCY = 8
# 训练的batch size
BATCH_SIZE = 16
# 数据集路径
DATASET_PATH = './poetry.txt'
# 每个epoch训练完成后随机生成SHOW_NUM首古诗作为展示
SHOW_NUM = 5
# 共训练多少个epoch
TRAIN_EPOCHS = 10
# 最佳权重保存路径
BEST_MODEL_PATH = './best_model.h5'

View File

@@ -0,0 +1,35 @@
import tensorflow as tf
from dataset import PoetryDataGenerator, tokenizer, poetry
import settings
import utils
model = tf.keras.Sequential([
tf.keras.layers.Input((None,)),
tf.keras.layers.Embedding(input_dim=tokenizer.vocab_size, output_dim=128),
tf.keras.layers.LSTM(128, dropout=0.5, return_sequences=True),
tf.keras.layers.LSTM(128, dropout=0.5, return_sequences=True),
tf.keras.layers.TimeDistributed(tf.keras.layers.Dense(tokenizer.vocab_size, activation='softmax')),
])
model.summary()
model.compile(optimizer=tf.keras.optimizers.Adam(), loss=tf.keras.losses.categorical_crossentropy)
class Evaluate(tf.keras.callbacks.Callback):
def __init__(self):
super().__init__()
self.lowest = 1e10
def on_epoch_end(self, epoch, logs=None):
if logs['loss'] <= self.lowest:
self.lowest = logs['loss']
model.save(settings.BEST_MODEL_PATH)
print()
for i in range(settings.SHOW_NUM):
print(utils.generate_random_poetry(tokenizer, model))
data_generator = PoetryDataGenerator(poetry, random=False)
model.fit_generator(data_generator.for_fit(),
steps_per_epoch=data_generator.steps,
epochs=settings.TRAIN_EPOCHS,
callbacks=[Evaluate()])

View File

@@ -0,0 +1,86 @@
import numpy as np
import settings
def generate_random_poetry(tokenizer, model, s=''):
"""
随机生成一首诗
:param tokenizer: 分词器
:param model: 用于生成古诗的模型
:param s: 用于生成古诗的起始字符串,默认为空串
:return: 一个字符串,表示一首古诗
"""
# 将初始字符串转成token
token_ids = tokenizer.encode(s)
# 去掉结束标记[SEP]
token_ids = token_ids[:-1]
while len(token_ids) < settings.MAX_LEN:
# 进行预测只保留第一个样例我们输入的样例数只有1的、最后一个token的预测的、不包含[PAD][UNK][CLS]的概率分布
output = model(np.array([token_ids, ], dtype=np.int32))
_probas = output.numpy()[0, -1, 3:]
del output
# print(_probas)
# 按照出现概率对所有token倒序排列
p_args = _probas.argsort()[::-1][:100]
# 排列后的概率顺序
p = _probas[p_args]
# 先对概率归一
p = p / sum(p)
# 再按照预测出的概率,随机选择一个词作为预测结果
target_index = np.random.choice(len(p), p=p)
target = p_args[target_index] + 3
# 保存
token_ids.append(target)
if target == 3:
break
return tokenizer.decode(token_ids)
def generate_acrostic(tokenizer, model, head):
"""
随机生成一首藏头诗
:param tokenizer: 分词器
:param model: 用于生成古诗的模型
:param head: 藏头诗的头
:return: 一个字符串,表示一首古诗
"""
# 使用空串初始化token_ids加入[CLS]
token_ids = tokenizer.encode('')
token_ids = token_ids[:-1]
# 标点符号,这里简单的只把逗号和句号作为标点
punctuations = ['', '']
punctuation_ids = {tokenizer.token_to_id(token) for token in punctuations}
# 缓存生成的诗的list
poetry = []
# 对于藏头诗中的每一个字,都生成一个短句
for ch in head:
# 先记录下这个字
poetry.append(ch)
# 将藏头诗的字符转成token id
token_id = tokenizer.token_to_id(ch)
# 加入到列表中去
token_ids.append(token_id)
# 开始生成一个短句
while True:
# 进行预测只保留第一个样例我们输入的样例数只有1的、最后一个token的预测的、不包含[PAD][UNK][CLS]的概率分布
output = model(np.array([token_ids, ], dtype=np.int32))
_probas = output.numpy()[0, -1, 3:]
del output
# 按照出现概率对所有token倒序排列
p_args = _probas.argsort()[::-1][:100]
# 排列后的概率顺序
p = _probas[p_args]
# 先对概率归一
p = p / sum(p)
# 再按照预测出的概率,随机选择一个词作为预测结果
target_index = np.random.choice(len(p), p=p)
target = p_args[target_index] + 3
# 保存
token_ids.append(target)
# 只有不是特殊字符时才保存到poetry里面去
if target > 3:
poetry.append(tokenizer.id_to_token(target))
if target in punctuation_ids:
break
return ''.join(poetry)