Code for the paper "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer"
Code for the paper "Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer"
T5X is the new and improved implementation of T5 (and more) in JAX and Flax.
T5 on Tensorflow with MeshTF is no longer actively developed. If you are new
to T5, we recommend starting with T5X.
The t5 library serves primarily as code for reproducing the experiments in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer][paper]. In the paper, we demonstrate how to achieve state-of-the-art results on multiple NLP tasks using a text-to-text transformer pre-trained on a large text corpus.
The bulk of the code in this repository is used for loading, preprocessing, mixing, and evaluating datasets. It also provides a way to fine-tune the pre-trained models released alongside the publication.
The t5 library can be used for future model development by providing useful modules for training and fine-tuning (potentially huge) models on mixtures of text-to-text tasks.
t5.data is a package for defining Task objects that provide tf.data.Datasets.
Each Task is made up of:
Additionally, you may optionally provide:
The data source can be an arbitrary function that provides a tf.data.Dataset, but we also provide simpler wrappers for datasets available in [TensorFlow Datasets (TFDS)][tfds] (a TfdsTask) or stored as text files with one example per line (a TextLineTask).
The text preprocessor converts the examples in the source dataset into the appropriate format for a text-to-text model with fields for inputs and targets. For example, the predefined t5.data.preprocessors.translate preprocessor converts inputs in the form
{'de': 'Das ist gut.', 'en': 'That is good.'}
to the form
{'inputs': 'translate German to English: Das ist gut.', 'targets': 'That is good.'}
In addition to text preprocessing, you can also use one or more token preprocessors to modify the inputs post-tokenization. We implemented our unsupervised pre-training objectives using these token preprocessors.
We provide many predefined preprocessors in t5.data.preprocessors, but you may also define your own.
The SentencePiece model is used to tokenize the input strings and decode the output tokens. You can create your own model with the google/sentencepiece library, or use our default one at t5.data.DEFAULT_SPM_PATH. If you create your own, you must use the flags --pad_id=0 --eos_id=1 --unk_id=2 --bos_id=-1 with spm_train to be compatible with our model code.
The metric function returns a score given the target and prediction from the model. You may also define a postprocess function to convert the target and prediction text to another format before calling the metric. We provide some predefined metrics in t5.evaluation.metrics.
Finally, t5.data contains a Mixture class that can be instantiated to combine multiple Task datasets for multi-task training using various functions for specifying the mixture rates.
t5.evaluation contains two core components:
t5.models contains shims for connecting T5 Tasks and Mixtures to a model implementation for training, evaluation, and inference.
Currently there are two shims available: One for the [Mesh TensorFlow Transformer][mtft] that we used in our paper and another for the Hugging Face Transformers library.
The Hugging Face API is currently experimental and subject to change, but provides a simple and easy way to load, fine-tune, and evaluate our pre-trained models using PyTorch on a single GPU.
If you want to use our largest models on TPUs and/or reproduce the results in our paper, you should use the MtfModel API and the t5_mesh_transformer binary.
If you are interested fine-tuning our models on a GPU in PyTorch, you should try the HfPyTorchModel API.
Since the HfPyTorchModel is experimental, the remainder of this README assumes usage of the MtfModel and its associated binary.
A usage example of HfPyTorchModel is available here.
The easiest way to try out T5 is with a free TPU in our Colab Tutorial.
Below we provide examples for how to pre-train, fine-tune, evaluate, and decode from a model from the command-line with our codebase. You can use these instructions to reproduce our results, fine-tune one of our released checkpoints with your own data and/or hyperparameters, or pre-train a model from scratch.
You may either use a new or pre-existing Task, or you may load examples from a preprocessed TSV file.
TaskDepending on your data source (see above), you will need to prepare your data appropriately.
Task
If using a vanilla task, just make sure any file(s) loaded by your dataset_fn are accessible to the TPU (i.e., are in a GCS bucket), and you should be good to go!
TfdsTask
Most of our predefined Tasks use [TensorFlow Datasets (TFDS)][tfds] as their data source. When you run our training binary (see instructions below) with a TfdsTask, the dataset will automatically be downloaded and prepared on its first use. After preparation is complete, the dataset is cached to your local storage to avoid this overhead in future runs. If working in the cloud, we recommend you set the --t5_tfds_data_dir flag to point to a persistent storage location, such as a [GCS bucket][gcs]. This is a requirement when training on TPU.
The [C4][c4] dataset we created for unsupervised pre-training is available in TensorFlow Datasets, but it requires a significant amount of bandwidth for downloading the raw [Common Crawl][cc] scrapes (~7 TB) and compute for its preparation (~335 CPU-days). We suggest you take advantage of the [Apache Beam][beam] support in TFDS, which enables distributed preprocessing of the dataset and can be run on [Google Cloud Dataflow][gcd]. With 500 workers, the job should complete in ~16 hours.
After defining MY_PROJECT and MY_BUCKET appropriately, you can build the dataset in DataFlow from GCP using the following commands:
pip install tfds-nightly[c4]
echo 'tfds-nightly[c4]' > /tmp/beam_requirements.txt
python -m tensorflow_datasets.scripts.download_and_prepare \
--datasets=c4/en \
--data_dir=gs://$MY_BUCKET/tensorflow_datasets \
--beam_pipeline_options="project=$MY_PROJECT,job_name=c4,staging_location=gs://$MY_BUCKET/binaries,temp_location=gs://$MY_BUCKET/temp,runner=DataflowRunner,requirements_file=/tmp/beam_requirements.txt,experiments=shuffle_mode=service,region=$MY_REGION"
Read more in the [TFDS Beam instructions][tfds_beam].
TextLineTask
A TextLineTask is useful when your data source is a text file (or files) with one example per line. You can then use a text preprocessor to convert each line into a dictionary of inputs and targets.
Make sure your files are accessible to the TPU (i.e., are in a GCS bucket), and you should be good to go!
Instead of defining a new Task, you may use a TSV file (or files) directly as your dataset where each line is formatted as <input>\t<target>.
However, there are a couple of caveats:
If you need any of these features, you must define a new Task, TfdsTask, or TextLineTask.
Similar to the above cases, your TSV file(s) must be accessible to the TPU (i.e., are in a GCS bucket).
To install the T5 package, simply run:
pip install t5[gcp]
You will first need to launch a Virtual Machine (VM) on Google Cloud. Details about launching the VM can be found at the Google Cloud Documentation.
In order to run training or eval on Cloud TPUs, you must set up the following variables based on your project, zone and GCS bucket appropriately. Please refer to the Cloud TPU Quickstart guide for more details.
export PROJECT=your_project_name
export ZONE=your_project_zone
export BUCKET=gs://yourbucket/
export TPU_NAME=t5-tpu
export TPU_SIZE=v3-8
export DATA_DIR="${BUCKET}/your_data_dir"
export MODEL_DIR="${BUCKET}/your_model_dir"
Please use the following command to create a TPU device in the Cloud VM.
ctpu up --name=$TPU_NAME --project=$PROJECT --zone=$ZONE --tpu-size=$TPU_SIZE \
--tpu-only --noconf
In the command below, we train a model on the GLUE Benchmark MRPC task from scratch. You can change the MIXTURE_NAME gin parameter to use any of the tasks or mixtures provided in our package.
t5_mesh_transformer \
--tpu="${TPU_NAME}" \
--gcp_project="${PROJECT}" \
--tpu_zone="${ZONE}" \
--model_dir="${MODEL_DIR}" \
--t5_tfds_data_dir="${DATA_DIR}" \
--gin_file="dataset.gin" \
--gin_file="models/bi_v1.gin" \
--gin_param="utils.tpu_mesh_shape.model_parallelism = 1" \
--gin_param="utils.tpu_mesh_shape.tpu_topology = '${TPU_SIZE}'" \
--gin_param="MIXTURE_NAME = 'glue_mrpc_v002'"
The full list of tasks and mixtures can be obtained by running:
python -c "import t5; print(t5.data.MixtureRegistry.names())"
You may also define additional tasks and mixtures in a new file and import it using the --module_import flag.
Alternatively, you could train with a TSV file where each line is formatted as <input>\t<target> (see above).
In order to fine-tune one of our pre-trained models, you need to pass the operative config of the pre-trained model to the training script. The operative config should be passed in as a gin_file flag. It specifies the model architecture and other hyperparameters. In addition, you need to specify the mixture to fine-tune on. For example, to fine-tune the T5-small model on the glue_mrpc_v002 mixture, please run:
t5_mesh_transformer \
--tpu="${TPU_NAME}" \
--gcp_project="${PROJECT}" \
--tpu_zone="${ZONE}" \
--model_dir="${MODEL_DIR}" \
--t5_tfds_data_dir="${DATA_DIR}" \
--gin_file="dataset.gin" \
--gin_param="utils.tpu_mesh_shape.model_parallelism = 1"
No open issues yet, or sync has not completed.