README_CN.md

---

TextBox (妙笔)

“李太白少时，梦所用之笔头上生花后天才赡逸，名闻天下。”——王仁裕《开元天宝遗事·梦笔头生花》

文档 | 模型 | 数据集 | 论文 | English Version

TextBox是基于Python和PyTorch开发的，用于在一个统一的、全面的、高效的框架中复现和开发文本生成算法，主要面向研究者使用。我们的库包括21种文本生成算法，涵盖了两个主要任务：

• 无条件（无输入）生成
• 条件（Seq2Seq）生成，包括机器翻译、文本摘要、对话系统和属性文本生成

我们支持9个文本生成的基准数据集，用户可以使用我们的库重新处理原始数据，或者简单地下载我们团队已经处理好的数据集。

图片: TextBox整体框架

特点

• 统一模块化框架： TextBox基于PyTorch构建，通过将不同的模型解耦为一组高度可重用的模块，实现高度模块化设计。
• 全面的基准模型、数据集和标准评估： TextBox还包含广泛的基准文本生成模型，涵盖VAE、GAN、基于RNN或Transformer的模型，以及预训练语言模型（PLM）等类别。
• 高度灵活及拓展性强的框架： TextBox在文本生成模型部分提供了多种常用的函数与模块的接口。例如RNN encoder-decoder, Transformer encoder-decoder以及各种预训练语言模型。
• 简单便捷的上手： TextBox提供了灵活的配置文件，生手不用修改源代码即可运行实验，研究者也可以通过修改少量配置文件进行量化实验。

安装

TextBox 需要的安装条件:

• Python >= 3.6.2

• torch >= 1.6.0. 请根据你的CUDA版本和NVIDIA驱动版本参考 官方指南 来安装合适的版本。

• GCC >= 5.1.0

通过 pip 安装

pip install textbox

如果你安装fast_bleu时遇到了问题，对于Linux请保证GCC >= 5.1.0；对于Windows，请使用fast_bleu_wheel4windows目录下的wheel文件进行安装；对于MacOS，请使用如下命令安装：

pip install fast-bleu --install-option="--CC=<path-to-gcc>" --install-option="--CXX=<path-to-g++>"

成功安装后fast_bleu，重新安装textbox即可。

通过源文件安装

git clone https://github.com/RUCAIBox/TextBox.git && cd TextBox
pip install -e . --verbose

快速上手

从源文件开始

下载TextBox源码后，你可以使用下述提供的脚本快速上手使用该库：

python run_textbox.py

上述脚本会在COCO数据集上运行RNN模型进行无条件生成。一般而言，这个运行示例会执行几分钟。我们可以获得如 example.log 这样的输出日志。

如果你想要改变运行的参数，如rnn_type，max_vocab_size，只需要按照你的需要设定额外的命令参数：

python run_textbox.py --rnn_type=lstm --max_vocab_size=4000

我们也支持修改相应数据集和模型 properties 文件夹中的YAML配置文件，并将其包含在命令行中。

如果你想修改模型、数据集或任务类型，只需通过修改相应的命令参数来运行脚本：

python run_textbox.py --model=[model_name] --dataset=[dataset_name]

model_name 是将被运行的模型，比如RNN或者BART。 我们实现了的模型可以在 模型 中找到。

如果你想要修改数据集，请参考 数据集。

从API开始

如果TextBox是由pip安装的，你可以创建一个新的python文件，下载数据集，并编写和运行以下代码：

from textbox.quick_start import run_textbox

