nlp on Marcin Zabłocki blog

Custom classifier on top of BERT-like Language Model - guide

Marcin Zabłocki — Mon, 23 Mar 2020 23:00:00 +0000

Last time I wrote about training the language models from scratch, you can find this post here. Now it’s time to take your pre-trained lamnguage model at put it into good use by fine-tuning it for real world problem, i.e text classification or sentiment analysis. In this post I will show how to take pre-trained language model and build custom classifier on top of it. As in the previous post - I cover all of the important parts in details and I point out possible mistakes.

Why custom model? There are two answers. First is that the fun in deep learning begins only when you can do something custom with your model. Experiments, experiments and more experiments! Second thing is that by implmenting some parts on your own, you gain better understaing of different parts of the modeling itself, but also the whole training/fine-tuning process.

TL;DR

This post covers:

taking existing pre-trained language model and understanding it’s output - here I use PolBERTa trained for Polish language
building custom classification head on top of the LM
using fast tokenizers to efficiently tokenize and pad input text as well as prepare attention masks
preparing reproducible training code with PyTorch Lightning
finding good starting learning rate for the model
validating the trained model on PolEmo 2.0 dataset (benchmark for Polish language sentiment analysis with 4 classes)

The whole post is available as a Colab Notebook, linked below in Additional links and resources section.

Libraries and preparation

In this post I will use pretty standard set of ML-related libraries, including: transformers & tokenizers from HuggingFace, PyTorch Lightning, pandas, scikit-learn and LR Finder.

python

!pip install transformers tokenizers pytorch-lightning

After running the cell below, you will probably have to restart the notebook (one time only).

python

!git clone https://github.com/davidtvs/pytorch-lr-finder.git && cd pytorch-lr-finder && python setup.py install

python

import torch
from torch import nn
from typing import List
import torch.nn.functional as F
from transformers import AutoTokenizer, AutoModelWithLMHead, AdamW, get_linear_schedule_with_warmup
import logging
import os
from functools import lru_cache
from tokenizers import ByteLevelBPETokenizer
from tokenizers.processors import BertProcessing
import pytorch_lightning as pl
from torch.utils.data import DataLoader, Dataset
import pandas as pd
from argparse import Namespace
from sklearn.metrics import classification_report
torch.__version__

Understanding pre-trained language model input / outputs

In order to build your classifier on top of pre-trained language model you must first understand it outputs. Usually it requries reading the related paper or (if possible) reading the documentation / github code. Fortunately, it’s easy for Transformers library - as the models are documented and the return value is described well. For my case (RoBERTa) the docs are here: https://huggingface.co/transformers/model_doc/roberta.html#robertamodel.

We’re interested what RobertaForMaskedLM returns. It this case, output format is the following:

python

sequence_output, pooled_output, (hidden_states), (attentions)

Let’s take a look at it ourselves:

python

tokenizer = AutoTokenizer.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")
model = AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")

We are actually interested only in base model, without the LM head which is:

python

base_model = model.roberta

Let’s see what does output look like by running the model forward pass with some input. We need to tokenize the input first:

python

text = "Mówiąc otwarcie jesteście reliktami. Ilu wiedźminów jeszcze żyje? Kilku? Góra kilkunastu. Nie stanowicie k<mask>. Poza tym, Zakon nie każe sobie płacić za ratowanie ludzi od złego."
enc = tokenizer.encode_plus(text)
enc.keys()

dict_keys(['input_ids', 'token_type_ids', 'attention_mask'])

Tokenizer returns input_ids which are the tokens of your input text. Then token_type_ids and attention_mask - first one is used when you use LM for next-sentence prediction task, second one is to tell the model which tokens are from your input text and which are just padding tokens.

For classification, you only need input_ids and attention_mask as token types are created on the fly by the internal model’s forward function (source).

Let’s run the model. The .unsqueeze(0) is just to simulate batch processing (in this case, it will be batch size of 1).

python

# sequence_output, pooled_output, (hidden_states), (attentions)
out = base_model(torch.tensor(enc["input_ids"]).unsqueeze(0), torch.tensor(enc["attention_mask"]).unsqueeze(0))
out[0].shape, out[1].shape

(torch.Size([1, 12, 768]), torch.Size([1, 768]))

We are interested only in first element of the tuple while building classifier as the second element is used for next-sentence prediction task and it is literally said in the Transformers docs not to use it:

This output is usually not a good summary of the semantic content of the input, you’re often better with averaging or pooling the sequence of hidden-states for the whole input sequence. [source]

Output shape description:

Keep in mind that for bigger batch sizes you have to pad to either longest text’s lenght or to model’s max supported length (which usually for BERT-like models is 512).

The second dimension contains internal representation of each of the input tokens, so for properly tokenized data, i.e:

python

t = "Mówiąc otwarcie jesteście reliktami."
enc = tokenizer.encode_plus(t)
token_representations = base_model(torch.tensor(enc["input_ids"]).unsqueeze(0))[0][0]
print(enc["input_ids"])
print(tokenizer.decode(enc["input_ids"]))
print(f"Length: {len(enc['input_ids'])}")
print(token_representations.shape)

[0, 355, 304, 1748, 71, 12029, 28398, 807, 2578, 19022, 18, 2]
<s> Mówiąc otwarcie jesteście reliktami.</s>
Length: 12
torch.Size([12, 768])

Representation for each token is as follows:

0 = <s> token is represented by token_representations[0]
355 = M token is represented by token_representations[1]

and so on…

Build custom classification head on top of LM

OK, we understand the model’s output and sizes, now we can build custom classification head on top of it. Have you found new awesome activation function? Or maybe you want to try-out different pooling methods to better fit your case? You’re free to experiment. In this example, I will try out custom Mish activation function on top of straightforward classification head inspired by RoBERTa’s original one (source) and DistilBERT’s custom initialization. Because why not!

python

# from https://github.com/digantamisra98/Mish/blob/b5f006660ac0b4c46e2c6958ad0301d7f9c59651/Mish/Torch/mish.py
@torch.jit.script
def mish(input):
    return input * torch.tanh(F.softplus(input))
  
class Mish(nn.Module):
    def forward(self, input):
        return mish(input)

python

class PolBERTaSentimentModel(nn.Module):
    def __init__(self, base_model, n_classes, base_model_output_size=768, dropout=0.05):
        super().__init__()
        self.base_model = base_model
        
        self.classifier = nn.Sequential(
            nn.Dropout(dropout),
            nn.Linear(base_model_output_size, base_model_output_size),
            Mish(),
            nn.Dropout(dropout),
            nn.Linear(base_model_output_size, n_classes)
        )
        
        for layer in self.classifier:
            if isinstance(layer, nn.Linear):
                layer.weight.data.normal_(mean=0.0, std=0.02)
                if layer.bias is not None:
                    layer.bias.data.zero_()

    def forward(self, input_, *args):
        X, attention_mask = input_
        hidden_states, _ = self.base_model(X, attention_mask=attention_mask)
        
        # here I use only representation of <s> token, but you can easily use more tokens,
        # maybe do some pooling / RNNs... go crazy here!
        return self.classifier(hidden_states[:, 0, :])

Hint

As a sanity check, I encourage you to run forward pass of the model on dummy data to validate the code. It also better to do this on CPU, as the stack traces of potential exceptions will be more meaningful. Mysterious CUDA erros very often are caused by simple Python-level errors.

python

