"妙笔"生花：一个易用、高效的文本生成开源库

文本生成作为近年来自然语言处理中的热门领域，受到了学界和工业界的广泛关注。随着研究的不断深入，文本生成领域下的子任务和相应的模型越来越丰富，一些优秀的开源框架也纷纷涌现。

现有的开源文本生成库大致可分为两类，一类倾向于模块化驱动，提供各个组件供研究者搭建不同的模型，但缺点在于许多特殊的baseline无法通过模块完成；另一类倾向于baseline驱动，方便用户快速调用baseline，但缺点在于整体结构不够统一，可扩展性有限。

因此，为了兼顾两者的优势，我们（中国人民大学AI BOX团队）推出了新的开源文本生成库TextBox（妙笔），在保证组件模块化的同时，提供了许多较新baseline的接口，使得研究者一方面可以基于模块搭建自己的模型，另一方面也可以快速调用已有的baseline完成实验。目前相关代码和相应论文都已上线。

下面让我们来看看到底有什么厉害的功能吧！

• 16种模型：覆盖无条件生成和有条件生成两个主流任务，后者主要包括机器翻译和自动摘要任务。
• 6个数据集：提供已处理好的数据集，也可以使用我们的库来处理原始数据。
• 多种评测方式：支持主流的自动评测指标，包括通用指标和特定任务指标。
• 易上手、易扩展：新手也可快速上手，一键复现，搭积木式的灵活开发，想要的模型统统安排！

论文地址：https://arxiv.org/pdf/2101.02046.pdf

项目GitHub地址：https://github.com/RUCAIBox/TextBox

核心特色

TextBox框架具有三大核心特色。

基于PyTorch的统一模块化框架

TextBox是基于深度学习框架PyTorch开发的文本生成库。为了统一开发多种文本生成模型并简化其开发难度，TextBox的设计是高度模块化的。它将不同的模型解耦成了一组可高度复用的模块，如数据模块、模型模块、评估模块和公共功能模块。因此，在本框架中，通过简单而灵活的配置，就可以根据内置的评估方法来比较不同的文本生成模型。同时，还可以像搭积木一样插入或者交换模块，在高度概念化的层面上开发新的文本生成模型。

全面的基准模型、数据集和评估。

除了模块化的设计之外，TextBox还包含了广泛的基准文本生成模型，其涵盖了变分自编码器（VAE）、生成式对抗网络（GAN）、循环神经网络（RNN），基于Transformer的模型，以及预训练语言模型（PLM）等类别。TextBox通过配置文件或命令行提供灵活的参数指定来运行、比较和测试这些传统的基准模型和目前最先进的模型。为了构建一个可重用的基准，TextBox整合了许多常用的数据集，可以针对不同的文本生成任务进行评估。例如，我们采用在GigaWord数据集上执行文本摘要任务，以及在WMT 2014 English to German数据集上执行机器翻译任务。TextBox支持一系列广泛使用的评估算法，用于测试和比较各类文本生成模型，例如评估语言模型的perplexity和neg-likelihood，评估翻译模型的BLEU，以及评估摘要生成模型的ROUGE。

高度灵活及拓展性强的框架

TextBox在文本生成模型部分提供了多种常用的函数与模块的接口。例如RNN encoder-decoder, Transformer encoder-decoder以及各种预训练语言模型。具有不同需求和专业背景的用户可以很方便地选择不同的接口来构建自己的模型，并且能够基于框架提供的统一评估标准来对模型进行评估。更重要的是，TextBox提供的接口可以与PyTorch接口完全兼容，允许无缝集成用户自定义的模块，使得用户能够根据需要集成任何外部组件。

快速上手

TextBox提供了较好的封装，我们将提供简单的tutorial大致介绍如何快速上手。

安装

目前TextBox支持源代码安装和Pip安装，用户可以通过下面的简单命令进行安装使用：

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

或

pip install textbox

一键运行

我们在Github提供了一系列一键运行的脚本，包括训练、单独测试：

python run_textbox.py --model=RNN --dataset=COCO --task_type=unconditional
python run_textbox.py --model=RNN --dataset=COCO --task_type=unconditional –-load_experiment=<ckpt_file> --test_only=true

通过上述命令就可以直接在COCO数据集上基于RNN模型进行无条件生成任务的训练和单独测试。

如果你想改变参数，比如rnn_type、max_vocab_size等，可以按需在命令行设置额外的参数：

python run_textbox.py --model=RNN --dataset=COCO --task_type=unconditional \
                      --rnn_type=lstm --max_vocab_size=4000

如果你想运行别的数据集、模型和任务，可以使用

python run_textbox.py --model=<model_name> --dataset=<dataset_name> --task_type=<task_type>

如果你使用pip安装，只需新建一个py文件，添加以下代码即可实现模型的训练和测试

from textbox.quick_start import run_textbox

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

使用预训练语言模型

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

• 从huggingface提供的模型源(https://huggingface.co/gpt2/tree/main)中下载GPT-2模型，包括config.json, merges.txt, pytorch_model.bin, tokenizer.json和vocab.json五个文件，将其放在同一个文件夹下，路径为pretrained_model/gpt2，与textbox文件夹同级。
• 下载好模型之后，直接通过脚本运行：

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

未来展望

至此，“妙笔”文本生成库的自我介绍就到此结束啦。未来我们将继续努力，开发更多新的模型，维护并保持版本稳定。有好的建议或意见也欢迎关注（点击右边链接关注➡️）AI Box公众号后台留言或者在Github的issues中提问，期待大家的试用体验和宝贵意见。

更多推荐

伯乐：一个易用、强大的 Pytorch 推荐系统开源库