run_textbox(config_dict={'model': 'RNN',
                         'dataset': 'COCO',
                         'data_path': './dataset')

这将在COCO数据集上进行RNN模型的训练和测试。

如果你想运行不同的模型、参数或数据集，可以使用与 从源文件开始 相同的操作。

使用预训练语言模型

TextBox支持部分预训练语言模型进行文本生成任务，下面以GPT-2为例，展示我们如何利用预训练语言模型进行fine-tuning。

1. 从Hugging Face提供的模型源 (https://huggingface.co/gpt2/tree/main) 中下载GPT-2模型，包括config.json, merges.txt, pytorch_model.bin, tokenizer.json和vocab.json五个文件，将其放在与textbox同级的文件夹下，例如pretrained_model/gpt2。

2. 下载好模型之后，直接通过脚本运行：

python run_textbox.py --model=GPT2 --dataset=COCO \
                      --pretrained_model_path=pretrained_model/gpt2

使用分布式数据并行

TextBox支持方便地使用多块GPU训练模型，你不需要修改模型，只需要运行脚本：

python -m torch.distributed.launch --nproc_per_node=[gpu_num] \
       run_textbox.py --model=[model_name] \
       --dataset=[dataset_name] --gpu_id=[gpu_ids] --DDP=True

gpu_num是你想用来训练的GPU数量（例如4）， gpu_ids是你使用的GPU ID列表（例如0,1,2,3）.

注意：我们仅支持使用分布式数据并行训练端到端的模型，我们将在未来支持非端到端的模型，例如GAN。

架构

上述图片展示了TextBox的整体架构。程序的运行需要从文件、命令行或参数字典中获取实验参数配置，数据集和模型会根据设置的配置进行初始化，之后执行模块负责对模型进行训练和评估。获取更多接口相关的细节可以参考说明文档。

模型

我们总共实现了包括无条件生成和条件生成在内的21个文本生成模型，可以参照下表：

类别	模型	引用
VAE	LSTMVAE	(Bowman et al., 2016)
	CNNVAE	(Yang et al., 2017)
	HybridVAE	(Semeniuta et al., 2017)
	CVAE	(Li et al., 2018)
GAN	SeqGAN	(Yu et al., 2017)
	TextGAN	(Zhang et al., 2017)
	RankGAN	(Lin et al., 2017)
	MaliGAN	(Che et al., 2017)
	LeakGAN	(Guo et al., 2018)
	MaskGAN	(Fedus et al., 2018)
PLM	GPT-2	(Radford et al., 2019)
	XLNet	(Yang et al., 2019)
	BERT2BERT	(Rothe et al., 2020)
	BART	(Lewis et al., 2020)
	T5	(Raffel et al., 2020)
	ProphetNet	(Qi et al., 2020)
Seq2Seq	RNN	(Sutskever et al., 2014)
	Transformer	(Vaswani et al., 2017b)
	Context2Seq	(Tang et al., 2016)
	Attr2Seq	(Dong et al., 2017)
	HRED	(Serban et al., 2016)

数据集

我们总共收集了9个在上述提及的6类文本生成任务中常用的数据集，这些数据集可以通过Google Drive 和 百度网盘 (密码: lwy6)来下载，数据集中包含原始数据以及处理过的数据。

在下表我们列出了9个数据集：

任务	数据集
Unconditional	Image COCO Caption
	EMNLP2017 WMT News
	IMDB Movie Review
Translation	IWSLT2014 German-English
	WMT2014 English-German
Summarization	GigaWord
Dialog	Persona Chat
Attribute to Text	Amazon Electronic
Poem Generation	Chinese Classical Poetry Corpus

下载好的数据集需要放到 `dataset` 目录下面，和我们项目中的结构类似。

我们也支持用户在自己的数据集上训练模型，只需要按照下面三个步骤操作即可：

1. 在 dataset 目录下面创建一个新的目录用于放置用户自己的数据集，数据集要求每行包含一个文本序列，例如 dataset/YOUR_DATASET;

2. 创建一个YAML参数配置文件用于对自己数据集的超参数进行配置，YAML的文件名称应与数据集名称相同，例如 textbox/properties/dataset/YOUR_DATASET.yaml.

   如果你想对数据集进行分割，请在YAML文件中设置 split_strategy: "load_split"，具体可以参考 COCO yaml 或者 IWSLT14_DE_EN yaml.

   如果你想按照比例自动对数据集进行划分，请在YAML文件中设置 split_strategy: "by_ratio" 和 split_ratio 这两个参数，具体可以参考 IMDB yaml.

3. 对于无条件文本生成，如果你设置了 "by_ratio" ，请将数据集命名为 corpus.txt ，如果你设置了 "load_split" ，请将数据集命名为 train.txt, valid.txt, dev.txt 。

   对于有条件文本生成，请将数据集命名为 train.[xx/yy], valid.[xx/yy], dev.[xx/yy] ， xx 或者 yy 是源文件或目标文件的后缀，应与YAML文件中的 source_suffix 和 target_suffix 保持一致。

实验结果

我们实现了多个文本生成模型，并在有条件文本生成和无条件文本生成任务上对他们的结果进行了比较。我们也展示了一些生成实例，更多的实例可以生成实例找到。

在前期实验中，我们的TextBox得到了以下结果。然而，这些算法是根据我们的理解和经验来实现和调整的，这可能没有达到它们的最佳性能。如果你能在某个具体算法上得到更好的结果，请告知我们。验证结果后，我们会更新该表。

无条件文本生成

Image COCO Caption

测试集负对数似然 (NLL), BLEU and Self-BLEU (SBLEU)指标结果展示：

Model	NLL	BLEU-2	BLEU-3	BLEU-4	BLEU-5	SBLEU-2	SBLEU-3	SBLEU-4	SBLEU-5
RNNVAE	33.02	80.46	51.5	25.89	11.55	89.18	61.58	32.69	14.03
CNNVAE	36.61	0.63	0.27	0.28	0.29	3.10	0.28	0.29	0.30
HybridVAE	56.44	31.96	3.75	1.61	1.76	77.79	26.77	5.71	2.49
SeqGAN	30.56	80.15	49.88	24.95	11.10	84.45	54.26	27.42	11.87
TextGAN	32.46	77.47	45.74	21.57	9.18	82.93	51.34	24.41	10.01
RankGAN	31.07	77.36	45.05	21.46	9.41	83.13	50.62	23.79	10.08
MaliGAN	31.50	80.08	49.52	24.03	10.36	84.85	55.32	28.28	12.09
LeakGAN	25.11	93.49	82.03	62.59	42.06	89.73	64.57	35.60	14.98
MaskGAN	95.93	58.07	21.22	5.07	1.88	76.10	43.41	20.06	9.37
GPT-2	26.82	75.51	58.87	38.22	21.66	92.78	75.47	51.74	32.39

部分生成实例展示：

模型	实例
RNNVAE	people playing polo to eat in the woods .
LeakGAN	a man is standing near a horse on a lush green grassy field .
GPT-2	cit a large zebra lays down on the ground.

EMNLP2017 WMT News

测试集NLL, BLEU, SBLEU指标结果展示：

Model	NLL	BLEU-2	BLEU-3	BLEU-4	BLEU-5	SBLEU-2	SBLEU-3	SBLEU-4	SBLEU-5
RNNVAE	142.23	58.81	19.70	5.57	2.01	72.79	27.04	7.85	2.73
CNNVAE	164.79	0.82	0.17	0.18	0.18	2.78	0.19	0.19	0.20
HybridVAE	177.75	29.58	1.62	0.47	0.49	59.85	10.3	1.43	1.10
SeqGAN	142.22	63.90	20.89	5.64	1.81	70.97	25.56	7.05	2.18
TextGAN	140.90	60.37	18.86	4.82	1.52	68.32	23.24	6.10	1.84
RankGAN	142.27	61.28	19.81	5.58	1.82	67.71	23.15	6.63	2.09
MaliGAN	149.93	45.00	12.69	3.16	1.17	65.10	20.55	5.41	1.91
LeakGAN	162.70	76.61	39.14	15.84	6.08	85.04	54.70	29.35	14.63
MaskGAN	303.00	63.08	21.14	5.40	1.80	83.92	47.79	19.96	7.51
GPT-2	88.01	55.88	21.65	5.34	1.40	75.67	36.71	12.67	3.88

部分生成实例展示：

模型	实例
RNNVAE	lewis holds us in total because they have had a fighting opportunity to hold any bodies when companies on his assault .
LeakGAN	we ' re a frustration of area , then we do coming out and play stuff so that we can be able to be ready to find a team in a game , but I know how we ' re going to say it was a problem .
GPT-2	russ i'm trying to build a house that my kids can live in, too, and it's going to be a beautiful house.

IMDB Movie Review

测试集NLL, BLEU, SBLEU指标结果展示：

Model	NLL	BLEU-2	BLEU-3	BLEU-4	BLEU-5	SBLEU-2	SBLEU-3	SBLEU-4	SBLEU-5
RNNVAE	445.55	29.14	13.73	4.81	1.85	38.77	14.39	6.61	5.16
CNNVAE	552.09	1.88	0.11	0.11	0.11	3.08	0.13	0.13	0.13
HybridVAE	318.46	38.65	2.53	0.34	0.31	70.05	17.27	1.57	0.59
SeqGAN	547.09	66.33	26.89	6.80	1.79	72.48	35.48	11.60	3.31
TextGAN	488.37	63.95	25.82	6.81	1.51	72.11	30.56	8.20	1.96
RankGAN	518.10	58.08	23.71	6.84	1.67	69.93	31.68	11.12	3.78
MaliGAN	552.45	44.50	15.01	3.69	1.23	57.25	22.04	7.36	3.26
LeakGAN	499.57	78.93	58.96	32.58	12.65	92.91	79.21	60.10	39.79
MaskGAN	509.58	56.61	21.41	4.49	0.86	92.09	77.88	59.62	42.36
GPT-2	348.67	72.52	41.75	15.40	4.22	86.21	58.26	30.03	12.56

部分生成实例展示（最大长度 max_length设置为100 ）：

模型	实例
RNNVAE	best brilliant known plot , sound movie , although unfortunately but it also like . the almost five minutes i will have done its bad numbers . so not yet i found the difference from with
LeakGAN	i saw this film shortly when I storms of a few concentration one before it all time . It doesn t understand the fact that it is a very good example of a modern day , in the <|unk|> , I saw it . It is so bad it s a <|unk|> . the cast , the stars , who are given a little
GPT-2	be a very bad, low budget horror flick that is not worth watching and, in my humble opinion, not worth watching any time. the acting is atrocious, there are scenes that you could laugh at and the story, if you can call it that, was completely lacking in logic and

序列到序列（seq2seq）文本生成

GigaWord （摘要）

使用beam搜索在测试集上的ROUGE指标（beam搜索大小 beam_size 设置为5）:

Model	ROUGE-1	ROUGE-2	ROUGE-L	ROUGE-W
RNN with Attention	36.32	17.63	38.36	25.08
Transformer	36.21	17.64	38.10	24.89
BART	39.34	20.07	41.25	27.13
BERT2BERT	38.16	18.89	40.06	26.21
ProphetNet	38.49	18.41	39.84	26.12
T5	38.83	19.68	40.76	26.73

部分生成实例展示：

文章	japan 's nec corp. and computer corp. of the united states said wednesday they had agreed to join forces in supercomputer sales .
真实摘要	nec in computer sales tie-up
RNN with Attention	nec computer corp .
Transformer	nec computer to join forces in chip sales
BART	nec computer corp.
BERT2BERT	nec computer form alliance for supercomputer sales
ProphetNet	nec computer to join forces in supercomputer sales
T5	nec computer to join forces in supercomputer sales

IWSLT2014 German-English（翻译）

测试集上的BLEU指标有三种解码策略：top-k采样、贪婪搜索和beam搜索（beam搜索大小 beam_size 设置为5）:

模型	解码策略	BLEU-2	BLEU-3	BLEU-4	BLEU
RNN with Attention	Top-k sampling	26.68	16.95	10.85	19.66
	Greedy search	33.74	23.03	15.79	26.23
	Beam search	35.68	24.94	17.42	28.23
Transformer	Top-k sampling	30.96	20.83	14.16	23.91
	Greedy search	35.48	24.76	17.41	28.10
	Beam search	36.88	26.10	18.54	29.49
BART	Beam search	29.02	19.58	13.48	22.42
BERT2BERT	Beam search	27.61	18.41	12.46	21.07

部分生成实例展示：

源语言（德语）	wissen sie , eines der großen < unk > beim reisen und eine der freuden bei der < unk > forschung ist , gemeinsam mit den menschen zu leben , die sich noch an die alten tage erinnern können . die ihre vergangenheit noch immer im wind spüren , sie auf vom regen < unk > steinen berühren , sie in den bitteren blättern der pflanzen schmecken .
真实目标语言（英语）	you know , one of the intense pleasures of travel and one of the delights of < unk > research is the opportunity to live amongst those who have not forgotten the old ways , who still feel their past in the wind , touch it in stones < unk > by rain , taste it in the bitter leaves of plants .
RNN with Attention	you know , one of the great < unk > trips is a travel and one of the friends in the world & apos ; s investigation is located on the old days that you can remember the past day , you & apos ; re < unk > to the rain in the < unk > chamber of plants .
Transformer	you know , one of the great < unk > about travel , and one of the pleasure in the < unk > research is to live with people who remember the old days , and they still remember the wind in the wind , but they & apos ; re touching the < unk > .

Persona Chat （对话）

使用beam搜索在测试集上的BLEU、Distinct指标（beam搜索大小 beam_size 设置为5）:

Model	Distinct-1	Distinct-2	BLEU-1	BLEU-2	BLEU-3	BLEU-4
RNN with Attention	0.24	0.72	17.51	4.65	2.11	1.47
Transformer	0.38	2.28	17.29	4.85	2.32	1.65
HRED	0.22	0.63	17.29	4.72	2.20	1.60

Amazon Electronic （属性文本生成）

使用beam搜索在测试集上的BLEU、Distinct指标（beam搜索大小 beam_size 设置为5）:

Model	Distinct-1	Distinct-2	BLEU-1	BLEU-2	BLEU-3	BLEU-4
Context2Seq	0.07	0.39	17.21	2.80	0.83	0.43
Attr2Seq	0.14	2.81	17.14	2.81	0.87	0.48

TextBox重要发布

发行版本	日期	特点
v0.2.1	15/04/2021	TextBox
v0.1.5	2021/01/11	Basic TextBox

贡献

如果您遇到错误或有任何建议，请通过 filing an issue.

我们欢迎关于修复错误、添加新特性的任何贡献。

我们希望所有的贡献者先在issue中提出问题，然后再提PR。

我们感谢@LucasTsui0725实现了HRED模型，@Richar-Du实现了CVAE模型。

引用

如果你觉得TextBox对你的科研工作有帮助，请引用我们的 论文:

@article{textbox,
    title={TextBox: A Unified, Modularized, and Extensible Framework for Text Generation},
    author={Junyi Li, Tianyi Tang, Gaole He, Jinhao Jiang, Xiaoxuan Hu, Puzhao Xie, Wayne Xin Zhao, Ji-Rong Wen},
    year={2021},
    journal={arXiv preprint arXiv:2101.02046}
}

项目团队

TextBox 由 AI Box 团队成员开发。

许可

TextBox 使用 MIT License。