Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Real-time collaboration for Jupyter Notebooks, Linux Terminals, LaTeX, VS Code, R IDE, and more,
all in one place. Commercial Alternative to JupyterHub.
Path: blob/main/examples/summarization-tf.ipynb
Views: 2535
If you're opening this Notebook on colab, you will probably need to install 🤗 Transformers and 🤗 Datasets as well as other dependencies. Uncomment the following cell and run it. Note the rouge-score and nltk dependencies - even if you've used 🤗 Transformers before, you may not have these installed!
If you're opening this notebook locally, make sure your environment has an install from the last version of those libraries.
To be able to share your model with the community and generate results like the one shown in the picture below via the inference API, there are a few more steps to follow.
First you have to store your authentication token from the Hugging Face website (sign up here if you haven't already!) then run the following cell and input your token:
Then you need to install Git-LFS and setup Git if you haven't already. Uncomment the following instructions and adapt with your name and email:
Make sure your version of Transformers is at least 4.16.0 since some of the functionality we use was introduced in that version:
You can find a script version of this notebook to fine-tune your model in a distributed fashion using multiple GPUs or TPUs here.
We also quickly upload some telemetry - this tells us which examples and software versions are getting used so we know where to prioritize our maintenance efforts. We don't collect (or care about) any personally identifiable information, but if you'd prefer not to be counted, feel free to skip this step or delete this cell entirely.
Fine-tuning a model on a summarization task
In this notebook, we will see how to fine-tune one of the 🤗 Transformers model for a summarization task. We will use the XSum dataset (for extreme summarization) which contains BBC articles accompanied with single-sentence summaries.
We will see how to easily load the dataset for this task using 🤗 Datasets and how to fine-tune a model on it using Keras.
Loading the dataset
We will use the 🤗 Datasets library to download the data and get the metric we need to use for evaluation (to compare our model to the benchmark). This can be easily done with the functions load_dataset and load_metric.
The dataset object itself is DatasetDict, which contains one key for the training, validation and test set:
To access an actual element, you need to select a split first, then give an index:
To get a sense of what the data looks like, the following function will show some examples picked randomly in the dataset.
The metric is an instance of datasets.Metric:
You can call its compute method with your predictions and labels, which need to be list of decoded strings:
Preprocessing the data
Before we can feed those texts to our model, we need to preprocess them. This is done by a 🤗 Transformers Tokenizer which will (as the name indicates) tokenize the inputs (including converting the tokens to their corresponding IDs in the pretrained vocabulary) and put it in a format the model expects, as well as generate the other inputs that the model requires.
To do all of this, we instantiate our tokenizer with the AutoTokenizer.from_pretrained method, which will ensure:
we get a tokenizer that corresponds to the model architecture we want to use,
we download the vocabulary used when pretraining this specific checkpoint.
That vocabulary will be cached, so it's not downloaded again the next time we run the cell.
By default, the call above will use one of the fast tokenizers (backed by Rust) from the 🤗 Tokenizers library.
You can directly call this tokenizer on one sentence or a pair of sentences:
Depending on the model you selected, you will see different keys in the dictionary returned by the cell above. They don't matter much for what we're doing here (just know they are required by the model we will instantiate later), you can learn more about them in this tutorial if you're interested.
Instead of one sentence, we can pass along a list of sentences:
To prepare the targets for our model, we need to tokenize them inside the as_target_tokenizer context manager. This will make sure the tokenizer uses the special tokens corresponding to the targets:
If you are using one of the five T5 checkpoints we have to prefix the inputs with "summarize:" (the model can also translate and it needs the prefix to know which task it has to perform).
We can then write the function that will preprocess our samples. We just feed them to the tokenizer with the argument truncation=True. This will ensure that an input longer that what the model selected can handle will be truncated to the maximum length accepted by the model. The padding will be dealt with later on (in a data collator) so we pad examples to the longest length in the batch and not the whole dataset.
This function works with one or several examples. In the case of several examples, the tokenizer will return a list of lists for each key:
To apply this function on all the pairs of sentences in our dataset, we just use the map method of our dataset object we created earlier. This will apply the function on all the elements of all the splits in dataset, so our training, validation and testing data will be preprocessed in one single command.
Even better, the results are automatically cached by the 🤗 Datasets library to avoid spending time on this step the next time you run your notebook. The 🤗 Datasets library is normally smart enough to detect when the function you pass to map has changed (and thus requires to not use the cache data). For instance, it will properly detect if you change the task in the first cell and rerun the notebook. 🤗 Datasets warns you when it uses cached files, you can pass load_from_cache_file=False in the call to map to not use the cached files and force the preprocessing to be applied again.
Note that we passed batched=True to encode the texts by batches together. This is to leverage the full benefit of the fast tokenizer we loaded earlier, which will use multi-threading to treat the texts in a batch concurrently.
Fine-tuning the model
Now that our data is ready, we can download the pretrained model and fine-tune it. Since our task is sequence-to-sequence (both the input and output are text sequences), we use the AutoModelForSeq2SeqLM class. Like with the tokenizer, the from_pretrained method will download and cache the model for us.
Note that we don't get a warning like in our classification example. This means we used all the weights of the pretrained model and there is no randomly initialized head in this case.
Next we set some parameters like the learning rate and the batch_sizeand customize the weight decay.
The last two arguments are to setup everything so we can push the model to the Hub at the end of training. Remove the two of them if you didn't follow the installation steps at the top of the notebook, otherwise you can change the value of push_to_hub_model_id to something you would prefer.
Then, we need a special kind of data collator, which will not only pad the inputs to the maximum length in the batch, but also the labels. Note that our data collators are designed to work for multiple frameworks, so ensure you set the return_tensors='np' argument to get NumPy arrays out - you don't want to accidentally get a load of torch.Tensor objects in the middle of your nice TF code! You could also use return_tensors='tf' to get TensorFlow tensors, but our TF dataset pipeline actually uses a NumPy loader internally, which is wrapped at the end with a tf.data.Dataset. As a result, np is usually more reliable and performant when you're using it!
We also want to compute ROUGE metrics, which will require us to generate text from our model. To speed things up, we can compile our generation loop with XLA. This results in a huge speedup - up to 100X! The downside of XLA generation, though, is that it doesn't like variable input shapes, because it needs to run a new compilation for each new input shape! To compensate for that, let's use pad_to_multiple_of for the dataset we use for text generation. This will reduce the number of unique input shapes a lot, meaning we can get the benefits of XLA generation with only a few compilations.
Next, we convert our datasets to tf.data.Dataset, which Keras understands natively. There are two ways to do this - we can use the slightly more low-level Dataset.to_tf_dataset() method, or we can use Model.prepare_tf_dataset(). The main difference between these two is that the Model method can inspect the model to determine which column names it can use as input, which means you don't need to specify them yourself. Make sure to specify the collator we just created as our collate_fn!
Now we initialize our loss and optimizer and compile the model. Note that most Transformers models compute loss internally - we can train on this as our loss value simply by not specifying a loss when we compile().
Now we can train our model. We can also add a few optional callbacks here, which you can remove if they aren't useful to you. In no particular order, these are:
PushToHubCallback will sync up our model with the Hub - this allows us to resume training from other machines, share the model after training is finished, and even test the model's inference quality midway through training!
TensorBoard is a built-in Keras callback that logs TensorBoard metrics.
KerasMetricCallback is a callback for computing advanced metrics. There are a number of common metrics in NLP like ROUGE which are hard to fit into your compiled training loop because they depend on decoding predictions and labels back to strings with the tokenizer, and calling arbitrary Python functions to compute the metric. The KerasMetricCallback will wrap a metric function, outputting metrics as training progresses.
If this is the first time you've seen KerasMetricCallback, it's worth explaining what exactly is going on here. The callback takes two main arguments - a metric_fn and an eval_dataset. It then iterates over the eval_dataset and collects the model's outputs for each sample, before passing the list of predictions and the associated list of labels to the user-defined metric_fn. If the predict_with_generate argument is True, then it will call model.generate() for each input sample instead of model.predict() - this is useful for metrics that expect generated text from the model, like ROUGE.
This callback allows complex metrics to be computed each epoch that would not function as a standard Keras Metric. Metric values are printed each epoch, and can be used by other callbacks like TensorBoard or EarlyStopping.
And now we can try training our model. By default, we only do a single epoch of training here, as the inputs are very long, which means training is quite slow. However, you may wish to experiment with larger pre-trained models and longer training runs if you want to maximize the quality of your summaries.
25505/25505 [==============================] - 4900s 192ms/step - loss: 2.7178 - val_loss: 2.3966 - rouge1: 29.5589 - rouge2: 8.6381 - rougeL: 23.3911 - rougeLsum: 23.3926 - gen_len: 18.8211
If you used the callback above, you can now share this model with all your friends, family, favorite pets: they can all load it with the identifier "your-username/the-name-you-picked" so for instance:
Inference
Now we've trained our model, let's see how we could load it and use it to summarize text in future! First, let's load it from the hub. This means we can resume the code from here without needing to rerun everything above every time.
Now let's try tokenizing a document from the training set. Don't forget to add 'summarize:' at the start if you're using a T5 model.
Not bad for a single epoch of training! Of course, the flood warning isn't much use to them after they've been flooded, but the summary correctly identified flooding in Dumfries and the Nith as the key event.
Using XLA in inference
If you just want to generate a few summaries, the code above is all you need. However, generation can be much faster if you use XLA, and if you want to generate data in bulk, you should probably use it! If you're using XLA, though, remember that you'll need to do a new XLA compilation for every input size you pass to the model. This means that you should keep your batch size constant, and consider padding inputs to the same length, or using pad_to_multiple_of in your tokenizer to reduce the number of different input shapes you pass. Let's show an example of that:
When using XLA generation, you'll notice that the first call to generate with a new input shape takes a long time because XLA has to compile your function, but subsequent calls are extremely quick. Also, XLA always generates to the maximum length, which can lead to a lot of padding tokens in your output! These are easy to remove, however:
Pipeline API
The pipeline API offers a convenient shortcut for all of this, but doesn't (yet!) support XLA generation:
Remember that if we're using a T5 model then we appended "summarize: " to the start of our input above. Don't forget to do that when you're getting summaries for new texts!
Easy!