classifier = PolBERTaSentimentModel(AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1").roberta, 3)

python

X = torch.tensor(enc["input_ids"]).unsqueeze(0)
attn = torch.tensor(enc["attention_mask"]).unsqueeze(0)

python

classifier((X, attn))

tensor([[0.5643, 0.2237, 0.1142]], grad_fn=<AddmmBackward>)

The output is correct - we’ve run the model with batch size and expected 3-class output (for CrossEntropyLoss).

Preparing fast tokenizers to use with text dataset

When you prepare your dataset you have two options while implementing Dataset class:

return tokenized text and class in __getitem__ function
return just text and class in __getitem__ function and implement Collate function - this function will be used by PyTorch DataLoader to prepare batches

Here I’ve chosen the second approach.

Default tokenizer loaded above (as for Transformers v2.5.1) uses Python implementation. In order to leverage full potential of parallel Rust tokenizers, we need to save the tokenizer’s internal data and then create instance of fast tokenizer with it.

python

!mkdir -p tokenizer

python

tokenizer.save_pretrained("tokenizer")

('tokenizer/vocab.json',
 'tokenizer/merges.txt',
 'tokenizer/special_tokens_map.json',
 'tokenizer/added_tokens.json')

python

!ls tokenizer

merges.txt  special_tokens_map.json  tokenizer_config.json  vocab.json

Once saved, implement CollateFn using fast tokenizers. As you can see, the whole batch can be tokenized in parallel and the tokenizer will also automatically return attention mask and pad input to max length of 512.

Important!

Ensure that all of the special tokens are in appropriate places and their IDs come from the actual tokenizer you’re using - they should not be hardcoded, as different models will have different IDs for padding, BOS, EOS and so on.

python

class TokenizersCollateFn:
    def __init__(self, max_tokens=512):
        t = ByteLevelBPETokenizer(
            "tokenizer/vocab.json",
            "tokenizer/merges.txt"
        )
        t._tokenizer.post_processor = BertProcessing(
            ("</s>", t.token_to_id("</s>")),
            ("<s>", t.token_to_id("<s>")),
        )
        t.enable_truncation(max_tokens)
        t.enable_padding(max_length=max_tokens, pad_id=t.token_to_id("<pad>"))
        self.tokenizer = t

    def __call__(self, batch):
        encoded = self.tokenizer.encode_batch([x[0] for x in batch])
        sequences_padded = torch.tensor([enc.ids for enc in encoded])
        attention_masks_padded = torch.tensor([enc.attention_mask for enc in encoded])
        labels = torch.tensor([x[1] for x in batch])
        
        return (sequences_padded, attention_masks_padded), labels

Getting data for classification

Here I use PolEmo 2.0 dataset, which is a benchmark for sentiment analysis task in Polish language.

python

!wget -O dataset_conll.zip https://clarin-pl.eu/dspace/bitstream/handle/11321/710/dataset_conll.zip?sequence=1&isAllowed=y

python

!unzip -n dataset_conll.zip

Archive:  dataset_conll.zip

python

train_path, val_path, test_path = [f"dataset_conll/all.text.{d}.txt" for d in ("train", "dev", "test")]
train_path, val_path, test_path

('dataset_conll/all.text.train.txt',
 'dataset_conll/all.text.dev.txt',
 'dataset_conll/all.text.test.txt')

This dataset has CSV-like format with __label__ text being the separator. Don’t ask. Let’s take a quick look at the data:

python

df = pd.read_csv(train_path, sep="__label__", header=None, names=["text", "label"], engine="python")
df.head()

	text	label
0	Na samym wejściu hotel śmierdzi . W pokojach j...	meta_minus_m
1	Ze swoimi dziećmi ( 10 lat i 2 latka ) chodzę ...	meta_plus_m
2	Chciał em zestaw do podstawowych prac coś przy...	meta_minus_m
3	Drogie Mamy . Mam o tym laktatorze takie same ...	meta_minus_m
4	własciwie nic się nie dowidział em , podczas b...	meta_minus_m

python

df.label.unique()

array(['meta_minus_m', 'meta_plus_m', 'meta_zero', 'meta_amb'],
      dtype=object)

You will need mapping between text label and integer:

python

label2int = {
  "meta_minus_m": 0,
  "meta_plus_m": 1,
  "meta_zero": 2,
  "meta_amb": 3
}

As you can see below, the Dataset itself is really simple as it returns only the text and label. The rest is handled by the TokenizersCollateFn above.

python

class PolEmo2Dataset(Dataset):
    def __init__(self, path):
        super().__init__()
        self.data_column = "text"
        self.class_column = "class"
        self.data = pd.read_csv(path, sep="__label__", header=None, names=[self.data_column, self.class_column],
                               engine="python")

    def __getitem__(self, idx):
        return self.data.loc[idx, self.data_column], label2int[self.data.loc[idx, self.class_column]]

    def __len__(self):
        return self.data.shape[0]

Sanity check:

python

ds = PolEmo2Dataset(train_path)
ds[13]

('Była m przekonana , że w Sopocie jest jeden hotel , który może zapewnić najlepszą jakość , tym większe było moje zaskoczenie krótkim pobytem służbowym w Małej Anglii . Miała m okazję mieszkać w pięknym i wygodnym pokoju o bardzo wysokim standardzie . Przemiła obsługa , co jest kluczowe i często niesieciowe hotele nie potrafią sobie z tym poradzić , tutaj nie było takiego problemu . Bardzo przyjemna strefa wellness oraz spa . Doskonała restauracja śniadaniowa , przepięknie urządzona , śniadanie bez zarzutu i kawa . . doskonała , każdy kawosz zrozumie . Brak widoku na morze został zupełnie zrekompensowany . Polecam i mam nadzieję , że taki poziom już pozostanie . Dziękuję za miłe wspomnienia . ',
 1)

Preparing reproducible training code with PyTorch Lightning

Most recently, PyTorch Lightning is my go to library for quick prototyping with PyTorch. You have clear API that is actually extension of the original PyTorch nn.Module one with all of the repeatable parts like training loop, validation loop, using GPUs, learning rate schedulers, gradient accumulation, tensorboard, checkpointing and many others handled for you.

There is only a thin API you have to implement, which is pl.LightningModule. Required parts are:

__init__ to initialize custom variables in your model
forward same as you would do in nn.Module - here I wrap my model inside of the pl.LightningModule but you could as well implement it directly here
training_step - take care of the forward pass for the model and return loss
validation_step and test_step - same as above but for the different learning phases
train_dataloader - DataLoader for your training dataset
val_dataloader and test_dataloader - same as above but for the different learning phases
configure_optimizers - prepare optimizer and (optionally) LR scheudles

python

class TrainingModule(pl.LightningModule):
    def __init__(self, hparams):
        super().__init__()
        self.model = PolBERTaSentimentModel(AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1").roberta, 4)
        self.loss = nn.CrossEntropyLoss()
        self.hparams = hparams

    def step(self, batch, step_name="train"):
        X, y = batch
        loss = self.loss(self.forward(X), y)
        loss_key = f"{step_name}_loss"
        tensorboard_logs = {loss_key: loss}

        return { ("loss" if step_name == "train" else loss_key): loss, 'log': tensorboard_logs,
               "progress_bar": {loss_key: loss}}

    def forward(self, X, *args):
        return self.model(X, *args)

    def training_step(self, batch, batch_idx):
        return self.step(batch, "train")
    
    def validation_step(self, batch, batch_idx):
        return self.step(batch, "val")

    def validation_end(self, outputs: List[dict]):
        loss = torch.stack([x["val_loss"] for x in outputs]).mean()
        return {"val_loss": loss}
        
    def test_step(self, batch, batch_idx):
        return self.step(batch, "test")
    
    def train_dataloader(self):
        return self.create_data_loader(self.hparams.train_path, shuffle=True)

    def val_dataloader(self):
        return self.create_data_loader(self.hparams.val_path)

    def test_dataloader(self):
        return self.create_data_loader(self.hparams.test_path)
                
    def create_data_loader(self, ds_path: str, shuffle=False):
        return DataLoader(
                    PolEmo2Dataset(ds_path),
                    batch_size=self.hparams.batch_size,
                    shuffle=shuffle,
                    collate_fn=TokenizersCollateFn()
        )
        
    @lru_cache()
    def total_steps(self):
        return len(self.train_dataloader()) // self.hparams.accumulate_grad_batches * self.hparams.epochs

    def configure_optimizers(self):
        optimizer = AdamW(self.model.parameters(), lr=self.hparams.lr)
        lr_scheduler = get_linear_schedule_with_warmup(
                    optimizer,
                    num_warmup_steps=self.hparams.warmup_steps,
                    num_training_steps=self.total_steps(),
        )
        return [optimizer], [{"scheduler": lr_scheduler, "interval": "step"}]

Hints and comments

hparams from __init__ is a Namespace-type object, which holds all of the models hyperparameters; PyTorch Lightning takes good care of them when checkpointing the model
*_step functions can be implemented with the same method with argument
*_dataloader functions can also be implemented in this way
in configure_optimizers the line: [{"scheduler": lr_scheduler, "interval": "step"}] takes care of the step-based LR scheduler - it does LR warmup and then linear decay, implementation comes from Transformers library

The code above is fairly simple and we don’t need anything more here. For more robust examples I send you to PyTorch Lightning docs.

Finding good starting learning rate for the model

There is a method of approximating “good” learning rate for the training. I first encountered it in Fast.ai Deep Learning course. This approach is called LR-Finder and boils down to plotting loss of the training batch (starting from really small LR) while increasing learning rate up to the poit where the loss blows up. Then, the “good” LR is one of the values that appears when the Loss/Learning Rate curve is decreasing.

Important!

Keep in mind that this method does not guarantee to provide good learning rate for every dataset. It’s just a method discovered to often work in practice.

Here I use implementation for PyTorch from here (installed at the top of this notebook).

You need to specify initial learning rate and number of steps. Other parameters are optional.

python

from torch_lr_finder import LRFinder
hparams_tmp = Namespace(
    train_path=train_path,
    val_path=val_path,
    test_path=test_path,
    batch_size=8,
    warmup_steps=100,
    epochs=3,
    lr=lr,
    accumulate_grad_batches=1,
)
module = TrainingModule(hparams_tmp)
criterion = nn.CrossEntropyLoss()
optimizer = AdamW(module.parameters(), lr=5e-7)
lr_finder = LRFinder(module, optimizer, criterion, device="cuda")
lr_finder.range_test(module.train_dataloader(), end_lr=100, num_iter=100, accumulation_steps=hparams_tmp.accumulate_grad_batches)
lr_finder.plot()
lr_finder.reset()

Stopping early, the loss has diverged
Learning rate search finished. See the graph with {finder_name}.plot()

From the plot above we can guess that something between 1e-5 and 1e-4 would be a good learning rate, as everyhing higher results in increased loss.

To ensure the right spot was chosen, you can plot it:

python

lr = 4e-5
lr

4e-05

python

lr_finder.plot(show_lr=lr)

Training the sentiment analysis classifier on PolEmo2.0

We know the learning rate, we have the code for the model. Now it’s time to run it.

python

hparams = Namespace(
    train_path=train_path,
    val_path=val_path,
    test_path=test_path,
    batch_size=8,
    warmup_steps=100,
    epochs=3,
    lr=lr,
    accumulate_grad_batches=1
)
module = TrainingModule(hparams)

python

# just to get rid of any leftovers before the training - useful when you run different cells multiple times
import gc; gc.collect()
torch.cuda.empty_cache()

Using Trainer from PyTorch Lightning

The pl.Trainer provides all-in-one solution for training, validation and testing. There is a lot of useful parameters for almost every aspect of training you can think of, everything is well documented here.

python

trainer = pl.Trainer(gpus=1, max_epochs=hparams.epochs, progress_bar_refresh_rate=10,
                     accumulate_grad_batches=hparams.accumulate_grad_batches)

python

trainer.fit(module)

Evaluating the trained model

As this dataset is small, I’ve used scikit-learn to get the overall classification report from the model.

Important!

Be aware that for big datasets this approach might not be optimal as it transfers the data from GPU to CPU for the whole dataset and puts every prediction in memory. For big datasets it’s better to implement metrics inside of the pl.LightningModule when the dataset is processed in batches.

Ensure that you disable gradients and put model into evaluation state (module.eval()) to disable dropouts during prediction.

python

with torch.no_grad():
    progress = ["/", "-", "\\", "|", "/", "-", "\\", "|"]
    module.eval()
    true_y, pred_y = [], []
    for i, batch_ in enumerate(module.test_dataloader()):
        (X, attn), y = batch_
        batch = (X.cuda(), attn.cuda())
        print(progress[i % len(progress)], end="\r")
        y_pred = torch.argmax(module(batch), dim=1)
        true_y.extend(y.cpu())
        pred_y.extend(y_pred.cpu())
print("\n" + "_" * 80)
print(classification_report(true_y, pred_y, target_names=label2int.keys(), digits=6))

python

    ________________________________________________________________________________
                  precision    recall  f1-score   support
    
    meta_minus_m   0.880878  0.828909  0.854103       339
     meta_plus_m   0.867521  0.894273  0.880694       227
       meta_zero   0.957983  0.966102  0.962025       118
        meta_amb   0.574324  0.625000  0.598592       136
    
        accuracy                       0.832927       820
       macro avg   0.820177  0.828571  0.823854       820
    weighted avg   0.837433  0.832927  0.834617       820

Additional links and resources

this notebook in Google Colab - https://colab.research.google.com/drive/1sajgpLTrTJDzRSlxycy8aE6ysxGqaT18
this notebook on GitHub - https://github.com/marrrcin/transformers-sentiment-analysis
PyTorch Lightning documentation - https://pytorch-lightning.readthedocs.io/en/0.7.1/

Training RoBERTa from scratch - the missing guide

Marcin Zabłocki — Sun, 15 Mar 2020 23:00:00 +0000

After hours of research and attempts to understand all of the necessary parts required for one to train custom BERT-like model from scratch using HuggingFace’s Transformers library I came to conclusion that existing blog posts and notebooks are always really vague and do not cover important parts or just skip them like they weren’t there - I will give a few examples, just follow the post.

I’ve decided to get my hands dirty and try to train transformer language model (BERT or other) for Polish language. Here you will find all of the missing pieces that I have encountered while doing it.

Many thanks to Egnyte Inc. for providing me with necessary resources to run the training on powerful machine with 4x NVIDIA P100 GPUs.

TL;DR

This post covers:

what mistakes to avoid
how to prepare own dataset with details - I use Polish Wikipedia and Wolne Lektury datasets
how to prepare training configuration
how to train new tokenizers using HugginFace’s Rust tokenizers
how to use run_language_modeling.py script from HuggingFace library to train your model - what parameters to set, how to deal with the script
how to monitor training process
how to share the model on HuggingFace’s model hub

Getting the dataset

The most important part when dealing with language models is to have solid dataset with text in the language you will be modeling. In my proof of concept I’ve started with Polish Wikipedia and then extended my dataset with Wolne Lektury (repository of public domain polish books and poems).

Preparing Polish Wikipedia dump

So there is a “just get wikipedia” part that noone covers in details but I find it important. First of all, on the wikipedia dumps page there are many files that you can download. I’ve decided to download dump from 2020-02-20 labeled as:

Articles, templates, media/file descriptions, and primary meta-pages, in multiple bz2 streams, 100 pages per stream

which consisted of 7 files with names like: plwiki-20200220-pages-articles-multistream*.xml*.bz2. Why? That kind of split allows to simpify preprocessing - I don’t have to have lots of RAM to processs smaller files. Another benefit is that I can just hold-out one or two of those files for validation phase after my language model learning is complete.

Every file is a huge XML containing articles with mediawiki-specific markup language. I used this tool to extract well formated JSONs from them: https://github.com/attardi/wikiextractor (clone the repo).

So now I have a list of files for both training and evaluation:

bash

    # train.txt
    plwiki-20200220-pages-articles-multistream1.xml-p1p169750
    plwiki-20200220-pages-articles-multistream2.xml-p169751p510662
    plwiki-20200220-pages-articles-multistream3.xml-p510663p1056310
    plwiki-20200220-pages-articles-multistream4.xml-p1056311p1831508
    plwiki-20200220-pages-articles-multistream5.xml-p1831509p3070393
    plwiki-20200220-pages-articles-multistream6.xml-p4570394p4720470

and

bash

    # eval.txt
    plwiki-20200220-pages-articles-multistream6.xml-p3070394p4570393

To extract them just run:

bash

cat train.txt | xargs -I@  python wikiextractor/WikiExtractor.py @ --bytes=100M --json --output="./plwiki-json/train/@"

After extracting you will get the following structure:

bash

├── eval
│   └── plwiki-20200220-pages-articles-multistream6.xml-p3070394p4570393
│       └── AA
│           ├── wiki_00
│           ├── wiki_01
│           ├── wiki_02
│           └── wiki_03
└── train
    ├── plwiki-20200220-pages-articles-multistream1.xml-p1p169750
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       └── wiki_02
    ├── plwiki-20200220-pages-articles-multistream2.xml-p169751p510662
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       └── wiki_02
    ├── plwiki-20200220-pages-articles-multistream3.xml-p510663p1056310
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       └── wiki_02
    ├── plwiki-20200220-pages-articles-multistream4.xml-p1056311p1831508
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       ├── wiki_02
    │       └── wiki_03
    ├── plwiki-20200220-pages-articles-multistream5.xml-p1831509p3070393
    │   └── AA
    │       ├── wiki_00
    │       ├── wiki_01
    │       ├── wiki_02
    │       └── wiki_03
    └── plwiki-20200220-pages-articles-multistream6.xml-p4570394p4720470
        └── AA
            └── wiki_00

Extracting text from wikipedia

After you get JSON lines from wikipedia dump, you need to transform articles in JSON text into plaintext consumable by the language model training script.

Important!

I’ve probably made a mistake here which is important later on! I first tokenized the text by sentences and outputted each sentence in a separate line because of the ambiguity of available resources about dataset preparation. I think that it affects the quality of both tokenization and training. Do not make that mistake.

After digging throughout the HuggingFace’s github I’ve finally got an answer that there should be one document per line. OK - but what if the text is really long? We will be dealing not only with long wiki articles but later on - with books that have thousands of characters and current SOTA language models have limit of around 512-768 tokens (depends on the model). The possible solution for that is to use sliding window over sentences in the input text. For more details see the following thread: https://github.com/huggingface/transformers/issues/2693

To extract the text and tokenize the input I’ve used nltk with custom polish abbreviations - thanks to Krzysztof Sopyła’s awesome gist: https://gist.github.com/ksopyla/f05fe2f48bbc9de895368b8a7863b5c3.

You can find my extraction notebook here: https://gist.github.com/marrrcin/e383b75a5d0dad42048847d97965e037. I decided to skip it here to not pollute the post.

Key takeaways from preprocessing:

tokenize article into sentences
use sliding window over sentences to handle long texts - I used sliding window of size 4

Once the extraction completes, you will have training file like this:

bash

Sumatra Północna Sumatra Północna (indonez. "Sumatera Utara") – prowincja w Indonezji w północnej części Sumatry. Obejmuje również wyspę Nias i wyspy Batu leżące na Oceanie Indyjskim. Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan.
"Sumatera Utara") – prowincja w Indonezji w północnej części Sumatry. Obejmuje również wyspę Nias i wyspy Batu leżące na Oceanie Indyjskim. Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan. Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba.
Obejmuje również wyspę Nias i wyspy Batu leżące na Oceanie Indyjskim. Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan. Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba. Bliżej wybrzeży nizinne tereny w dużym stopniu pokryte bagnami.
Powierzchnia 70 787 km²; 12 391 tys. mieszkańców (2005); stolica Medan. Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba. Bliżej wybrzeży nizinne tereny w dużym stopniu pokryte bagnami. Ludność prowincji tworzą Malajowie wyznający islam, Batakowie będący w większości chrześcijanami (gł. protestantami) oraz przybysze z Jawy, Chin, Indii.
Większą część powierzchni zajmują góry Barisan (Sinabung 2460 m n.p.m.) z licznymi wulkanami i jeziorem Toba. Bliżej wybrzeży nizinne tereny w dużym stopniu pokryte bagnami. Ludność prowincji tworzą Malajowie wyznający islam, Batakowie będący w większości chrześcijanami (gł. protestantami) oraz przybysze z Jawy, Chin, Indii. Główne miasta: Medan, Binjai, Pematang Siantar, Tebing Tinggi

Prepare Wolne Lektury dump

For extraction of polish books I have crawled https://wolnelektury.pl/katalog/ site. It’s pretty straightforward process - all of the books can be downloaded in well formatted TXT files. The preprocessing steps are exactly the same as for wiki, namely: tokenize book into sentences, use sliding window to transform long books into smaller chunks without dropping any data - this time I used sliding window of 8 sentences. Gist of preprocessing notebook: https://gist.github.com/marrrcin/bcc115fbadf79eba9d9c8ca711da9e20.

Training new tokenizer

At this point, I’ve decided to go with RoBERTa model. Model you choose determines the tokenizer that you will have to train. For RoBERTa it’s a ByteLevelBPETokenizer, for BERT it would be BertWordPieceTokenizer (both from tokenizers library). Training the tokenizer is super fast thanks to the Rust implementation that guys at HuggingFace have prepared (great job!).

python


from pathlib import Path

from tokenizers import ByteLevelBPETokenizer
from glob import glob
paths = list(
    glob("your_input_dataset_files/*.txt")
)
# Initialize a tokenizer
tokenizer = ByteLevelBPETokenizer(lowercase=False)

# Customize training
tokenizer.train(files=paths, vocab_size=32000, min_frequency=3, special_tokens=[
    "<s>",
    "<pad>",
    "</s>",
    "<unk>",
    "<mask>",
])

# Save files to disk
import os
OUT_DIR = "polish_tokenizer_bpe_32k"
os.makedirs(OUT_DIR, exist_ok=True)
tokenizer.save(OUT_DIR, "pl")

Here, the special_tokens part is super important - those tokens need to be exactly in the same order, because otherwise the model will crash!.

For BERT tokenizer however, you will have:

python

from tokenizers import BertWordPieceTokenizer, ByteLevelBPETokenizer
# Initialize a tokenizer
tokenizer = BertWordPieceTokenizer(lowercase=False, handle_chinese_chars=False)

Adjusting the `run_language_modeling.py` script

The version I used and modified comes from v.2.5.1 release of HuggingFace transformers: https://github.com/huggingface/transformers/blob/v2.5.1/examples/run_language_modeling.py

Dataset - check, tokenizer - check. Now it’s time for the model training. Initially, I wanted to use the run_language_modeling.py script from HuggingFace git repo without changes, but it’s it does not support fast Rust-based tokenizers (as for 2020-03-16). Here I mention the minimal changes I made to make it work.

Fast tokenizers support in `run_language_modeling.py`

In order to plug fast tokenizer I’ve trained above into this script I had to modify the LineByLineTextDataset that’s provided there. The final version looks like this:

python

class LineByLineTextDataset(Dataset):
    def __init__(self, t: PreTrainedTokenizer, args, file_path: str, block_size=512):
        assert os.path.isfile(file_path)
        logger.info("Creating features from dataset file at %s", file_path)
        
        # -------------------------- CHANGES START
        bert_tokenizer = os.path.join(args.tokenizer_name, "vocab.txt")
        if os.path.exists(bert_tokenizer):
            logger.info("Loading BERT tokenizer")
            from tokenizers import BertWordPieceTokenizer
            tokenizer = BertWordPieceTokenizer(os.path.join(args.tokenizer_name, "vocab.txt"), handle_chinese_chars=False, lowercase=False)
            tokenizer.enable_truncation(512)
        else:
            from tokenizers import ByteLevelBPETokenizer
            from tokenizers.processors import BertProcessing
            logger.info("Loading RoBERTa tokenizer")
            
            tokenizer = ByteLevelBPETokenizer(
                os.path.join(args.tokenizer_name, "vocab.json"),
                os.path.join(args.tokenizer_name, "merges.txt")
            )
            tokenizer._tokenizer.post_processor = BertProcessing(
                ("</s>", tokenizer.token_to_id("</s>")),
                ("<s>", tokenizer.token_to_id("<s>")),
            )
            tokenizer.enable_truncation(max_length=512)

        logger.info("Reading file %s", file_path)
        with open(file_path, encoding="utf-8") as f:
            lines = [line for line in f.read().splitlines() if (len(line) > 0 and not line.isspace())]

        logger.info("Running tokenization")
        self.examples = tokenizer.encode_batch(lines)
        
        # -------------------------- CHANGES END

    def __len__(self):
        return len(self.examples)

    def __getitem__(self, i):
        return torch.tensor(self.examples[i].ids, dtype=torch.long)

So the only changes are in the way of processing input args.tokenizer_name parameter - if the folder has vocab.txt it’s BERT tokenizer, RoBERTa otherwise. Thanks to the tokenizer.encode_batch(lines) the whole dataset is tokenized in parallel.

Important!

If you don’t have enough RAM to fit the whole dataset into memory, you have to implement your own dataset. My input dataset - Polish wikipedia with sliding window of 4 + Wolne Lektury with sliding window of 8 took ~235GB of RAM.

Preparing training configuration

This part is the most tricky one, as there are a lot of magic numbers in the posts / notebooks I’ve found, even in the official one https://huggingface.co/blog/how-to-train. They just skip it… Start with creating the directory for the configuration, say MyRoBERTaConfig. Copy the tokenizer files and rename them. Be sure to finally have:

bash

    config.json
    tokenizer_config.json
    merges.txt
    vocab.json

in this directory.

Training configuration for RoBERTa

You know what? You don’t have to specify any model-specific parameters if you want to go with the defaults! The only thing that you actually need to specify is the vocabulary size, because it determines the output layer size (Language Model is predicting the tokens after all…). For my RoBERTa I have the following config.json:

javascript

{"architectures": ["RobertaForMaskedLM"], "max_position_embeddings": 514, "vocab_size": 32000}

The architectures is self explanatory - this is the language model you will be training, vocab_size comes from the tokenizer you have created.

Important!

The tricky one is max_position_embeddings - no one will tell you that, but if you don’t set this here, the script will just crash. Awesome. I believe that for BERT model it’s not required.

Tokenizer configuration for RoBERTa is simple (tokenizer_config.json). However it’s probably not required, because we’ve overridden the LineByLineTextDataset dataset with tokenizer.

javascript

{"max_len": 512}

Training your own RoBERTa language model

Here I will show two types of run_language_modeling.py script execution - one for training from scratch and the second - for fine tuning.

Training from scratch

This is the usual case for any new language model. First, the script:

bash

#!/bin/bash
export TRAIN_FILE=plwiki.train.txt
export EVAL_FILE=plwiki.eval.txt
# run_language_modeling_with_tokenizers.py -- it's the version with support for fast tokenizers, see above
python run_language_modeling_with_tokenizers.py \
    --train_data_file $TRAIN_FILE \
    --eval_data_file $EVAL_FILE \
    --output_dir ./MyRoBERTa \
    --model_type roberta \
    --mlm \
    --config_name ./MyRoBERTaConfig \
    --tokenizer_name ./MyRoBERTaConfig \
    --do_train \
    --do_eval \
    --line_by_line \
    --learning_rate 1e-5 \
    --num_train_epochs 5 \
    --save_total_limit 20 \
    --save_steps 5000 \
    --per_gpu_train_batch_size 8 \
    --warmup_steps=10000 \
    --logging_steps=100 \
    --gradient_accumulation_steps=4 \
    --seed 666 --block_size=512

Arguments explanation: (only the non-trivial)

--output_dir - here you will get all of the model checkpoints, saved as often as specified in --save_steps argument
--mlm - configure the script to run Masked Language Model task. The script will fortunately yell at you if the --model_type is not appropriate for this task
--line_by_line - we say to use LineByLineTextDataset, the script will use my overridden version with tokenizers
--save_total_limit - how many checkpoint to keep on the disk - the oldest one will be removed once this limit is exceeded. It’s good to set it here in order to not clog your disk
--per_gpu_train_batch_size - setting this parameter might require a few failed runs fro you. If you don’t know how big batch size you can fit into your GPU, just sample your training data, to say 10k lines and launch the script a few times. Monitor the GPU usage by using watch -n 1 nvidia-smi command. Set it to as high number as possible. For Nvidia P100 and 512 --block_size, it’s 8.
--gradient_accumulation_steps - this parameter configures how many steps you will pass in forward pass of the network before running backward pass, it’s useful when you GPU cannot handle large batches
--warmup_steps - configures learning rate scheduler; as per BERT paper recommendation - it’s good to do warmup and then warmdown. Number of steps depends on your training dataset size
--logging_steps - tells how often to store Tensorboard logs - the script only outputs loss and learning rate

Training from checkpoint (finetuning)

After I’ve pretrained my model on wikipedia only, I decided to finetune it further on the Wolne Lektury dataset. Training script has produced a few snapshots, from which every one has the following structure:

bash

├── checkpoint-20000
│   ├── config.json
│   ├── merges.txt
│   ├── optimizer.pt
│   ├── pytorch_model.bin
│   ├── scheduler.pt
│   ├── special_tokens_map.json
│   ├── tokenizer_config.json
│   ├── training_args.bin
│   └── vocab.json

Important!

If you’re changing the input dataset, you cannot use existing scheduler.pt and optimizer.pt as you’re not “resuming” the training. You should remove (but backup first!) those files if you want to fine tune on another dataset.'

OK, now the script itself:

bash

#!/bin/bash
export TRAIN_FILE=wolne-lektury.sliding8.txt
export EVAL_FILE=plwiki.eval.txt

python run_language_modeling_with_tokenizers.py \
    --train_data_file $TRAIN_FILE \
    --eval_data_file $EVAL_FILE \
    --output_dir ./MyRoBERTa_with_another_dataset \
    --model_type roberta \
    --model_name_or_path ./MyRoBERTa/checkpoint-30000 \
    --mlm \
    --config_name ./MyRoBERTaConfig \
    --tokenizer_name ./MyRoBERTaConfig \
    --do_train \
    --do_eval \
    --line_by_line \
    --learning_rate 5e-5 \
    --num_train_epochs 3 \
    --save_total_limit 20 \
    --save_steps 5000 \
    --per_gpu_train_batch_size 8 \
    --warmup_steps=5000 \
    --logging_steps=100 \
    --gradient_accumulation_steps=4 \
    --mlm_probability=0.2 \
    --seed 666 --block_size=512

Arguments - notable mentions:

--mlm_probability=0.2 - this parameter controls the percentage of the tokens you mask during training; default is 0.15, I’ve decided to change it to make the training more difficult to the model
--model_name_or_path ./MyRoBERTa/checkpoint-30000 - this should point to the checkpoint of your previously trained model; remember the note above!

Training phase

I highly recommend you to launch the script using screen tool, which allows to easily spawn and manage multiple SSH sessions - you will be able to disconnect from the server while the training runs. And it will run for days… so brace yourself. You can use screen -UL command (-U for enabling UTF-8 mode for progress bars and -L to log session output to file - it’s helpful when something crashes).

During training, the script will output Tensorboard logs into runs directory. You can monitor them by running tensorboard command:

bash

tensorboard --logdir=runs

It will launch small webserver on localhost:6006 and you will be able to monitor the training.

Training summary for Polish RoBERTa a.k.a PolBERTa

I’ve run my training in three phases:

From scratch on Polish Wikipedia only (1 sentence per line) for 370k steps using learning rate 1e-5 with 10k warmup steps.
Fine tuning on Wolne Lektury only (8 sentences per line) for 60k steps, starting after checkpoint from 1. Learning rate 5e-5 with 5k warmup steps.
Fine tuning on Wolne Lektury (8 sentences per line) and Polish Wikipedia (4 sentences per line) for 30k steps. Learning rate 5-e5 with 10k warmup steps.

Tensorboard for all of the runs can be viewed here https://tensorboard.dev/experiment/aYrAE9uMTxGKRtLUOFjprg/

Exporting the model

If you want to share the model with the NLP community (which I highly encourage you to do!) you need to export it in appropriate format. Prepare the path for your model’s latest checkpoint and then run the following code:

python

from transformers import AutoModelWithLMHead, AutoTokenizer
import os
directory = "/path/to/your/model/checkpoint-30000"
model = AutoModelWithLMHead.from_pretrained(directory)
tokenizer = AutoTokenizer.from_pretrained(directory)
out = "MyBERTa-base-cased-v1"
os.makedirs(out, exist_ok=True)
model.save_pretrained(out)
tokenizer.save_pretrained(out)

Then use transformers-cli util to upload the model:

bash

transformers-cli upload ./MyBERTa-base-cased-v1/

Using my pre-trained model

I’ve uploaded my pre-trained RoBERTa to HuggingFace’s model hub: marrrcin/PolBERTa-base-polish-cased-v1 Usage:

python

from transformers import AutoModelWithLMHead, AutoTokenizer, pipeline

model = AutoModelWithLMHead.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")
tokenizer = AutoTokenizer.from_pretrained("marrrcin/PolBERTa-base-polish-cased-v1")

fill_mask = pipeline('fill-mask', model=model, tokenizer=tokenizer)
fill_mask("Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad <mask>, a postępujesz tak jakbyś miał niecałe dziesięć.")

Output:

javascript

[{'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad sto, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.1013106107711792,
  'token': 1675},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad tysiąc, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.09420681744813919,
  'token': 31144},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad połowę, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.0891065001487732,
  'token': 13627},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad dwadzieścia, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.07204979658126831,
  'token': 18811},
 {'sequence': '<s> Wiadomym jest mi, że masz lat blisko czterdzieści, wyglądasz na blisko trzydzieści, wyobrażasz sobie, że masz nieco ponad godzinę, a postępujesz tak jakbyś miał niecałe dziesięć.</s>',
  'score': 0.059559762477874756,
  'token': 20775}]

