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/question_answering-tf.ipynb
Views: 2535
If you're opening this Notebook on colab, you will probably need to install 🤗 Transformers and 🤗 Datasets. Uncomment the following cell and run it.
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 uncomment 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 the functionality 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 question-answering task
In this notebook, we will see how to fine-tune one of the 🤗 Transformers model to a question answering task, which is the task of extracting the answer to a question from a given context. We will see how to easily load a dataset for these kinds of tasks and use Keras to fine-tune a model on it. Note that this model does not generate new text! Instead, it selects a span of the input passage as the answer.
This notebook is built to run on any question answering task with the same format as SQUAD (version 1 or 2), with any model checkpoint from the Model Hub as long as that model has a version with a token classification head and a fast tokenizer (check on this table if this is the case). It might, however, need some small adjustments if you decide to use a different dataset than the one used here. Depending on your model and the GPU you are using, you might need to adjust the batch size to avoid out-of-memory errors. Set those three parameters, then the rest of the notebook should run smoothly:
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.
For our example here, we'll use the SQUAD dataset. The notebook should work with any question answering dataset in the 🤗 Datasets library. If you're using your own dataset in a JSON or CSV file (see the Datasets documentation on how to load them), it might need some adjustments to the column names.
The datasets object itself is DatasetDict, which contains one key for the training, validation and test set.
We can see the training, validation and test sets all have a column for the context, the question and the answers to those questions.
To access an actual element, you need to select a split first, then give an index:
We can see the answers are indicated by their start position in the text (here at character 515) and their full text, which is a substring of the context as we mentioned above.
To get a sense of what the data looks like, the following function will show some examples picked randomly from the dataset and decoded back to strings.
Preprocessing the training 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 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.
The following assertion ensures that our tokenizer is a fast tokenizer (backed by Rust) from the 🤗 Tokenizers library. Those fast tokenizers are available for almost all models, and we will need some of the special features they have for our preprocessing.
You can check which type of models have a fast tokenizer available and which don't in the big table of models.
You can directly call this tokenizer on two sentences (one for the answer, one for the context):
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.
Now one specific thing for the preprocessing in question answering is how to deal with very long documents. We usually truncate them in other tasks, when they are longer than the model maximum sentence length, but here, removing part of the the context might result in losing the answer we are looking for. To deal with this, we will allow one (long) example in our dataset to give several input features, each of length shorter than the maximum length of the model (or the one we set as a hyper-parameter). Also, just in case the answer lies at the point we split a long context, we allow some overlap between the features we generate controlled by the hyper-parameter doc_stride:
Let's find one long example in our dataset:
Without any truncation, we get the following length for the input IDs:
Now, if we just truncate, we will lose information (and possibly the answer to our question):
Note that we never want to truncate the question, only the context, and so we use the only_second truncation method. Our tokenizer can automatically return a list of features capped by a certain maximum length, with the overlap we talked about above, we just have to tell it to do so with return_overflowing_tokens=True and by passing the stride:
Now we don't have one list of input_ids, but several:
And if we decode them, we can see the overlap:
It's going to take some work to properly label the answers here: we need to find in which of those features the answer actually is, and where exactly in that feature. The models we will use require the start and end positions of these answers in the tokens, so we will also need to to map parts of the original context to some tokens. Thankfully, the tokenizer we're using can help us with that by returning an offset_mapping:
This gives the corresponding start and end character in the original text for each token in our input IDs. The very first token ([CLS]) has (0, 0) because it doesn't correspond to any part of the question/answer, then the second token is the same as the characters 0 to 3 of the question:
So we can use this mapping to find the position of the start and end tokens of our answer in a given feature. We just have to distinguish which parts of the offsets correspond to the question and which part correspond to the context, this is where the sequence_ids method of our tokenized_example can be useful:
It returns None for the special tokens, then 0 or 1 depending on whether the corresponding token comes from the first sentence past (the question) or the second (the context). Now with all of this, we can find the first and last token of the answer in one of our input feature (or if the answer is not in this feature):
And we can double check that it is indeed the correct answer:
For this notebook to work with any kind of model, we need to account for the special case where the model expects padding on the left (in which case we switch the order of the question and the context):
Now let's put everything together in one function we will apply to our training set. In the case of impossible answers (the answer is in another feature given by an example with a long context), we set the cls index for both the start and end position. We could also simply discard those examples from the training set if the flag allow_impossible_answers is False. Since the preprocessing is already complex enough as it is, we've kept is simple for this part.
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 sentences (or pairs of sentences) in our dataset, we just use the map method of the 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. Since our preprocessing changes the number of samples, we need to remove the old columns when applying it.
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 for training, we can download the pretrained model and fine-tune it. Since our task is question answering, we use the TFAutoModelForQuestionAnswering class. Like with the tokenizer, the from_pretrained method will download and cache the model for us:
The warning is telling us we are throwing away some weights (the vocab_transform and vocab_layer_norm layers) and randomly initializing some other (the pre_classifier and classifier layers). This is absolutely normal in this case, because we are removing the head used to pretrain the model on a masked language modeling objective and replacing it with a new head for which we don't have pretrained weights, so the library warns us we should fine-tune this model before using it for inference, which is exactly what we are going to do.
To train our model, we will need to define a few more things. The first 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.
We also tweak the learning rate, use the batch_size defined at the top of the notebook and customize the number of epochs for training, as well as the weight decay.
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. It also supplies a default data collator that will work fine for us, as our samples are already padded to the same length and ready to go.
Next, we can create an optimizer and specify a loss function. The create_optimizer function gives us a very solid AdamW optimizer with weight decay and a learning rate schedule, but it needs us to compute the number of training steps to build that schedule.
Note that most Transformers models compute loss internally, so we actually don't have to specify anything there! You can of course set your own loss function if you want, but by default our models will choose the 'obvious' loss that matches their task, such as cross-entropy in the case of language modelling. The built-in loss will also correctly handle things like masking the loss on padding tokens, or unlabelled tokens in the case of masked language modelling, so we recommend using it unless you're an advanced user!
In addition, because the outputs and loss for this model class are quite straightforward, we can use built-in Keras metrics - these are liable to misbehave in other contexts (for example, they don't know about the masking in masked language modelling) but work well here.
We can also use jit_compile to compile the model with XLA. In other cases, we should be careful about that - if our inputs might have variable sequence lengths, we may end up having to do a new XLA compilation for each possible length, because XLA compilation expects a static input shape! In this notebook, however, we have padded all examples to exactly the same length. This makes it perfect for XLA, which will give us a nice performance boost.
We will evaluate our model and compute metrics in the next section (this is a very long operation, so we will only compute the evaluation loss during training). For now, let's just train our model. We can also add a callback to sync up our model with the Hub - this allows us to resume training from other machines and even test the model's inference quality midway through training! If you don't want to do this, simply remove the callbacks argument in the call to fit().
Evaluation
Evaluating our model will require a bit more work, as we will need to map the predictions of our model back to parts of the context. The model itself predicts logits for the start and end position of our answers: if we take a batch from our validation dataset, here is the output our model gives us:
The output of the model is a dict-like object that contains the loss (since we provided labels), the start and end logits. We won't need the loss for our predictions, let's have a look a the logits:
We have one logit for each feature and each token. The most obvious thing to predict an answer for each feature is to take the index for the maximum of the start logits as a start position and the index of the maximum of the end logits as an end position.
This will work great in a lot of cases, but what if this prediction gives us something impossible: the start position could be greater than the end position, or point to a span of text in the question instead of the answer. In that case, we might want to look at the second best prediction to see if it gives a possible answer and select that instead.
However, picking the second best answer is not as easy as picking the best one: is it the second best index in the start logits with the best index in the end logits? Or the best index in the start logits with the second best index in the end logits? And if that second best answer is not possible either, it gets even trickier for the third best answer.
To classify our answers, we will use the score obtained by adding the start and end logits. We won't try to order all the possible answers and limit ourselves to with a hyper-parameter we call n_best_size. We'll pick the best indices in the start and end logits and gather all the answers this predicts. After checking if each one is valid, we will sort them by their score and keep the best one. Here is how we would do this on the first feature in the batch:
And then we can sort the valid_answers according to their score and only keep the best one. The only point left is how to check a given span is inside the context (and not the question) and how to get back the text inside. To do this, we need to add two things to our validation features:
the ID of the example that generated the feature (since each example can generate several features, as seen before);
the offset mapping that will give us a map from token indices to character positions in the context.
That's why we will re-process the validation set with the following function, slightly different from prepare_train_features:
And like before, we can apply that function to our validation set easily:
And turn the dataset into a tf.data.Dataset as before.
Now we can grab the predictions for all features by using the model.predict method:
We can now refine the test we had before: since we set None in the offset mappings when it corresponds to a part of the question, it's easy to check if an answer is fully inside the context. We also eliminate very long answers from our considerations (with an hyper-parameter we can tune)
We can compare to the actual ground-truth answer:
Our model's most likely answer is correct!
As we mentioned in the code above, this was easy on the first feature because we knew it comes from the first example. For the other features, we will need a map between examples and their corresponding features. Also, since one example can give several features, we will need to gather together all the answers in all the features generated by a given example, then pick the best one. The following code builds a map from example index to its corresponding features indices:
We're almost ready for our post-processing function. The last bit to deal with is the impossible answer (when squad_v2 = True). The code above only keeps answers that are inside the context, we need to also grab the score for the impossible answer (which has start and end indices corresponding to the index of the CLS token). When one example gives several features, we have to predict the impossible answer when all the features give a high score to the impossible answer (since one feature could predict the impossible answer just because the answer isn't in the part of the context it has access too), which is why the score of the impossible answer for one example is the minimum of the scores for the impossible answer in each feature generated by the example.
We then predict the impossible answer when that score is greater than the score of the best non-impossible answer. All combined together, this gives us this post-processing function:
And we can apply our post-processing function to our raw predictions:
Then we can load the metric from the datasets library.
Then we can call compute on it. We just need to format predictions and labels a bit as it expects a list of dictionaries and not one big dictionary. In the case of squad_v2, we also have to set a no_answer_probability argument (which we set to 0.0 here as we have already set the answer to empty if we picked it).
If you ran the callback above, you can now share this model with all your friends, family or 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 answer questions 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 get some sample text and ask a question. Feel free to substitute your own text and question!
The outputs are logits, so let's use argmax to find the largest logit, which represents the model's best guess for the right answer.
Well, these are definitely tokens. Let's decode them back to text:
Pipeline API
An alternative way to quickly perform inference with any model on the hub is to use the Pipeline API, which abstracts away all the steps we did manually above. It will perform the preprocessing, forward pass and postprocessing all in a single object.
Let's showcase this for our trained model: