The Dataset

In this chapter we will be using a subset of the Cross-lingual TRansfer Evaluation of Multilingual Encoders (XTREME)² benchmark called Wikiann³ or PAN-X. This dataset consists of Wikipedia articles in many languages, including the four most commonly spoken languages in Switzerland: German (62.9%), French (22.9%), Italian (8.4%), and English (5.9%). Each article is annotated with LOC (location), PER (person) and ORG (organization) tags in the “inside-outside-beginning” (IOB2) format, where a B- prefix indicates the beginning of an entity, and consecutive positions of the same entity are given an I- prefix. An O tag indicates that the token does not belong to any entity. For example, the following sentence

Jeff Dean is a computer scientist at Google in California

would be labeled in IOB2 format as shown in Table 4-1.

To load PAN-X with HuggingFace Datasets we first need to manually download the file AmazonPhotos.zip from XTREME’s Amazon Cloud Drive, and place it in a local directory (data in our example). Having done that, we can then load a PAN-X corpus using one of the two-letter ISO 639-1 language codes supported in the XTREME benchmark (see Table 5 of the paper for a list of the 40 available language codes). For example, to load the German corpus we use the “de” code as follows:

from datasets import load_dataset

load_dataset("xtreme", "PAN-X.de", data_dir="data")

DatasetDict({
    validation: Dataset({
        features: ['tokens', 'ner_tags', 'langs'],
        num_rows: 10000
    })
    test: Dataset({
        features: ['tokens', 'ner_tags', 'langs'],
        num_rows: 10000
    })
    train: Dataset({
        features: ['tokens', 'ner_tags', 'langs'],
        num_rows: 20000
    })
})

In this case, load_dataset returns a DatasetDict where each key corresponds to one of the splits, and each value is a Dataset object with features and num_rows attributes. To make a representative Swiss corpus, we’ll sample the German (de), French (fr), Italian (it), and English (en) corpora from PAN-X according to their spoken proportions. This will create a language imbalance that is very common in real-world datasets, where acquiring labeled examples in a minority language can be expensive due to the lack of domain experts who are fluent in that language.

To keep track of each language, let’s create a Python defaultdict that stores the language code as the key and a PAN-X corpus of type DatasetDict as the value:

from collections import defaultdict
from datasets import DatasetDict

langs = ["de", "fr", "it", "en"]
fracs = [0.629, 0.229, 0.084, 0.059]
# return a DatasetDict if a key doesn't exist
panx_ch = defaultdict(DatasetDict)

for lang, frac in zip(langs, fracs):
    # load monolingual corpus
    ds = load_dataset("xtreme", f"PAN-X.{lang}", data_dir="data")
    # shuffle and downsample each split according to spoken proportion
    for split in ds.keys():
        panx_ch[lang][split] = (
            ds[split]
            .shuffle(seed=0)
            .select(range(int(frac * ds[split].num_rows))))

Here we’ve used the Dataset.shuffle function to make sure we don’t accidentally bias our dataset splits, while Dataset.select allows us to downsample each corpus according to the values in fracs. Let’s have a look at how many examples we have per language in the training sets by accessing the Dataset.num_rows attribute:

import pandas as pd

pd.DataFrame({lang: [panx_ch[lang]["train"].num_rows] for lang in langs},
             index=["Number of training examples"])

By design, we have more examples in German than all other languages combined, so we’ll use it as a starting point from which to perform zero-shot cross-lingual transfer to French, Italian, and English. Let’s inspect one of the examples in the German corpus:

panx_ch["de"]["train"][0]

{'langs': ['de',
  'de',
  'de',
  'de',
  'de',
  'de',
  'de',
  'de',
  'de',
  'de',
  'de',
  'de'],
 'ner_tags': [0, 0, 0, 0, 5, 6, 0, 0, 5, 5, 6, 0],
 'tokens': ['2.000',
  'Einwohnern',
  'an',
  'der',
  'Danziger',
  'Bucht',
  'in',
  'der',
  'polnischen',
  'Woiwodschaft',
  'Pommern',
  '.']}

As with our previous encounters with Dataset objects, the keys of our example correspond to the column names of an Apache Arrow table, while the values denote the entry in each column. In particular, we see that the ner_tags column corresponds to the mapping of each entity to an integer. This is a bit cryptic to the human eye, so let’s create a new column with the familiar LOC, PER, and ORG tags. To do this, the first thing to notice is that our Dataset object has a features attribute that specifies the underlying data types associated with each column:

panx_ch["de"]["train"].features

{'tokens': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None),
 'ner_tags': Sequence(feature=ClassLabel(num_classes=7, names=['O', 'B-PER',
 > 'I-PER', 'B-ORG', 'I-ORG', 'B-LOC', 'I-LOC'], names_file=None, id=None),
 > length=-1, id=None),
 'langs': Sequence(feature=Value(dtype='string', id=None), length=-1, id=None)}

The Sequence class specifies that the field contains a list of features, which in the case of ner_tags corresponds to a list of ClassLabel features. Let’s pick out this feature from the training set as follows:

tags = panx_ch["de"]["train"].features["ner_tags"].feature
tags

ClassLabel(num_classes=7, names=['O', 'B-PER', 'I-PER', 'B-ORG', 'I-ORG',
 > 'B-LOC', 'I-LOC'], names_file=None, id=None)

One handy property of the ClassLabel feature is that it has conversion methods to convert from the class name to an integer and vice versa. For example, we can find the integer associated with the B-PER tag by using the ClassLabel.str2int function as follows:

tags.str2int("B-PER")

1

Similarly, we can map back from an integer to the corresponding class name:

tags.int2str(1)

'B-PER'

Let’s use the ClassLabel.int2str function to create a new column in our training set with class names for each tag. We’ll use the Dataset.map function to return a dict with the key corresponding to the new column name and the value as a list of class names:

def create_tag_names(batch):
    return {"ner_tags_str": [tags.int2str(idx) for idx in batch["ner_tags"]]}

panx_de = panx_ch["de"].map(create_tag_names)

Now that we have our tags in human-readable format, let’s see how the tokens and tags align for the first example in the training set:

de_example = panx_de["train"][0]
df = pd.DataFrame([de_example["tokens"], de_example["ner_tags_str"]],
                  ['Tokens', 'Tags'])
display_df(df, header=None)

The presence of the LOC tags make sense since the sentence “2,000 Einwohnern an der Danziger Bucht in der polnischen Woiwodschaft Pommern” means “2,000 inhabitants at the Gdansk Bay in the Polish voivodeship of Pomerania” in English, and Gdansk Bay is a bay in the Baltic sea, while “voivodeship” corresponds to a state in Poland.

As a sanity check that we don’t have any unusual imbalance in the tags, let’s calculate the frequencies of each entity across each split:

from itertools import chain
from collections import Counter

split2freqs = {}

for split in panx_de.keys():
    tag_names = []
    for row in panx_de[split]["ner_tags_str"]:
        tag_names.append([t.split("-")[1] for t in row if t.startswith("B")])
    split2freqs[split] = Counter(chain.from_iterable(tag_names))

pd.DataFrame.from_dict(split2freqs, orient="index")

This looks good - the distribution of the PER, LOC, and ORG frequencies are roughly the same for each split, so the validation and test sets should provide a good measure of our NER tagger’s ability to generalize. Next, let’s look at a few popular multilingual Transformers and how they can be adapted to tackle our NER task.

Multilingual Transformers

Multilingual Transformers involve similar architectures and training procedures as their monolingual counterparts, except that the corpus used for pretraining consists of documents in many languages. A remarkable feature of this approach is that despite receiving no explicit information to differentiate among the languages, the resulting linguistic representations are able to generalize well across languages for a variety of downstream tasks. In some cases, this ability to perform cross-lingual transfer can produce results that are competitive with monolingual models, which circumvents the need to train one model per language!

To measure the progress of cross-lingual transfer for NER, the CoNLL-2002 and CoNLL-2003 datasets are often used as a benchmark for English, Dutch, Spanish, and German. This benchmark consists of news articles annotated with the same LOC, PER, and ORG categories as PAN-X, but contains an additional MISC label for miscellaneous entities that do not belong to the previous three groups. Multilingual Transformer models are then evaluated in three different ways:

en

Fine-tune on the English training data and then evaluate on each language’s test set.

each

Fine-tune and evaluate on monolingual training data to measure per-language performance.

all

Fine-tune on all the training data to evaluate multilingual learning.

We will adopt a similar evaluation strategy for our NER task and we’ll use XLM-RoBERTa (or XLM-R for short) which, as of this book’s writing, is the current state-of-the-art Transformer model for multilingual applications. But first, let’s take a look at the two models that inspired its development: mBERT and XLM.

XLM

In the Cross-lingual Language Model Pretraining paper, Guillaume Lample and Alexis Conneau from Facebook AI Research investigated three pretraining objectives for cross-lingual language (XLM) models. One of these objectives is the masked language modeling (MLM) objective from BERT, but instead of receiving complete sentences as input, XLM receives sentences that can be truncated arbitrarily (there is also no next-sentence prediction task). To increase the number of tokens associated with low-resource languages, the sentences are sampled from a monolingual corpus ${\left\{C_i\right\}}_{i=1,...,N}$ according to the multinomial distribution, with probabilities

$$q_i=\frac{p_i^{\alpha }}{{\sum }_{j=1}^N{p}_j^{\alpha }}\text{where}{p}_i=\frac{n_i}{{\sum }_{k=1}^N{n}_k}$$

and $\alpha =0.5$ and $n_i$ is the number of sentences in a monolingual corpus $C_i$. Another difference from BERT is the use of Byte-Pair-Encoding instead of WordPiece for tokenization, which the authors observe improves the alignment of the language embeddings across languages. The paper also introduces translation language modelling (TLM) as a new pretraining objective, which concatenates pairs of sentences from two languages and randomly masks the tokens as in MLM. To predict a masked token in one language, the model can attend to tokens in the translated pair which encourages the alignment of the cross-lingual representations. A comparison of the two methods is shown in Figure 4-1.

Figure 4-1. The MLM (top) and TLM (bottom) pretraining objectives of XLM. Figure from the XLM paper.

Note

There are several variants of XLM based on the choice of pretraining objective and number of languages to be trained on. For the purposes of this discussion, we’ll use XLM to denote the model trained on the same 100 languages used for mBERT.

XLM-R

Like its predecessors, XLM-R uses MLM as a pretraining objective for 100 languages, but, as shown in Figure 4-2, is distinguished by the huge size of the corpus used for pretraining: Wikipedia dumps for each language and 2.5 terabytes of Common Crawl data from the web. This corpus is several orders of magnitude larger than the ones used in previous models and provides a significant boost in signal for low-resource languages like Burmese and Swahili, where only a small number of Wikipedia articles exist.

Figure 4-2. Amount of data for the languages that appear in both the Wiki-100 corpus used for mBERT and XLM, and the CommonCrawl corpus used for XLM-R. Figure from the XLM paper.

The RoBERTa part of the model’s name refers to the fact that the pretraining approach is the same as monolingual RoBERTa models. In the RoBERTa paper,⁵ the authors improved on several aspects of BERT, in particular by removing the next sentence prediction task altogether. XLM-R also drops the language embeddings used in XLM and uses SentencePiece⁶ to tokenize the raw texts directly. Besides its multilingual nature, a notable difference between XLM-R and RoBERTa is the size of the respective vocabularies: 250,000 tokens versus 55,000!

The Table 4-2 summarizes the main architectural differences between all the multilingual Transformers.

Table 4-2. Summary of multilingual models.

Model	Languages	Tokenizer	Layers	Hidden States	Attention Heads	Vocabulary Size	Parameters
mBERT	104	WordPiece	12	768	12	110k	172M
XLM	100	BytePairEncoding	16	1280	16	200k	570M
XLM-R (Base)	100	SentencePiece	12	768	12	250k	270M
XLM-R (Large)	100	SentencePiece	24	1024	16	250k	550M

The performance of mBERT and XLM-R on the CoNLL benchmark is also shown in Figure 4-3. We see that when trained on all the languages, the XLM-R models significantly outperform mBERT and earlier state-of-the-art approaches.

Figure 4-3. F1-scores on the CoNLL benchmark for NER. Figure from the XLM paper.

From this research it becomes apparent that XLM-R is the best choice for multilingual NER. In the next section we explore how to fine-tune XLM-R for this task on a new dataset.

Training a Named Entity Recognition Tagger

In Chapter 1, we saw that for text classification, BERT uses the special [CLS] token to represent an entire sequence of text. As shown in the left diagram of Figure 4-4, this representation is then fed through a fully-connected or dense layer to output the distribution of all the discrete label values. BERT and other encoder Transformers take a similar approach for NER, except that the representation of every input token is fed into the same fully-connected layer to output the entity of the token. For this reason, NER is often framed as a token classification task and the process looks something like the right diagram of Figure 4-4.

So far, so good, but how should we handle subwords in a token classification task? For example, the last name “Sparrow” in Figure 4-4 is tokenized by WordPiece into the subwords “Spa” and “##rrow”, so which one (or both) should be assigned the I-PER label?

In the BERT paper,⁷ the authors used the representation from first subword (i.e. “Spa” in our example) and this is the convention we’ll adopt here. Although we could have chosen to include the representation from the “##rrow” subword by assigning it a copy of the I-LOC label, this introduces extra complexity when subwords are associated with a B- entity because then we need to copy these tags and this violates the IOB2 format.

Figure 4-4. Fine-tuning BERT for text classification (left) and named entity recognition (right).

Fortunately, all this intuition from BERT carries over to XLM-R since the architecture is based on RoBERTa, which is identical to BERT! However, there are some slight differences, especially around the choice of tokenizer. Let’s see how the two differ.

from transformers import AutoTokenizer

bert_model_name = "bert-base-cased"
xlmr_model_name = "xlm-roberta-base"
bert_tokenizer = AutoTokenizer.from_pretrained(bert_model_name)
xlmr_tokenizer = AutoTokenizer.from_pretrained(xlmr_model_name)

text = "Jack Sparrow loves New York!"
bert_tokens = bert_tokenizer(text).tokens()
xlmr_tokens = xlmr_tokenizer(text).tokens()

"".join(xlmr_tokens).replace("_", " ")

'<s> Jack Sparrow loves New York!</s>'

The Anatomy of the Transformers Model Class

As we have seen in previous chapters, the Transformers library is organized around dedicated classes for each architecture and task. The list of supported tasks can be found in the Transformers documentation, and as of this book’s writing includes

• Sequence classification

• Extractive question answering

• Language modeling

• Named entity recognition

• Summarization

• Translation

and the associated classes are named according to a ModelNameForTask convention. Most of the time, we load these models using the ModelNameForTask.from_pretrained function and since the architecture can usually be guessed from the name alone (e.g. bert-base-uncased), Transformers provides a convenient set of AutoClasses to automatically load the relevant configuration, vocabulary, or weights. In practice, these AutoClasses are extremely useful because it means that we can switch to a completely different architecture in our experiments by simply changing the model name!

However, this approach has its limitations, and to motivate going deeper in the Transformers API consider the following scenario. Suppose you work for a consulting company that is engaged with many customer projects each year. By studying how these projects evolve, you’ve noticed that the initial estimates for person-months, number of required people, and the total project timespan are extremely inaccurate. After thinking about this problem, you have the idea that feeding the written project descriptions to a Transformer model might yield much better estimates of these quantities.

So you set up a meeting with your boss and, with an artfully crafted Powerpoint presentation, you pitch that you could increase the accuracy of the project estimates and thus increase the efficiency of the staff and revenue by making more accurate offers. Impressed with your colorful presentation and talk of efficiency and profits, your boss generously agrees to give you one week to build a proof-of-concept. Happy with the outcome, you start working straight away and decide that the only thing you need is regression model to predict the three variables (person-months, number of people, and timespan). You fire up your favorite GPU and open a notebook. You execute from transformers import BertForRegression and color escapes your face as dreaded red color fills your screen: ImportError: cannot import name 'BertForRegression'. Oh no, there is no BERT model for regression! How should you complete the project in one week if you have to implement the whole model yourself?! Where should you even start?

Don’t panic! The Transformers library is designed to enable you the easily extend existing models for your specific use-case. With it you have access to various utilities such as loading weights of pretrained models or task specific helper functions. This lets you build custom models for specific objectives with very little overhead.

Bodies and Heads

The main concept that makes Transformers so versatile is the split of the architecture into a body and head. We have already seen that when we switch from the pretraining task to the downstream task, we need to replace the last layer of the model with one that is suitable for the task. This last layer is called the model head and is the part that is task specific. The rest of the model is called the body and includes the token embeddings and Transformer layers that are task agnostic. This structure is reflected in the Transformers code as well: The body of a model is implemented in a class such as BertModel or GPT2Model that returns the hidden states of the last layer. Task specific models such as BertForMaskedLM or BertForSequenceClassification use the base model and add the necessary head on top of the hidden states as shown in figure Figure 4-5.

Figure 4-5. The BertModel class only contains the body of the model while the BertForTask classes combine the body with a dedicated head for a given task.

import torch.nn as nn
from transformers import XLMRobertaConfig
from transformers.models.roberta.modeling_roberta import (
    RobertaModel, RobertaPreTrainedModel)

class XLMRobertaForTokenClassification(RobertaPreTrainedModel):
    config_class = XLMRobertaConfig

    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels
        # load model body
        self.roberta = RobertaModel(config, add_pooling_layer=False)
        # setup token classification head
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
        # load and initialize weights
        self.init_weights()

    def forward(self, input_ids=None, attention_mask=None,
                token_type_ids=None, labels=None, **kwargs):
        pass

from transformers.modeling_outputs import TokenClassifierOutput

def forward(self, input_ids=None, attention_mask=None, token_type_ids=None,
            labels=None, **kwargs):
    # use model body to get encoder representations
    outputs = self.roberta(input_ids, attention_mask=attention_mask,
                           token_type_ids=token_type_ids, **kwargs)
    # apply classifier to encoder representation
    sequence_output = self.dropout(outputs[0])
    logits = self.classifier(sequence_output)
    # calculate losses
    loss = None
    if labels is not None:
        loss_fct = nn.CrossEntropyLoss()
        loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
    # return model output object
    return TokenClassifierOutput(loss=loss, logits=logits,
                                 hidden_states=outputs.hidden_states,
                                 attentions=outputs.attentions)

XLMRobertaForTokenClassification.forward = forward

index2tag = {idx: tag for idx, tag in enumerate(tags.names)}
tag2index = {tag: idx for idx, tag in enumerate(tags.names)}

from transformers import AutoConfig

xlmr_config = AutoConfig.from_pretrained(xlmr_model_name,
                                         num_labels=tags.num_classes,
                                         id2label=index2tag, label2id=tag2index)

import torch

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
xlmr_model = (XLMRobertaForTokenClassification
              .from_pretrained(xlmr_model_name, config=xlmr_config)
              .to(device))

input_ids = xlmr_tokenizer.encode(text, return_tensors="pt").to(device)

df = pd.DataFrame([xlmr_tokenizer.all_special_tokens,
                   xlmr_tokenizer.all_special_ids],
                  index=["Special Token", "Special Token ID"])
display_df(df, header=None)

outputs = xlmr_model(input_ids).logits
predictions = torch.argmax(outputs, dim=-1)
print(f"Number of tokens in sequence: {len(xlmr_tokens)}")
print(f"Shape of outputs: {outputs.shape}")

Number of tokens in sequence: 10
Shape of outputs: torch.Size([1, 10, 7])

preds = [tags.names[p] for p in predictions[0].cpu().numpy()]
df = pd.DataFrame([xlmr_tokens, preds], index=["Tokens", "Tags"])
display_df(df, header=None)

def tag_text(text, tags, model, tokenizer):
    # get tokens with special characters
    tokens = tokenizer.tokenize(tokenizer.decode(tokenizer.encode(text)))
    # encode the sequence into IDs
    inputs = tokenizer.encode(text, return_tensors="pt").to(device)
    # get predictions as distribution over 7 possible classes
    outputs = model(inputs)[0]
    # take argmax to get most likely class per token
    predictions = torch.argmax(outputs, dim=2)
    # convert to DataFrame
    preds = [tags.names[p] for p in predictions[0].cpu().numpy()]
    df = pd.DataFrame([tokens, preds], index=["Tokens", "Tags"])
    display_df(df, header=None)

function(examples: Dict[str, List]) -> Dict[str, List]

words, labels = de_example["tokens"], de_example["ner_tags"]

tokenized_input = xlmr_tokenizer(de_example["tokens"], is_split_into_words=True)
tokens = xlmr_tokenizer.convert_ids_to_tokens(tokenized_input["input_ids"])

word_ids = tokenized_input.word_ids()

previous_word_idx = None
label_ids = []

for word_idx in word_ids:
    if word_idx is None:
        label_ids.append(-100)
    elif word_idx != previous_word_idx:
        label_ids.append(labels[word_idx])
    else:
        label_ids.append(-100)
    previous_word_idx = word_idx

def tokenize_and_align_labels(examples):
    tokenized_inputs = xlmr_tokenizer(examples["tokens"], truncation=True,
                                      is_split_into_words=True)
    labels = []

    for idx, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=idx)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None or word_idx == previous_word_idx:
                label_ids.append(-100)
            else:
                label_ids.append(label[word_idx])
            previous_word_idx = word_idx

        labels.append(label_ids)

    tokenized_inputs["labels"] = labels
    return tokenized_inputs

single_sample = panx_de["train"].select(range(1))
single_sample_encoded = single_sample.map(tokenize_and_align_labels,
                                          batched=True)

print(" ".join(token for token in single_sample[0]["tokens"]))
print(xlmr_tokenizer.decode(single_sample_encoded["input_ids"][0]))

2.000 Einwohnern an der Danziger Bucht in der polnischen Woiwodschaft Pommern .
<s> 2.000 Einwohnern an der Danziger Bucht in der polnischen Woiwodschaft
 > Pommern.</s>

original_labels = single_sample["ner_tags_str"][0]
reconstructed_labels = [index2tag[idx] for idx
                        in single_sample_encoded["labels"][0] if idx != -100]

def encode_panx_dataset(corpus):
    return corpus.map(tokenize_and_align_labels, batched=True,
                      remove_columns=['langs', 'ner_tags', 'tokens'])

panx_de_encoded = encode_panx_dataset(panx_ch["de"])
panx_de_encoded["train"]

Dataset({
    features: ['attention_mask', 'input_ids', 'labels'],
    num_rows: 12580
})

from seqeval.metrics import classification_report

y_true = [["O", "O", "O", "B-MISC", "I-MISC", "I-MISC", "O"],
          ["B-PER", "I-PER", "O"]]
y_pred = [["O", "O", "B-MISC", "I-MISC", "I-MISC", "I-MISC", "O"],
          ["B-PER", "I-PER", "O"]]
print(classification_report(y_true, y_pred))

              precision    recall  f1-score   support

        MISC       0.00      0.00      0.00         1
         PER       1.00      1.00      1.00         1

   micro avg       0.50      0.50      0.50         2
   macro avg       0.50      0.50      0.50         2
weighted avg       0.50      0.50      0.50         2

import numpy as np

def align_predictions(predictions, label_ids):
    preds = np.argmax(predictions, axis=2)
    batch_size, seq_len = preds.shape
    labels_list, preds_list = [], []

    for batch_idx in range(batch_size):
        example_labels, example_preds = [], []
        for seq_idx in range(seq_len):
            # ignore label IDs = -100
            if label_ids[batch_idx, seq_idx] != -100:
                example_labels.append(index2tag[label_ids[batch_idx][seq_idx]])
                example_preds.append(index2tag[preds[batch_idx][seq_idx]])

        labels_list.append(example_labels)
        preds_list.append(example_preds)

    return preds_list, labels_list

from transformers import TrainingArguments

num_epochs = 3
batch_size = 24
logging_steps = len(panx_de_encoded["train"]) // batch_size
training_args = TrainingArguments(output_dir="results",
                                  num_train_epochs=num_epochs,
                                  per_device_train_batch_size=batch_size,
                                  per_device_eval_batch_size=batch_size,
                                  evaluation_strategy="epoch", save_steps=1e6,
                                  weight_decay=0.01, disable_tqdm=False,
                                  logging_steps=logging_steps)

from seqeval.metrics import f1_score

def compute_metrics(eval_pred):
    y_pred, y_true = align_predictions(eval_pred.predictions,
                                       eval_pred.label_ids)
    return {"f1": f1_score(y_true, y_pred)}

from transformers import DataCollatorForTokenClassification

data_collator = DataCollatorForTokenClassification(xlmr_tokenizer)

from transformers import Trainer

trainer = Trainer(model=xlmr_model, args=training_args,
                  data_collator=data_collator, compute_metrics=compute_metrics,
                  train_dataset=panx_de_encoded["train"],
                  eval_dataset=panx_de_encoded["validation"],
                  tokenizer=xlmr_tokenizer)

trainer.train();

trainer.save_model("models/xlm-roberta-base-finetuned-panx-de")

text_de = "Jeff Dean ist ein Informatiker bei Google in Kalifornien"
tag_text(text_de, tags, trainer.model, xlmr_tokenizer)

Error Analysis

Before we dive deeper into the multilingual aspects of XLM-R let’s take a minute to investigate the errors of our model. As we saw in Chapter 1, a thorough error analysis of your model is one of the most important aspects when training and debugging Transformers (and machine learning models in general). There are several failure modes where it might look like the model is performing well while in practice it has some serious flaws. Examples where Transformers can fail include:

• We can accidentally mask too many tokens and also mask some of our labels to get a really promising loss drop.

• The compute_metrics function can have a bug that overestimates the true performance.

• We might include the zero class or O entity in NER as a normal class which will heavily skew the accuracy and $F_1$-score since it is the majority class by a large margin.

When the model performs much worse than expected, looking at the errors can also yield useful insights and reveal bugs which would be hard to spot by just looking at the code. Even if the model performs well and there are no bugs in the code, error analysis is still a useful tool to understand the strength and weaknesses of the model. These are aspects we always need to keep in mind when we deploy a model in a production environment.

We will again use one of the most powerful tools at our disposal which is to look at the validation examples with highest loss. We can reuse much of the function we built to analyze the sequence classification model in Chapter 1 but in contrast we now calculate a loss per token in the sample sequence.

Let’s first load our fine-tuned model

xlmr_model = (XLMRobertaForTokenClassification
              .from_pretrained("models/xlm-roberta-base-finetuned-panx-de")
              .to(device))

and define a function that we can iterate over the validation set:

from torch.nn.functional import cross_entropy

def forward_pass_with_label(batch):
    # convert dict of lists to list of dicts
    features = [dict(zip(batch, t)) for t in zip(*batch.values())]
    # pad inputs and labels
    batch = data_collator(features)
    input_ids = batch["input_ids"].to(device)
    attention_mask = batch["attention_mask"].to(device)
    labels = batch["labels"].to(device)

    with torch.no_grad():
        output = xlmr_model(input_ids, attention_mask)
        batch["predicted_label"] = torch.argmax(output.logits, axis=-1)

    loss = cross_entropy(output.logits.view(-1, 7),
                         labels.view(-1), reduction="none")
    loss = loss.view(len(input_ids), -1)
    batch["loss"] = loss

    # datasets requires list of NumPy array data types
    for k, v in batch.items():
        batch[k] = v.cpu().numpy()

    return batch

We now apply this function to the whole validation set using Dataset.map and load all the data into a DataFrame for further analysis:

valid_set = panx_de_encoded["validation"]
valid_set = valid_set.map(forward_pass_with_label, batched=True, batch_size=32)
valid_set.set_format("pandas")
df = valid_set[:]

The tokens and the labels are still encoded with their IDs, so let’s map the tokens and labels back to strings to make it easier to read the results. For the padding tokens with label -100 we assign a special label IGN so we can filter them later:

index2tag[-100] = "IGN"
df["input_tokens"] = df["input_ids"].apply(
    lambda x: xlmr_tokenizer.convert_ids_to_tokens(x))
df["predicted_label"] = df["predicted_label"].apply(
    lambda x: [index2tag[i] for i in x])
df["labels"] = df["labels"].apply(lambda x: [index2tag[i] for i in x])

Each column contains a list of tokens, labels, predicted labels, and so on for each sample. Let’s have a look at the tokens individually by unpacking these lists. The pandas.Series.explode function allows us to do exactly that in one line by creating a row for each element in the original rows list. Since all the lists in one row have the same length we can do this in parallel for all columns. We also drop the padding tokens since their loss is zero anyway:

df_tokens = df.apply(pd.Series.explode)
df_tokens = df_tokens.query("labels != 'IGN'")
df_tokens["loss"] = df_tokens["loss"].astype(float)

With the data in this shape we can now group it by the input tokens and aggregate the losses for each token with the count, mean, and sum. Finally, we sort the aggregated data by the sum of the losses and see which tokens have accumulated most loss in the validation set:

(
    df_tokens.groupby("input_tokens")[["loss"]]
    .agg(["count", "mean", "sum"])
    .droplevel(level=0, axis=1)  # get rid of multi-level columns
    .sort_values(by="sum", ascending=False)
    .reset_index()
    .head(20)
)

We can observe several patterns in this list:

• The whitespace token has the highest total loss which is not surprising since it is also the most common token in the list. On average it seems to be well below most tokens in the list.

• Words like in, von, der, and und appear relatively frequently. They often appear together with named entities and are sometimes part of them which explains why the model might mix them up.

• Parentheses, slashes, and capital letters at the beginning of words are rarer but have a relatively high average loss. We will investigate them further.

• At the end of list we see some subwords that appear rarely but have a very high average loss. For example _West shows that these tokens appear in almost any class, and thus pose a classification challenge to the model:

df_tokens.query("input_tokens == '_West'")["labels"].value_counts()

O        23
B-LOC     6
I-ORG     6
B-ORG     5
I-LOC     4
I-PER     3
B-PER     1
Name: labels, dtype: int64

We can also group the label IDs and look at the losses for each class. We see that B-ORG has the highest average loss which means that determining the beginning of an organization poses a challenge to our model:

(
    df_tokens.groupby("labels")[["loss"]]
    .agg(["count", "mean", "sum"])
    .droplevel(level=0, axis=1)
    .sort_values(by="mean", ascending=False)
    .reset_index()
)

We can break this down further by plotting the confusion matrix of the token classification, where we see that the beginning of an organization is often confused with the subsequent I-ORG token:

plot_confusion_matrix(
    df_tokens["labels"], df_tokens["predicted_label"], tags.names
)

Now that we’ve examined the errors at the token level let’s move on and look at sequences with high losses. For this calculation, we revisit “unexploded” DataFrame and calculate the total loss by summing over the loss per token. To do this let’s first write a function that helps us display the token sequence with the labels and the losses:

def display_samples(df):
    for _, row in df.iterrows():
        labels, preds, tokens, losses = [], [], [], []
        for i, mask in enumerate(row["attention_mask"]):
            if mask == 1:
                labels.append(row["labels"][i])
                preds.append(row["predicted_label"][i])
                tokens.append(row["input_tokens"][i])
                losses.append(f"{row['loss'][i]:.2f}")
        df_tmp = pd.DataFrame({"tokens": tokens, "labels": labels,
                               "preds": preds, "losses": losses}).T
        display_df(df_tmp, header=None, max_cols=10)

df["total_loss"] = df["loss"].apply(sum)
display_samples(df.sort_values(by="total_loss", ascending=False).head(3))

It is apparent that something is wrong with the labels of these samples; for example, the United Nations is labeled as a person! It turns out the annotations for the Wikiann dataset were generated through an automated process. Such annotations are often referred to as “silver-standard” (in contrast to the “gold-standard” of human-generated annotations), and it is no surprise that there are cases where the automated approach failed to produce sensible labels. However, such failure modes are not unique to automatic approaches; even when humans carefully annotate data, mistakes can occur when the concentration of the annotators fades or they simply misunderstood the sentence.

Another thing we noticed when looking at the tokens with the most loss were the parentheses and slashes. Lets look at a few examples of sequences with an opening parenthesis:

display_samples(df.loc[df["input_tokens"].apply(lambda x: "_(" in x)].head(3))

Since Wikiann is a dataset created from Wikipedia, we can see that the entities contain parentheses from the introductory sentence of each article where the name of the article is described. In the first example, the parenthesis simply states that the Hama is an “Unternehmen” or company in English. In general we would not include the parenthesis and its content as part of the named entity but this seems to be the way the automatic extraction annotated the documents. In the other examples the parenthesis contains a geographic specification. While this is indeed a location as well we might want disconnect them from the original location in the annotations. These are important details to know when we roll-out the model since it might have implication on the downstream performance of the whole pipeline the model is part of.

With a relatively simple analysis we found weaknesses in both our model and the dataset. In a real use-case we would iterate on this step and clean up the dataset, re-train the model and analyze the new errors until we are satisfied with the performance.

Now we analysed the errors on a single language but we are also interested in the performance across the languages. In the next section we perform some experiments to see how well the cross-lingual transfer in XLM-R works.

Evaluating Cross-Lingual Transfer

Now that we have fine-tuned XLM-R on German, we can evaluate its ability to transfer to other languages via the Trainer.predict function that generates predictions on Dataset objects. For example, to get the predictions on the validation set we can run the following:

panx_de_encoded["validation"].reset_format()
preds_valid = trainer.predict(panx_de_encoded["validation"])

The output of Trainer.predict is a trainer_utils.PredictionOutput object which contains arrays of predictions and label_ids, along with the metrics we passed to the trainer. For example, the metrics on the validation set can be accessed as follows:

preds_valid.metrics

{'eval_loss': 0.13674472272396088, 'eval_f1': 0.863226026230625}

The predictions and label IDs are not quite in a form suitable for seqeval’s classification report, so let’s align them using our align_predictions function and print out the classification report with the following function:

def generate_report(trainer, dataset):
    preds = trainer.predict(dataset)
    preds_list, label_list = align_predictions(
        preds.predictions, preds.label_ids)
    print(classification_report(label_list, preds_list, digits=4))
    return preds.metrics["eval_f1"]

To keep track of our performance per language, our function also returns the micro-averaged $F_1$-score. Let’s use this function to examine the performance on the test set and keep track of our scores in a dict:

f1_scores = defaultdict(dict)
f1_scores["de"]["de"] = generate_report(trainer, panx_de_encoded["test"])

              precision    recall  f1-score   support

         LOC     0.8596    0.8950    0.8769      3180
         ORG     0.7979    0.7765    0.7871      2573
         PER     0.9162    0.9225    0.9194      3071

   micro avg     0.8619    0.8700    0.8659      8824
   macro avg     0.8579    0.8647    0.8611      8824
weighted avg     0.8613    0.8700    0.8655      8824

These are pretty good results for a NER task. Our metrics are in the ballpark of 85% and we can see that the model seems to struggle the most on the ORG entities, probably because ORG entities are the least common in the training data and many organization names are rare in XLM-R’s vocabulary. How about on other languages? To warm up, let’s see how our model fine-tuned on German fares on French:

text_fr = "Jeff Dean est informaticien chez Google en Californie"
tag_text(text_fr, tags, trainer.model, xlmr_tokenizer)

Not bad! Although the name and organization are the same in both languages, the model did manage to correctly label the French translation of “Kalifornien”. Next, let’s quantify how well our German model fares on the whole French test set by writing a simple function that encodes a dataset and generates the classification report on it:

def evaluate_zero_shot_performance(lang, trainer):
    panx_ds = encode_panx_dataset(panx_ch[lang])
    return generate_report(trainer, panx_ds["test"])

f1_scores["de"]["fr"] = evaluate_zero_shot_performance("fr", trainer)

              precision    recall  f1-score   support

         LOC     0.7239    0.7239    0.7239      1130
         ORG     0.6371    0.6407    0.6389       885
         PER     0.7207    0.7703    0.7447      1045

   micro avg     0.6981    0.7157    0.7068      3060
   macro avg     0.6939    0.7116    0.7025      3060
weighted avg     0.6977    0.7157    0.7064      3060

Although we see a drop of about 15 points in the micro-averaged metrics, remember that our model has not seen a single labeled French example! In general, the size of the performance drop is related to how “far away” the languages are from each other. Although German and French are grouped as Indo-European languages, they technically belong to the different languages families of “Germanic” and “Romance” respectively.

Next, let’s evaluate the performance on Italian. Since Italian is also a Romance language, we expect to get a similar result as we found on French:

f1_scores["de"]["it"] = evaluate_zero_shot_performance("it", trainer)

              precision    recall  f1-score   support

         LOC     0.7143    0.7042    0.7092       426
         ORG     0.5961    0.6185    0.6071       346
         PER     0.6736    0.8094    0.7353       362

   micro avg     0.6647    0.7116    0.6874      1134
   macro avg     0.6613    0.7107    0.6839      1134
weighted avg     0.6652    0.7116    0.6864      1134

Indeed, our expectations are borne out by the macro-averaged metrics. Finally, let’s examine the performance on English which belongs to the Germanic language family:

f1_scores["de"]["en"] = evaluate_zero_shot_performance("en", trainer)

              precision    recall  f1-score   support

         LOC     0.4847    0.6148    0.5421       283
         ORG     0.6034    0.6449    0.6235       276
         PER     0.6512    0.6932    0.6716       264

   micro avg     0.5722    0.6501    0.6086       823
   macro avg     0.5798    0.6510    0.6124       823
weighted avg     0.5779    0.6501    0.6109       823

Surprisingly, our model fares worst on English even though we might intuitively expect German to be more similar than French. Let’s next examine the trade-offs between zero-shot cross-lingual transfer and fine-tuning directly on the target language.

When Does Zero-Shot Transfer Make Sense?

So far we’ve seen that fine-tuning XLM-R on the German corpus yields an $F_1$-score of around 85%, and without any additional training is able to achieve modest performance on the other languages in our corpus. The question is: how good are these results and how do they compare against an XLM-R model fine-tuned on a monolingual corpus?

In this section we will explore this question for the French corpus by fine-tuning XLM-R on training sets of increasing size. By tracking the performance this way, we can determine at which point zero-shot cross-lingual transfer is superior, which in practice can be useful for guiding decisions about whether to collect more labeled data.

Since we want to train several models, we’ll use the model_init feature of the Trainer class so that we can instantiate a fresh model with each call to Trainer.train:

def model_init():
    return (XLMRobertaForTokenClassification
            .from_pretrained(xlmr_model_name, config=xlmr_config)
            .to(device))

For simplicity, we’ll also keep the same hyperparameters from the fine-tuning run on the German corpus, except that we’ll tweak TrainingArguments.logging_steps to account for the changing training set sizes. We can wrap this altogether in a simple function that takes a DatasetDict object corresponding to a monolingual corpus, downsamples it by num_samples, and fine-tunes XLM-R on that sample to return the metrics from the best epoch:

def train_on_subset(dataset, num_samples):
    train_ds = dataset["train"].shuffle(seed=42).select(range(num_samples))
    valid_ds = dataset["validation"]
    test_ds = dataset["test"]
    training_args.logging_steps = len(train_ds) // batch_size
    trainer = Trainer(model_init=model_init, args=training_args,
        data_collator=data_collator, compute_metrics=compute_metrics,
        train_dataset=train_ds, eval_dataset=valid_ds, tokenizer=xlmr_tokenizer)

    trainer.train()
    metrics = trainer.predict(test_ds).metrics
    return pd.DataFrame.from_dict(
        {"num_samples": [len(train_ds)], "f1_score": [metrics["eval_f1"]]})

As we did with fine-tuning on the German corpus, we also need to encode the French corpus into input IDs, attention masks, and label IDs:

panx_fr_encoded = encode_panx_dataset(panx_ch["fr"])

Next let’s check that our function works by running it on a small training set of 250 examples:

metrics_df = train_on_subset(panx_fr_encoded, 250)
metrics_df

We can see that with only 250 examples, fine-tuning on French under-performs the zero-shot transfer from German by a large margin. Let’s now increase our training set sizes to 500, 1,000, 2,000, and 4,000 examples to get an idea of how the performance increases:

for num_samples in [500, 1000, 2000, 4000]:
    metrics_df = metrics_df.append(
        train_on_subset(panx_fr_encoded, num_samples), ignore_index=True)

We can compare how fine-tuning on French samples compares to zero-shot cross-lingual transfer from German by plotting the $F_1$-scores on the test set as a function of increasing training set size:

From the plot we can see that zero-shot transfer remains competitive until about 750 training examples, after which fine-tuning on French reaches a similar level of performance to what we obtained when fine-tuning on German. Nevertheless, this result is not to be sniffed at! In our experience, getting domain experts to label even hundreds of documents can be costly; especially for NER where the labeling process is fine-grained and time consuming.

There is one final technique we can try to evaluate multilingual learning: fine-tune on multiple languages at once! Let’s see how we can do this in the next section.

Fine-tuning on Multiple Languages at Once

So far we’ve seen that zero-shot cross-lingual transfer from German to French or Italian produces a drop of around 15 points in performance. One way to mitigate this is by fine-tuning on multiple languages at the same time! To see what type of gains we can get, let’s first use the concatenate_datasets function from Datasets to concatenate the German and French corpora together:

from datasets import concatenate_datasets

def concatenate_splits(corpora):
    multi_corpus = DatasetDict()
    for split in corpora[0].keys():
        multi_corpus[split] = concatenate_datasets(
            [corpus[split] for corpus in corpora]).shuffle(seed=42)
    return multi_corpus

panx_de_fr_encoded = concatenate_splits([panx_de_encoded, panx_fr_encoded])

For training, we’ll also use the same hyperparameters from the previous sections, so we can simply update the logging steps, model, and datasets in the trainer:

