The transformers library provides a clean and well-documented interface for many popular transformer models. Not only it makes the source code easier to read and understand, it also provided a standardize way to interact with the model. You have seen in the previous post how to use a model such as DistilBERT for natural language processing tasks. In this post, you will learn how to fine-tune the model for your own purpose. This expands the use of the model from inference to training. Specifically, you will learn:
- How to prepare the dataset for training
- How to train a model using a helper library
Let’s get started.
Fine-Tuning DistilBERT for Question Answering
Photo by Lea Fabienne. Some rights reserved.
Overview
This post is divided into three parts; they are:
- Fine-tuning DistilBERT for Custom Q&A
- Dataset and Preprocessing
- Running the Training
Fine-tuning DistilBERT for Custom Q&A
The simplest way to use a model in the transformers library is to create a pipeline, which hides many details about how to interact with it.
One reason you may not want to create a pipeline but to set up the model individually is that you want to fine-tune the model on your own dataset. This is impossible with a pipeline because you need to examine the model’s raw output with a loss function, which is usually hidden from the pipeline.
Usually, the pre-trained model is created using a general-purpose dataset. However, it may not work well for a specific domain, especially if the language in the domain is significantly different from the general usage. This is where fine-tuning may be attempted.
The difficulty in fine-tuning is probably the availability of a good dataset. This is usually very expensive and time-consuming to create. For illustration purposes, in the following, we use a general-purpose and publicly available dataset called SQuAD (Stanford Question Answering Dataset).
Thanks to the highly generalized and cleaned design of the transformers library, fine-tuning the model is straightforward. Below is an example of how to fine-tune the model on the SQuAD dataset:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 |
rom transformers import DistilBertTokenizerFast, DistilBertForQuestionAnswering, \ Trainer, TrainingArguments from datasets import load_dataset
# Load the SQuAD dataset dataset = load_dataset(“squad”)
# Load tokenizer and model model_name = “distilbert-base-uncased” tokenizer = DistilBertTokenizerFast.from_pretrained(model_name) model = DistilBertForQuestionAnswering.from_pretrained(model_name)
# Tokenize the dataset def preprocess_function(examples): questions = [q.strip() for q in examples[“question”]] inputs = tokenizer( questions, examples[“context”], max_length=384, truncation=“only_second”, return_offsets_mapping=True, padding=“max_length”, )
offset_mapping = inputs.pop(“offset_mapping”) answers = examples[“answers”] start_positions = [] end_positions = []
for i, offsets in enumerate(offset_mapping): answer = answers[i] start_char = answer[“answer_start”][0] end_char = start_char + len(answer[“text”][0]) sequence_ids = inputs.sequence_ids(i)
# Find the start and end of the context context_start = sequence_ids.index(1) context_end = len(sequence_ids) – 1 – sequence_ids[::–1].index(1)
# If the answer is not fully inside the context, label it (0, 0) if offsets[context_start][0] > end_char or offsets[context_end][1] < start_char: start_positions.append(0) end_positions.append(0) else: # Otherwise find the start and end token positions idx = context_start while idx <= context_end and offsets[idx][0] <= start_char: idx += 1 start_positions.append(idx – 1)
idx = context_end while idx >= context_start and offsets[idx][1] >= end_char: idx -= 1 end_positions.append(idx + 1)
inputs[“start_positions”] = start_positions inputs[“end_positions”] = end_positions return inputs
# Apply preprocessing to the dataset tokenized_datasets = dataset.map(preprocess_function, batched=True, remove_columns=dataset[“train”].column_names) # Define training arguments training_args = TrainingArguments( output_dir=“./results”, evaluation_strategy=“epoch”, learning_rate=2e–5, per_device_train_batch_size=16, per_device_eval_batch_size=16, num_train_epochs=3, weight_decay=0.01, ) # Initialize Trainer trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_datasets[“train”], eval_dataset=tokenized_datasets[“validation”], tokenizer=tokenizer, ) # Train the model and save the results trainer.train() model.save_pretrained(“./fine-tuned-distilbert-squad”) tokenizer.save_pretrained(“./fine-tuned-distilbert-squad”) |
This code is a bit complex. Let’s break it down step by step.
Dataset and Preprocessing
The SQuAD dataset is a popular dataset for question answering and it is available on the Hugging Face hub. You can load it using the load_dataset() function from Hugging Face’s datasets library.
|
from datasets import load_dataset
dataset = load_dataset(“squad”) |
Every dataset is different. But this particular dataset is dictionary-like with keys “title”, “context”, “question”, and “answers”. The “context” is a piece of moderately long text. The “question” is a question sentence. The “answers” is a dictionary with the key “text” and “answer_start“. The “text” maps to a short string that is the answer to the question. The “answer_start” maps to the start position of the answer in the context. The “title” can be ignored, as it provides the title of the article that the context is extracted from.
To use the dataset for training, you need to know how the model expects the input and what kind of output it produces. In the case of DistilBERT for question answering, the model is fixed by the implementation of the DistilBertForQuestionAnswering class unless you decide to write your own model implementation. In this class, the model expects the input as a sequence of integer token IDs and the output is two vectors of logits, one for the start position and one for the end position of the answer.
You can find the details of the input and output format of the model in the previous post. Or you can find the details in the DistilBertForQuestionAnswering class documentation.
In order to use the dataset for training, you need to do some preprocessing. This is to transform the dataset into a format that matches the model’s input and output format. The dataset object loaded from the Hugging Face hub allows you to do this with the map() method, in which the transformation is implemented as a custom function, preprocess_function().
|
... tokenized_datasets = dataset.map(preprocess_function, batched=True, remove_columns=dataset[“train”].column_names) |
Note that the preprocess_function() is to accept a batch from the dataset as you used batched=True in the map() method.
In the preprocess_function(), the tokenizer is invoked with the questions from examples["question"] and the context from examples["context"]. The question is stripped of extra spaces and the context is truncated to fit in the maximum length of 384 tokens. The use of a tokenizer in this function is different from what you have seen in the previous post:
|
... inputs = tokenizer( questions, examples[“context”], max_length=384, truncation=“only_second”, return_offsets_mapping=True, padding=“max_length”, ) |
Firstly, the tokenizer is invoked with a batch of questions and the context. For potentially ragged input, the tokenizer will pad the input to the maximum length of the batch. Secondly, with return_offsets_mapping=True, the tokenizer returns a dictionary with the keys “input_ids“, “attention_mask“, and “offset_mapping“. The “input_ids” is the sequence of integer token IDs. The “attention_mask” is a binary mask that indicates which tokens are real (1) and which are padded (0). The “offset_mapping” is what is added by setting return_offsets_mapping=True and it is a list of tuples that indicates the character positions (start and end offsets) of each token in the original text.
The input_ids from the tokenizer output concatenates the question and the context in the format of:
|
[CLS] question [SEP] context [SEP] |
which is what the model expects. The answer from the dataset is a string, and the character offset from which the answer can be found in the original context. This is different from what the model produces, namely, the logits of token positions. Therefore, you used a for-loop in preprocess_function() to recreate the start and end token positions of the answer.
In this code, the tokenizer is invoked with additional arguments. Setting return_offsets_mapping=True will make the returned object contain offset_mapping, a list of tuples identifying the start and end positions of each token in each input text.
First, the offset_mapping is popped from the object returned by the tokenizer since it is not needed for the training. Then for each answer, you identified the character start and end offset from the context. You can verify this with the code like the following:
|
... start_char = answer[“answer_start”][0] end_char = start_char + len(answer[“text”][0]) assert answer[“text”] == context[start_char:end_char] |
Even if you know the character offset, the model operates on token positions.
Recall that the tokenizer concatenated the question and the context. Fortunately, the tokenizer provided the clue to identify the start and end of the context from its output. In inputs.sequence_ids(i), it is a Python list of integers or None, corresponding to the element i of the batch. The list holds None for the position where a special token is found and an integer for which the position is a token from the actual input. In your use case, you invoked the tokenizer with question first and context the second, therefore integer 0 corresponds to the question and 1 corresponds to the context.
Therefore, you can identify the token start and end offset of the context by checking where the integer 1 first and last appears in the sequence_ids list:
|
... sequence_ids = inputs.sequence_ids(i) context_start = sequence_ids.index(1) context_end = len(sequence_ids) – 1 – sequence_ids[::–1].index(1) |
Given you know the start and end token positions of the context, you still need to check if the answer is covered by any token. This is done in a loop by checking each token one by one. You use a for-loop to iterate over each pair of offsets and check if the start and end character positions of the answer are within any token. If so, the position of the token is remembered as the start_positions and end_positions. For answers not found (e.g., due to the context was clipped), they are set to 0.
|
... # If the answer is not fully inside the context, label it (0, 0) if offsets[context_start][0] > end_char or offsets[context_end][1] < start_char: start_positions.append(0) end_positions.append(0) else: # Otherwise find the start and end token positions idx = context_start while idx <= context_end and offsets[idx][0] <= start_char: idx += 1 start_positions.append(idx – 1)
idx = context_end while idx >= context_start and offsets[idx][1] >= end_char: idx -= 1 end_positions.append(idx + 1) |
At the end of the preprocess_function(), the object inputs is returned. It is dictionary-like with keys input_ids, attention_masks, start_positions, and end_positions. You must not change the name of these keys because the DistilBertForQuestionAnswering class expects such arguments in the forward() method.
The DistilBERT model expects you to call it with the arguments input_ids. And if you call with a large batch, attention_masks is required as well to tell which token in the input are paddings. If you call with the optional start and end positions, the cross-entropy loss will be computed as well. This is how the transformers library is designed to help you call the model in inference and training with the same interface.
Running the Training
To run this code, you need to install the following packages:
|
pip install torch datasets transformers accelerate |
While you can expect the requirements of torch, transformers, and datasets. The accelerate package is a dependency when you use the Trainer class from the transformers library.
You may expect training a complex model like DistilBERT to require a lot of code. Indeed, it is not easy since you need to decide what optimizer to use, the number of epochs to train, and the hyperparameters such as batch size, learning rate, weight decay, etc. You even need to handle the checkpointing so that you can resume the training in case of interruption.
This is why the Trainer class is introduced. You just need to set up the training arguments, then set up the Trainer with the dataset, and then run the training:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 |
... training_args = TrainingArguments( output_dir=“./results”, eval_strategy=“epoch”, learning_rate=2e–5, per_device_train_batch_size=16, per_device_eval_batch_size=16, num_train_epochs=3, weight_decay=0.01, ) trainer = Trainer( model=model, args=training_args, train_dataset=tokenized_datasets[“train”], eval_dataset=tokenized_datasets[“validation”], processing_class=tokenizer, ) trainer.train() |
The Trainer will handle the checkpointing, the logging, and the evaluation in one function call. You just need to save the fine-tuned model (together with the tokenizer since they are loaded together) in the Hugging Face format once the training is complete:
|
... model.save_pretrained(“./fine-tuned-distilbert-squad”) tokenizer.save_pretrained(“./fine-tuned-distilbert-squad”) |
That’s all you need to do. Even if you did not specify using a GPU for the training, the Trainer will automatically discover the GPU on your system and use it to speed up the process. The code above, although not very long, is the complete code for fine-tuning DistilBERT on the SQuAD dataset.
If you run this code, you will expect the following output:
|
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 |
Some weights of DistilBertForQuestionAnswering were not initialized from the model checkpoint at distilbert-base-uncased and are newly initialized: [‘qa_outputs.bias’, ‘qa_outputs.weight’] You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference. Map: 100%|████████████████████████████████| 10570/10570 [00:01<00:00, 5387.60 examples/s] ‘loss’: 2.9462, ‘grad_norm’: 13.834440231323242, ‘learning_rate’: 1.9391171993911722e-05, ‘epoch’: 0.09 ‘loss’: 1.7333, ‘grad_norm’: 14.540811538696289, ‘learning_rate’: 1.8782343987823442e-05, ‘epoch’: 0.18 ‘loss’: 1.5268, ‘grad_norm’: 15.629022598266602, ‘learning_rate’: 1.8173515981735163e-05, ‘epoch’: 0.27 ‘loss’: 1.4487, ‘grad_norm’: 20.17080307006836, ‘learning_rate’: 1.756468797564688e-05, ‘epoch’: 0.37 ‘loss’: 1.3957, ‘grad_norm’: 21.543432235717773, ‘learning_rate’: 1.69558599695586e-05, ‘epoch’: 0.46 ‘loss’: 1.3816, ‘grad_norm’: 15.349509239196777, ‘learning_rate’: 1.634703196347032e-05, ‘epoch’: 0.55 ‘loss’: 1.314, ‘grad_norm’: 14.986817359924316, ‘learning_rate’: 1.573820395738204e-05, ‘epoch’: 0.64 ‘loss’: 1.2313, ‘grad_norm’: 15.443862915039062, ‘learning_rate’: 1.5129375951293761e-05, ‘epoch’: 0.73 ‘loss’: 1.2613, ‘grad_norm’: 10.729198455810547, ‘learning_rate’: 1.4520547945205482e-05, ‘epoch’: 0.82 ‘loss’: 1.1976, ‘grad_norm’: 18.681406021118164, ‘learning_rate’: 1.39117199391172e-05, ‘epoch’: 0.91 ‘eval_loss’: 1.142066240310669, ‘eval_runtime’: 14.8679, ‘eval_samples_per_second’: 710.926, ‘eval_steps_per_second’: 44.458, ‘epoch’: 1.0 ‘loss’: 1.1858, ‘grad_norm’: 15.170207023620605, ‘learning_rate’: 1.330289193302892e-05, ‘epoch’: 1.0 ‘loss’: 0.962, ‘grad_norm’: 14.375147819519043, ‘learning_rate’: 1.2694063926940641e-05, ‘epoch’: 1.1 ‘loss’: 0.9994, ‘grad_norm’: 13.867342948913574, ‘learning_rate’: 1.2085235920852361e-05, ‘epoch’: 1.19 ‘loss’: 0.9912, ‘grad_norm’: 13.35099983215332, ‘learning_rate’: 1.147640791476408e-05, ‘epoch’: 1.28 ‘loss’: 0.976, ‘grad_norm’: 18.943002700805664, ‘learning_rate’: 1.08675799086758e-05, ‘epoch’: 1.37 ‘loss’: 0.9687, ‘grad_norm’: 12.70341968536377, ‘learning_rate’: 1.025875190258752e-05, ‘epoch’: 1.46 ‘loss’: 0.949, ‘grad_norm’: 10.327693939208984, ‘learning_rate’: 9.64992389649924e-06, ‘epoch’: 1.55 ‘loss’: 0.9482, ‘grad_norm’: 17.166929244995117, ‘learning_rate’: 9.04109589041096e-06, ‘epoch’: 1.64 ‘loss’: 0.9248, ‘grad_norm’: 23.135452270507812, ‘learning_rate’: 8.432267884322679e-06, ‘epoch’: 1.74 ‘loss’: 0.9289, ‘grad_norm’: 15.964847564697266, ‘learning_rate’: 7.823439878234399e-06, ‘epoch’: 1.83 ‘loss’: 0.9605, ‘grad_norm’: 10.738043785095215, ‘learning_rate’: 7.214611872146119e-06, ‘epoch’: 1.92 ‘eval_loss’: 1.0946319103240967, ‘eval_runtime’: 14.7779, ‘eval_samples_per_second’: 715.256, ‘eval_steps_per_second’: 44.729, ‘epoch’: 2.0 ‘loss’: 0.9376, ‘grad_norm’: 22.791458129882812, ‘learning_rate’: 6.605783866057839e-06, ‘epoch’: 2.01 ‘loss’: 0.7745, ‘grad_norm’: 15.398698806762695, ‘learning_rate’: 5.996955859969558e-06, ‘epoch’: 2.1 ‘loss’: 0.7458, ‘grad_norm’: 17.4672908782959, ‘learning_rate’: 5.388127853881279e-06, ‘epoch’: 2.19 ‘loss’: 0.7636, ‘grad_norm’: 13.833612442016602, ‘learning_rate’: 4.779299847792998e-06, ‘epoch’: 2.28 ‘loss’: 0.7803, ‘grad_norm’: 11.179983139038086, ‘learning_rate’: 4.170471841704719e-06, ‘epoch’: 2.37 ‘loss’: 0.7666, ‘grad_norm’: 9.601215362548828, ‘learning_rate’: 3.5616438356164386e-06, ‘epoch’: 2.47 ‘loss’: 0.7784, ‘grad_norm’: 24.625328063964844, ‘learning_rate’: 2.9528158295281586e-06, ‘epoch’: 2.56 ‘loss’: 0.7389, ‘grad_norm’: 13.041014671325684, ‘learning_rate’: 2.343987823439878e-06, ‘epoch’: 2.65 ‘loss’: 0.7636, ‘grad_norm’: 12.822973251342773, ‘learning_rate’: 1.7351598173515982e-06, ‘epoch’: 2.74 ‘loss’: 0.7625, ‘grad_norm’: 12.254212379455566, ‘learning_rate’: 1.1263318112633182e-06, ‘epoch’: 2.83 ‘loss’: 0.727, ‘grad_norm’: 8.469372749328613, ‘learning_rate’: 5.17503805175038e-07, ‘epoch’: 2.92 ‘eval_loss’: 1.1390912532806396, ‘eval_runtime’: 14.8303, ‘eval_samples_per_second’: 712.731, ‘eval_steps_per_second’: 44.571, ‘epoch’: 3.0 ‘train_runtime’: 1106.5639, ‘train_samples_per_second’: 237.489, ‘train_steps_per_second’: 14.843, ‘train_loss’: 1.0775353780946775, ‘epoch’: 3.0 100%|██████████████████████████████████████████████| 16425/16425 [18:26<00:00, 14.84it/s] |
This takes some time to run, even if you use a decent GPU. However, you are fine-tuning a pre-trained model on the new dataset. This is orders of magnitude faster and easier than training from scratch.
Once you finished the training, you can load the model in your other project by using the path:
|
from transformers import DistilBertTokenizerFast, DistilBertForQuestionAnswering
model_path = “./fine-tuned-distilbert-squad” tokenizer = DistilBertTokenizerFast.from_pretrained(model_path) model = DistilBertForQuestionAnswering.from_pretrained(model_path) ... |
Please make sure that model_path is the correct path to find the saved model files from your project.
Further Reading
Below are some links to the documentation of the classes and methods used in this post:
Summary
In this post, you have learned how to fine-tune DistilBERT for a custom question-answering task. Even DistilBERT and question-answering is used as an example, you can apply the same process to other models and tasks. In particular, you learned:
- How to prepare the dataset for training
- How to train or fine-tune the model using the
Trainerinterface from the transformers library