![](https://i-blog.csdnimg.cn/blog_migrate/b169d0aec0105f3bfe5ee1bb509c9ac9.gif#pic_center)
✅作者简介:人工智能专业本科在读,喜欢计算机与编程,写博客记录自己的学习历程。
🍎个人主页:小嗷犬的个人主页
🍊个人网站:小嗷犬的技术小站
🥭个人信条:为天地立心,为生民立命,为往圣继绝学,为万世开太平。
本文目录
简介(Introduction)
本项目是一个基于 Transformer 的中文对联生成器,使用 PyTorch 构建模型,使用 Gradio 构建 Web UI。
数据集:https://www.kaggle.com/datasets/marquis03/chinese-couplets-dataset
GitHub 仓库:https://github.com/Marquis03/Chinese-Couplets-Generator-based-on-Transformer
Gitee 仓库:https://gitee.com/marquis03/Chinese-Couplets-Generator-based-on-Transformer
项目结构(Structure)
.
├── config
│ ├── __init__.py
│ └── config.py
├── data
│ ├── fixed_couplets_in.txt
│ └── fixed_couplets_out.txt
├── dataset
│ ├── __init__.py
│ └── dataset.py
├── img
│ ├── history.png
│ ├── lr_schedule.png
│ └── webui.gif
├── model
│ ├── __init__.py
│ └── model.py
├── trained
│ ├── vocab.pkl
│ └── CoupletsTransformer_best.pth
├── utils
│ ├── __init__.py
│ └── EarlyStopping.py
├── LICENSE
├── README.md
├── requirements.txt
├── train.py
└── webui.py
部署(Deployment)
克隆项目(Clone Project)
git clone https://github.com/Marquis03/Chinese-Couplets-Generator-based-on-Transformer.git
cd Chinese-Couplets-Generator-based-on-Transformer
安装依赖(Requirements)
pip install -r requirements.txt
训练模型(Train Model)
python train.py
Kaggle Notebook: https://www.kaggle.com/code/marquis03/chinese-couplets-generator-based-on-transformer
启动 Web UI(Start Web UI)
python webui.py
项目演示(Demo)
Web UI
学习率变化(Learning Rate Schedule)
训练历史(Training History)
代码(Code)
配置参数(Config)
该部分用于配置项目的参数,包括全局参数、路径参数、模型参数、训练参数和日志参数。
对应项目文件为 config/config.py
。
import os
import sys
import time
import torch
from loguru import logger
class Config:
def __init__(self):
# global
self.seed = 0
self.cuDNN = True
self.debug = False
self.num_workers = 0
self.str_time = time.strftime("%Y-%m-%dT%H%M%S", time.localtime(time.time()))
# path
self.project_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
self.dataset_dir = os.path.join(self.project_dir, "data")
self.in_path = os.path.join(self.dataset_dir, "fixed_couplets_in.txt")
self.out_path = os.path.join(self.dataset_dir, "fixed_couplets_out.txt")
self.log_dir = os.path.join(self.project_dir, "logs")
self.save_dir = os.path.join(self.log_dir, self.str_time)
self.img_save_dir = os.path.join(self.save_dir, "images")
self.model_save_dir = os.path.join(self.save_dir, "checkpoints")
for path in (
self.log_dir,
self.save_dir,
self.img_save_dir,
self.model_save_dir,
):
if not os.path.exists(path):
os.makedirs(path)
# model
self.d_model = 256
self.num_head = 8
self.num_encoder_layers = 2
self.num_decoder_layers = 2
self.dim_feedforward = 1024
self.dropout = 0.1
# train
self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
self.batch_size = 128
self.val_ratio = 0.1
self.epochs = 20
self.warmup_ratio = 0.12
self.lr_max = 1e-3
self.lr_min = 1e-4
self.beta1 = 0.9
self.beta2 = 0.98
self.epsilon = 10e-9
self.weight_decay = 0.01
self.early_stop = True
self.patience = 4
self.delta = 0
# log
logger.remove()
level_std = "DEBUG" if self.debug else "INFO"
logger.add(
sys.stdout,
colorize=True,
format="[<green>{time:YYYY-MM-DD HH:mm:ss,SSS}</green>|<level>{level: <8}</level>|<cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan>] >>> <level>{message}</level>",
level=level_std,
)
logger.add(
os.path.join(self.save_dir, f"{self.str_time}.log"),
format="[{time:YYYY-MM-DD HH:mm:ss,SSS}|{level: <8}|{name}:{function}:{line}] >>> {message}",
level="INFO",
)
logger.info("### Config:")
for key, value in self.__dict__.items():
logger.info(f"### {key:20} = {value}")
数据集(Dataset)
该部分用于定义词典、数据集以及相关函数,包括数据的加载、词典的构建、数据集的封装和数据集的加载器。
对应项目文件为 dataset/dataset.py
。
from collections import Counter
import torch
import numpy as np
from torch.utils.data import Dataset, DataLoader
def load_data(filepaths, tokenizer=lambda s: s.strip().split()):
raw_in_iter = iter(open(filepaths[0], encoding="utf8"))
raw_out_iter = iter(open(filepaths[1], encoding="utf8"))
return list(zip(map(tokenizer, raw_in_iter), map(tokenizer, raw_out_iter)))
class Vocab(object):
UNK = "<unk>" # 0
PAD = "<pad>" # 1
BOS = "<bos>" # 2
EOS = "<eos>" # 3
def __init__(self, data=None, min_freq=1):
counter = Counter()
for lines in data:
counter.update(lines[0])
counter.update(lines[1])
self.word2idx = {Vocab.UNK: 0, Vocab.PAD: 1, Vocab.BOS: 2, Vocab.EOS: 3}
self.idx2word = {0: Vocab.UNK, 1: Vocab.PAD, 2: Vocab.BOS, 3: Vocab.EOS}
idx = 4
for word, freq in counter.items():
if freq >= min_freq:
self.word2idx[word] = idx
self.idx2word[idx] = word
idx += 1
def __len__(self):
return len(self.word2idx)
def __getitem__(self, word):
return self.word2idx.get(word, 0)
def __call__(self, word):
if not isinstance(word, (list, tuple)):
return self[word]
return [self[w] for w in word]
def to_tokens(self, indices):
if not isinstance(indices, (list, tuple, np.ndarray, torch.Tensor)):
return self.idx2word[int(indices)]
return [self.idx2word[int(i)] for i in indices]
def pad_sequence(sequences, batch_first=False, padding_value=0):
max_len = max([s.size(0) for s in sequences])
out_tensors = []
for tensor in sequences:
padding_content = [padding_value] * (max_len - tensor.size(0))
tensor = torch.cat([tensor, torch.tensor(padding_content)], dim=0)
out_tensors.append(tensor)
out_tensors = torch.stack(out_tensors, dim=1)
if batch_first:
out_tensors = out_tensors.transpose(0, 1)
return out_tensors.long()
class CoupletsDataset(Dataset):
def __init__(self, data, vocab):
self.data = data
self.vocab = vocab
self.PAD_IDX = self.vocab[self.vocab.PAD]
self.BOS_IDX = self.vocab[self.vocab.BOS]
self.EOS_IDX = self.vocab[self.vocab.EOS]
def __len__(self):
return len(self.data)
def __getitem__(self, index):
raw_in, raw_out = self.data[index]
in_tensor_ = torch.LongTensor(self.vocab(raw_in))
out_tensor_ = torch.LongTensor(self.vocab(raw_out))
return in_tensor_, out_tensor_
def collate_fn(self, batch):
in_batch, out_batch = [], []
for in_, out_ in batch:
in_batch.append(in_)
out_ = torch.cat(
[
torch.LongTensor([self.BOS_IDX]),
out_,
torch.LongTensor([self.EOS_IDX]),
],
dim=0,
)
out_batch.append(out_)
in_batch = pad_sequence(in_batch, True, self.PAD_IDX)
out_batch = pad_sequence(out_batch, True, self.PAD_IDX)
return in_batch, out_batch
def get_loader(self, batch_size, shuffle=False, num_workers=0):
return DataLoader(
self,
batch_size=batch_size,
shuffle=shuffle,
num_workers=num_workers,
collate_fn=self.collate_fn,
pin_memory=True,
)
模型(Model)
该部分用于定义模型,包括 TokenEmbedding、PositionalEncoding 和 CoupletsTransformer。
对应项目文件为 model/model.py
。
import math
import torch
import torch.nn as nn
class TokenEmbedding(nn.Module):
def __init__(self, vocab_size, emb_size):
super(TokenEmbedding, self).__init__()
self.embedding = nn.Embedding(vocab_size, emb_size)
self.emb_size = emb_size
def forward(self, tokens):
return self.embedding(tokens) * math.sqrt(self.emb_size)
class PositionalEncoding(nn.Module):
def __init__(self, d_model, dropout=0.1, max_len=5000):
super(PositionalEncoding, self).__init__()
self.dropout = nn.Dropout(p=dropout)
pe = torch.zeros(max_len, d_model)
position = torch.arange(0, max_len, dtype=torch.float).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, d_model, 2).float() * (-math.log(10000.0) / d_model)
)
pe[:, 0::2] = torch.sin(position * div_term)
pe[:, 1::2] = torch.cos(position * div_term)
pe = pe.unsqueeze(0)
self.register_buffer("pe", pe)
def forward(self, x):
x = x + self.pe[:, : x.size(1)]
return self.dropout(x)
class CoupletsTransformer(nn.Module):
def __init__(
self,
vocab_size,
d_model=512,
nhead=8,
num_encoder_layers=6,
num_decoder_layers=6,
dim_feedforward=2048,
dropout=0.1,
):
super(CoupletsTransformer, self).__init__()
self.name = "CoupletsTransformer"
self.token_embedding = TokenEmbedding(vocab_size, d_model)
self.pos_embedding = PositionalEncoding(d_model, dropout)
self.transformer = nn.Transformer(
d_model=d_model,
nhead=nhead,
num_encoder_layers=num_encoder_layers,
num_decoder_layers=num_decoder_layers,
dim_feedforward=dim_feedforward,
dropout=dropout,
batch_first=True,
)
self.fc = nn.Linear(d_model, vocab_size)
self._reset_parameters()
def _reset_parameters(self):
for p in self.parameters():
if p.dim() > 1:
nn.init.xavier_uniform_(p)
def forward(self, src, tgt, padding_value=0):
src_embed = self.token_embedding(src) # [batch_size, src_len, embed_dim]
src_embed = self.pos_embedding(src_embed) # [batch_size, src_len, embed_dim]
tgt_embed = self.token_embedding(tgt) # [batch_size, tgt_len, embed_dim]
tgt_embed = self.pos_embedding(tgt_embed) # [batch_size, tgt_len, embed_dim]
tgt_mask = self.transformer.generate_square_subsequent_mask(tgt.size(-1)).to(
tgt.device
)
src_key_padding_mask = src == padding_value # [batch_size, src_len]
tgt_key_padding_mask = tgt == padding_value # [batch_size, tgt_len]
outs = self.transformer(
src=src_embed,
tgt=tgt_embed,
tgt_mask=tgt_mask,
src_key_padding_mask=src_key_padding_mask,
tgt_key_padding_mask=tgt_key_padding_mask,
memory_key_padding_mask=src_key_padding_mask,
) # [batch_size, tgt_len, embed_dim]
logits = self.fc(outs) # [batch_size, tgt_len, vocab_size]
return logits
def encoder(self, src):
src_embed = self.token_embedding(src)
src_embed = self.pos_embedding(src_embed)
memory = self.transformer.encoder(src_embed)
return memory
def decoder(self, tgt, memory):
tgt_embed = self.token_embedding(tgt)
tgt_embed = self.pos_embedding(tgt_embed)
outs = self.transformer.decoder(tgt_embed, memory=memory)
return outs
def generate(self, text, vocab):
self.eval()
device = next(self.parameters()).device
max_len = len(text)
src = torch.LongTensor(vocab(list(text))).unsqueeze(0).to(device)
memory = self.encoder(src)
l_out = [vocab.BOS]
for i in range(max_len):
tgt = torch.LongTensor(vocab(l_out)).unsqueeze(0).to(device)
outs = self.decoder(tgt, memory)
prob = self.fc(outs[:, -1, :])
next_token = vocab.to_tokens(prob.argmax(1).item())
if next_token == vocab.EOS:
break
l_out.append(next_token)
return "".join(l_out[1:])
工具(Utils)
该部分用于定义工具函数,包括 EarlyStopping。
对应项目文件为 utils/EarlyStopping.py
。
class EarlyStopping(object):
def __init__(self, patience=7, delta=0):
self.patience = patience
self.counter = 0
self.best_score = None
self.early_stop = False
self.val_loss_min = float("inf")
self.delta = delta
def __call__(self, val_loss, model):
score = -val_loss
if self.best_score is None:
self.best_score = score
elif score < self.best_score + self.delta:
self.counter += 1
if self.counter >= self.patience:
self.early_stop = True
else:
self.best_score = score
self.counter = 0
训练(Train)
该部分用于定义训练函数,包括训练、验证和保存模型。
对应项目文件为 train.py
。
import os
import gc
import time
import math
import random
import joblib
import warnings
warnings.filterwarnings("ignore")
import numpy as np
import pandas as pd
import seaborn as sns
from loguru import logger
from tqdm.auto import tqdm
import matplotlib.pyplot as plt
from sklearn.model_selection import train_test_split
sns.set_theme(
style="darkgrid", font_scale=1.2, font="SimHei", rc={"axes.unicode_minus": False}
)
import torch
from torch import nn, optim
from torch.optim.lr_scheduler import LambdaLR
from config import Config
from model import CoupletsTransformer
from dataset import load_data, Vocab, CoupletsDataset
from utils import EarlyStopping
def train_model(
config, model, train_loader, val_loader, optimizer, criterion, scheduler
):
model = model.to(config.device)
best_loss = float("inf")
history = []
model_path = os.path.join(config.model_save_dir, f"{model.name}_best.pth")
if config.early_stop:
early_stopping = EarlyStopping(patience=config.patience, delta=config.delta)
for epoch in tqdm(range(1, config.epochs + 1), desc=f"All"):
train_loss = train_one_epoch(
config, model, train_loader, optimizer, criterion, scheduler
)
val_loss = evaluate(config, model, val_loader, criterion)
perplexity = math.exp(val_loss)
history.append((epoch, train_loss, val_loss))
msg = f"Epoch {epoch}/{config.epochs}, Train Loss: {train_loss:.6f}, Val Loss: {val_loss:.6f}, Perplexity: {perplexity:.4f}"
logger.info(msg)
if val_loss < best_loss:
logger.info(
f"Val loss decrease from {best_loss:>10.6f} to {val_loss:>10.6f}"
)
torch.save(model.state_dict(), model_path)
best_loss = val_loss
if config.early_stop:
early_stopping(val_loss, model)
if early_stopping.early_stop:
logger.info(f"Early stopping at epoch {epoch}")
break
logger.info(f"Save best model with val loss {best_loss:.6f} to {model_path}")
model_path = os.path.join(config.model_save_dir, f"{model.name}_last.pth")
torch.save(model.state_dict(), model_path)
logger.info(f"Save last model with val loss {val_loss:.6f} to {model_path}")
history = pd.DataFrame(
history, columns=["Epoch", "Train Loss", "Val Loss"]
).set_index("Epoch")
history.plot(
subplots=True, layout=(1, 2), sharey="row", figsize=(14, 6), marker="o", lw=2
)
history_path = os.path.join(config.img_save_dir, "history.png")
plt.savefig(history_path, dpi=300)
logger.info(f"Save history to {history_path}")
def train_one_epoch(config, model, train_loader, optimizer, criterion, scheduler):
model.train()
train_loss = 0
for src, tgt in tqdm(train_loader, desc=f"Epoch", leave=False):
src, tgt = src.to(config.device), tgt.to(config.device)
output = model(src, tgt[:, :-1], config.PAD_IDX)
output = output.contiguous().view(-1, output.size(-1))
tgt = tgt[:, 1:].contiguous().view(-1)
loss = criterion(output, tgt)
train_loss += loss.item()
optimizer.zero_grad()
loss.backward()
optimizer.step()
scheduler.step()
return train_loss / len(train_loader)
def evaluate(config, model, val_loader, criterion):
model.eval()
val_loss = 0
with torch.no_grad():
for src, tgt in tqdm(val_loader, desc=f"Val", leave=False):
src, tgt = src.to(config.device), tgt.to(config.device)
output = model(src, tgt[:, :-1], config.PAD_IDX)
output = output.contiguous().view(-1, output.size(-1))
tgt = tgt[:, 1:].contiguous().view(-1)
loss = criterion(output, tgt)
val_loss += loss.item()
return val_loss / len(val_loader)
def test_model(model, data, vocab):
model.eval()
for src_text, tgt_text in data:
src_text, tgt_text = "".join(src_text), "".join(tgt_text)
out_text = model.generate(src_text, vocab)
logger.info(f"\nInput: {src_text}\nTarget: {tgt_text}\nOutput: {out_text}")
def seed_everything(seed):
os.environ["PYTHONHASHSEED"] = str(seed)
random.seed(seed)
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed_all(seed)
def main():
config = Config()
# Set random seed
seed_everything(config.seed)
logger.info(f"Set random seed to {config.seed}")
# Set cuDNN
if config.cuDNN:
torch.backends.cudnn.enabled = True
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
# Load data
data = load_data([config.in_path, config.out_path])
if config.debug:
data = data[:1000]
logger.info(f"Load {len(data)} couplets")
# Build vocab
vocab = Vocab(data)
vocab_size = len(vocab)
logger.info(f"Build vocab with {vocab_size} tokens")
vocab_path = os.path.join(config.model_save_dir, "vocab.pkl")
joblib.dump(vocab, vocab_path)
logger.info(f"Save vocab to {vocab_path}")
# Build dataset
data_train, data_val = train_test_split(
data, test_size=config.val_ratio, random_state=config.seed, shuffle=True
)
train_dataset = CoupletsDataset(data_train, vocab)
val_dataset = CoupletsDataset(data_val, vocab)
config.PAD_IDX = train_dataset.PAD_IDX
logger.info(f"Build train dataset with {len(train_dataset)} samples")
logger.info(f"Build val dataset with {len(val_dataset)} samples")
# Build dataloader
train_loader = train_dataset.get_loader(
config.batch_size, shuffle=True, num_workers=config.num_workers
)
val_loader = val_dataset.get_loader(
config.batch_size, shuffle=False, num_workers=config.num_workers
)
logger.info(f"Build train dataloader with {len(train_loader)} batches")
logger.info(f"Build val dataloader with {len(val_loader)} batches")
# Build model
model = CoupletsTransformer(
vocab_size=vocab_size,
d_model=config.d_model,
nhead=config.num_head,
num_encoder_layers=config.num_encoder_layers,
num_decoder_layers=config.num_decoder_layers,
dim_feedforward=config.dim_feedforward,
dropout=config.dropout,
)
logger.info(f"Build model with {model.name}")
# Build optimizer
optimizer = optim.AdamW(
model.parameters(),
lr=1,
betas=(config.beta1, config.beta2),
eps=config.epsilon,
weight_decay=config.weight_decay,
)
# Build criterion
criterion = nn.CrossEntropyLoss(ignore_index=config.PAD_IDX, reduction="mean")
# Build scheduler
lr_max, lr_min = config.lr_max, config.lr_min
T_max = config.epochs * len(train_loader)
warm_up_iter = int(T_max * config.warmup_ratio)
def WarmupExponentialLR(cur_iter):
gamma = math.exp(math.log(lr_min / lr_max) / (T_max - warm_up_iter))
if cur_iter < warm_up_iter:
return (lr_max - lr_min) * (cur_iter / warm_up_iter) + lr_min
else:
return lr_max * gamma ** (cur_iter - warm_up_iter)
scheduler = LambdaLR(optimizer, lr_lambda=WarmupExponentialLR)
df_lr = pd.DataFrame(
[WarmupExponentialLR(i) for i in range(T_max)],
columns=["Learning Rate"],
)
plt.figure(figsize=(10, 6))
sns.lineplot(data=df_lr, linewidth=2)
plt.title("Learning Rate Schedule")
plt.xlabel("Iteration")
plt.ylabel("Learning Rate")
lr_img_path = os.path.join(config.img_save_dir, "lr_schedule.png")
plt.savefig(lr_img_path, dpi=300)
logger.info(f"Save learning rate schedule to {lr_img_path}")
# Garbage collect
gc.collect()
torch.cuda.empty_cache()
# Train model
train_model(
config, model, train_loader, val_loader, optimizer, criterion, scheduler
)
# Test model
test_model(model, data_val[:10], vocab)
if __name__ == "__main__":
main()
Web UI
该部分用于定义 Web UI,包括输入、输出和启动 Web UI。
对应项目文件为 webui.py
。
import random
import joblib
import torch
import gradio as gr
from dataset import Vocab
from model import CoupletsTransformer
data_path = "./data/fixed_couplets_in.txt"
vocab_path = "./trained/vocab.pkl"
model_path = "./trained/CoupletsTransformer_best.pth"
vocab = joblib.load(vocab_path)
vocab_size = len(vocab)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = CoupletsTransformer(
vocab_size,
d_model=256,
nhead=8,
num_encoder_layers=2,
num_decoder_layers=2,
dim_feedforward=1024,
dropout=0.1,
).to(device)
model.load_state_dict(torch.load(model_path))
model.eval()
example = (
line.replace(" ", "").strip() for line in iter(open(data_path, encoding="utf8"))
)
example = [line for line in example if len(line) > 5]
example = random.sample(example, 300)
def generate_couplet(vocab, model, src_text):
if not src_text:
return "上联不能为空"
out_text = model.generate(src_text, vocab)
return out_text
input_text = gr.Textbox(
label="上联",
placeholder="在这里输入上联",
max_lines=1,
lines=1,
show_copy_button=True,
autofocus=True,
)
output_text = gr.Textbox(
label="下联",
placeholder="在这里生成下联",
max_lines=1,
lines=1,
show_copy_button=True,
)
demo = gr.Interface(
fn=lambda x: generate_couplet(vocab, model, x),
inputs=input_text,
outputs=output_text,
title="中文对联生成器",
description="输入上联,生成下联",
allow_flagging="never",
submit_btn="生成下联",
clear_btn="清空",
examples=example,
examples_per_page=50,
)
demo.launch()