[Paper Review] TextBox: A Unified, Modularized, and Extensible Framework for Text Generation

Junyi Li, Tianyi Tang+8 more|arXiv (Cornell University)|Jan 6, 2021

Topic Modeling43 references4 citations

TL;DR

TextBox is a unified, modular, and extensible PyTorch-based framework for text generation that supports 21 models across VAE, GAN, RNN, Transformer, and pretrained language models on 9 benchmark datasets. It enables easy model comparison and extension through configurable, reusable components, with state-of-the-art performance on summarization tasks (e.g., BART achieving ROUGE-1 39.34).

ABSTRACT

In this paper, we release an open-source library, called TextBox, to provide a unified, modularized, and extensible text generation framework. TextBox aims to support a broad set of text generation tasks and models. In our library, we implement 21 text generation models on 9 benchmark datasets, covering the categories of VAE, GAN, and pretrained language models. Meanwhile, our library maintains sufficient modularity and extensibility by properly decomposing the model architecture, inference, and learning process into highly reusable modules, which allows users to easily incorporate new models into our framework. The above features make TextBox specially suitable for researchers and practitioners to quickly reproduce baseline models and develop new models. TextBox is implemented based on PyTorch, and released under Apache License 2.0 at https://github.com/RUCAIBox/TextBox.

Motivation & Objective

• To address the lack of a unified, modular, and extensible framework for text generation that supports diverse models and tasks.
• To improve reproducibility and standardization of text generation research by providing consistent implementation and evaluation protocols.
• To reduce the time and effort required to implement and compare baseline models across different text generation tasks.
• To facilitate the development of new models by enabling plug-and-play integration of reusable components.
• To support a broad range of text generation tasks, including text summarization, machine translation, and dialogue, with standardized evaluation metrics.

Proposed method

• The framework is built on PyTorch and organizes model components into modular units: data, model, evaluation, and common components.
• It uses a configurable pipeline where users define experiments via configuration files or command-line arguments, enabling flexible model and hyperparameter setup.
• The data module supports both single-sequence (unconditional) and paired-sequence (conditional) tasks through custom Dataset and DataLoader classes.
• Model components are decoupled into reusable modules such as RNN/Transformer encoders, decoders, attention mechanisms, and pretrained language models.
• Evaluation is standardized with support for metrics like BLEU, ROUGE, perplexity, and negative log-likelihood across all models and datasets.
• The framework ensures compatibility with PyTorch APIs, allowing seamless integration of user-defined modules and functions.

Experimental results

Research questions

• RQ1 How can a unified framework improve reproducibility and reduce implementation overhead in text generation research?
• RQ2 To what extent can a modular and extensible architecture support diverse text generation models across multiple tasks?
• RQ3 Can standardized evaluation protocols across multiple models and datasets improve fair comparison and benchmarking?
• RQ4 How does the modularity of TextBox enable rapid prototyping and integration of new models or components?
• RQ5 What performance gains can be achieved using state-of-the-art pretrained models like BART and T5 on text summarization tasks?

Key findings

• BART achieved the highest ROUGE-1 score of 39.34 on the GigaWord text summarization dataset, outperforming RNN and Transformer baselines.
• Pretrained models such as BART, T5, and ProphetNet significantly outperformed non-pretrained models like RNN and Transformer on text summarization, with ROUGE-L scores ranging from 39.84 to 41.25.
• Beam search strategy yielded better performance than greedy or top-k decoding in machine translation, with improved BLEU scores across IWSLT2014 German-to-English translation.
• The framework successfully integrated 21 diverse text generation models, including VAE, GAN, RNN, Transformer, and PLM-based models, across 9 benchmark datasets.
• The modular design enabled seamless configuration and comparison of models through configuration files, reducing implementation complexity.
• The library demonstrated strong extensibility, supporting future integration of new models, inputs (e.g., graphs, tables), and distributed training.