training_args.logging_steps = len(panx_de_fr_encoded["train"]) // batch_size
trainer.train_dataset = panx_de_fr_encoded["train"]
trainer.eval_dataset = panx_de_fr_encoded["validation"]
trainer.train();

This model gives a similar $F_1$-score to our first model that was fine-tuned on German. How does it fare with cross-lingual transfer? First, let’s examine the performance on Italian:

evaluate_zero_shot_performance("it", trainer);

              precision    recall  f1-score   support

         LOC     0.7143    0.7042    0.7092       426
         ORG     0.5961    0.6185    0.6071       346
         PER     0.6736    0.8094    0.7353       362

   micro avg     0.6647    0.7116    0.6874      1134
   macro avg     0.6613    0.7107    0.6839      1134
weighted avg     0.6652    0.7116    0.6864      1134

Wow, this is a 10 point improvement compared to our German model which scored an $F_1$-score of around 70% on Italian! Given the similarities between French and Italian, this is perhaps not so surprising; how does the model perform on English?

evaluate_zero_shot_performance("en", trainer);

              precision    recall  f1-score   support

         LOC     0.4847    0.6148    0.5421       283
         ORG     0.6034    0.6449    0.6235       276
         PER     0.6512    0.6932    0.6716       264

   micro avg     0.5722    0.6501    0.6086       823
   macro avg     0.5798    0.6510    0.6124       823
weighted avg     0.5779    0.6501    0.6109       823

Here we also have a significant boost in zero-shot performance by 7-8 points, with most of the gain coming from a dramatic improvement of the PER tokens! Apparently the Norman conquest of 1066 left a long-lasting effect on the English language.

Let’s round out our analysis by comparing the performance of fine-tuning on each language separately against multilingual learning on all the corpora. Since we have already fine-tuned on the German corpus, we can fine-tune on the remaining languages with our train_on_subset function, but where num_samples is equal to the number of examples in the training set:

corpora = [panx_de_encoded]

# exclude German from iteration
for lang in langs[1:]:
    # fine-tune on monolingual corpus
    ds_encoded = encode_panx_dataset(panx_ch[lang])
    metrics = train_on_subset(ds_encoded, ds_encoded["train"].num_rows)
    # collect F1-scores in common dict
    f1_scores[lang][lang] = metrics["f1_score"][0]
    # add monolingual corpus to list of corpora to concatenate
    corpora.append(ds_encoded)

Now that we’ve fine-tuned on each language’s corpus, the next step is to concatenate all the splits together to create a multilingual corpus of all four languages. As we did with the previous German and French analysis, we can use our concatenate_splits function to do this step for us on the list of coropora we generate in the previous step:

corpora_encoded = concatenate_splits(corpora)

Now that we have our multilingual corpus, we run the familiar steps with the trainer

training_args.logging_steps = len(corpora_encoded["train"]) // batch_size
trainer.train_dataset = corpora_encoded["train"]
trainer.eval_dataset = corpora_encoded["validation"]
trainer.train();

The final step is generate the predictions from the trainer on each language’s test set. This will give us an insight into how well multilingual learning is really working. We’ll collect the $F_1$-scores in our f1_scores dictionary and then create a DataFrame that summarizes the main results from our multilingual experiments:

for idx, lang in enumerate(langs):
    f1_scores["all"][lang] = (trainer
                              .predict(corpora[idx]["test"])
                              .metrics["eval_f1"])

scores_data = {"de": f1_scores["de"],
               "each": {lang: f1_scores[lang][lang] for lang in langs},
               "all": f1_scores["all"]}
f1_scores_df = pd.DataFrame.from_dict(scores_data, orient="index").round(4)
f1_scores_df.index.name = "Fine-tune on"

From these results we can draw a few general conclusions:

• Multilingual learning can provide significant gains in performance, especially if the low-resource languages for cross-lingual transfer belong to similar language families. In our experiments we can see that German, French, and Italian achieve similar performance in the all category suggesting that these languages are more similar to each other than English.

• As a general strategy, it is a good idea to focus attention on cross-lingual transfer within language families, especially when dealing with different scripts like Japanese.

Building a Pipeline for Inference

Although the Trainer object is useful for training and evaluation, in production we would like to be able to pass raw text as input and receive the model’s predictions as output. Fortunately, there is a way to do that using the Transformers pipeline abstraction!

For named entity recognition, we can use the TokenClassificationPipeline so we just need to load the model and tokenizer and wrap them as follows:

from transformers import TokenClassificationPipeline

pipeline = TokenClassificationPipeline(trainer.model.to("cpu"), xlmr_tokenizer,
                                       grouped_entities=True)

Note that we set the model’s device to cpu since it is generally faster to run inference on CPUs. Once the pipeline is loaded, we can then pass raw text to retrieve the predictions in a structured format:

pipeline(text_de)

[{'entity_group': 'PER',
  'score': 0.9977577924728394,
  'word': 'Jeff Dean',
  'start': 0,
  'end': 9},
 {'entity_group': 'ORG',
  'score': 0.984943151473999,
  'word': 'Google',
  'start': 35,
  'end': 41},
 {'entity_group': 'LOC',
  'score': 0.9276596009731293,
  'word': 'Kalifornien',
  'start': 45,
  'end': 56}]

By inspecting the output we see each word is given both a predicted entity, confidence score, and indices to locate it in the span of text.

Conclusion

In this chapter we saw how one can tackle NLP task on a multilingual corpus using a single Transformer pretrained on 100 languages: XLM-R. Although we were able to show that cross-lingual transfer from German to French is competitive when only a small number of labeled examples are available for fine-tuning, this good performance generally does not occur if the target language is significantly different from German or was not one of the 100 languages used during pretraining. For such cases, poor performance can be understood from a lack of model capacity in both the vocabulary and space of cross-lingual representations. Recent proposals like MAD-X⁸ are designed precisely for these low-resource scenarios, and since MAD-X is built on top of Transformers you can easily adapt the code in this chapter to work with it!

In this chapter we saw that cross-lingual transfer helps improve the performance on tasks in a langauge where label are scarce. In the next chapter we will see how we can deal with few labels in cases where we can’t use cross-lingual transfer, for example if there is no language with many labels.