Follow up and TODO

Here are a few follow up things and my planned future work on Polish RoBERTa:

prepare model card for HuggingFace's model hub
evaluate PolBERTa on PolEmo2.0 and other downstream tasks
re-train PolBERTa again, using bigger vocab size for tokenizer
fine tune model further on another open datasets

If you find any bug in the post or something is not clear - feel free to post a comment, I will be glad to improve this short practical guide.

Thanks to…

I would like to thank Egnyte Inc. for providing me with required resources to train the language model and also Darek Kłeczek (creator of original PolBERT) for tips related to training BERT.

Additional links and resources

PolBERT - original BERT-base-uncased LM trained by Darek Kłeczek
HuggingFaces original guide for training new language models

nlp on Marcin Zabłocki blog

Custom classifier on top of BERT-like Language Model - guide

TL;DR

Libraries and preparation

Understanding pre-trained language model input / outputs

Build custom classification head on top of LM

Preparing fast tokenizers to use with text dataset

Getting data for classification

Preparing reproducible training code with PyTorch Lightning

Hints and comments

Finding good starting learning rate for the model

Training the sentiment analysis classifier on PolEmo2.0

Using Trainer from PyTorch Lightning

Evaluating the trained model

Additional links and resources

Training RoBERTa from scratch - the missing guide

TL;DR

Getting the dataset

Preparing Polish Wikipedia dump

Extracting text from wikipedia

Prepare Wolne Lektury dump

Training new tokenizer

Adjusting the run_language_modeling.py script

Fast tokenizers support in run_language_modeling.py

Preparing training configuration

Training configuration for RoBERTa

Training your own RoBERTa language model

Training from scratch

Training from checkpoint (finetuning)

Training phase

Training summary for Polish RoBERTa a.k.a PolBERTa

Exporting the model

Using my pre-trained model

Follow up and TODO

Thanks to…

Additional links and resources

Adjusting the `run_language_modeling.py` script

Fast tokenizers support in `run_language_modeling.py`