Hands-on guide to constructing language model for MLM tasks from scratch using Python and Transformers library

In recent times, large language models(LLMs) have taken all the eye from the machine learning community. Before LLMs got here in, we had a vital research phase on various language modeling techniques, including masked language modeling, causal language modeling, and sequence-to-sequence language modeling.

From the above list, masked language models corresponding to BERT became more usable in downstream NLP tasks corresponding to classification and clustering. Due to libraries corresponding to Hugging Face Transformers, adapting these models for downstream tasks became more accessible and manageable. Also due to the open-source community, we have now loads of language models to select from covering widely used languages and domains.

When adapting existing language models to your specific use cases, sometimes we will use existing models without further tuning (so-called fine-tuning). For instance, for those who want an English sentiment/intent detection model, you may go into HuggingFace.co and find an appropriate model on your use case.

Nevertheless, you may only expect this for among the tasks encountered in the actual world. That is where we want an extra technique called fine-tuning. First, you will need to select a base model that can be fine-tuned. Here, you will need to watch out concerning the chosen model and your goal language’s lexical similarity.

Nevertheless, for those who cannot find an appropriate model retrained on the specified language, consider constructing one from scratch. On this tutorial, we’ll implement the BERT model for the masked language model.

Though describing BERT architecture is out of the scope of this tutorial, for the sake of clarity, let’s undergo it very narrowly. BERT, or Bidirectional Encoder Representations from Transformers, belongs to the encoder-only transformer family. It was introduced within the 12 months 2018 by researchers at Google.

Paper abstract:

We introduce a brand new language representation model called BERT, which stands for Bidirectional Encoder Representations from Transformers. Unlike recent language representation models, BERT is designed to pre-train deep bidirectional representations from unlabeled text by jointly conditioning on each left and right context in all layers. Consequently, the pre-trained BERT model may be fine-tuned with only one additional output layer to create state-of-the-art models for a wide selection of tasks, corresponding to query answering and language inference, without substantial task-specific architecture modifications.
BERT is conceptually easy and empirically powerful. It obtains latest state-of-the-art results on eleven natural language processing tasks, including pushing the GLUE rating to 80.5% (7.7% point absolute improvement), MultiNLI accuracy to 86.7% (4.6% absolute improvement), SQuAD v1.1 query answering Test F1 to 93.2 (1.5 point absolute improvement) and SQuAD v2.0 Test F1 to 83.1 (5.1 point absolute improvement).
Paper: https://arxiv.org/abs/1810.04805

Within the above, we will see an interesting keyword, Bidirectional. Bidirectional nature gives human-like power to BERT. Assume you’ve got to fill in a blank just like the below one,

“War may sometimes be a crucial evil. But irrespective of how crucial, it’s at all times an _____, never an excellent.”

To guess the word comes right into a blank position, you have to be aware of a number of things: words before empty, words after blank, and the general context of the sentence. Adapting this human nature, BERT works the identical way. During training, we hide some words and ask BERT to try predicting those. When training is finished, BERT can predict masked tokens based on their before and after words. To do that, the model should allocate different attention to words presented within the input sequence, which can significantly impact predicting masked tokens.

As you may see here, the model sees an appropriate word for the hidden position as evil and the primary sentence’s evil as crucial to make this prediction. This can be a noticeable point and implies that the model understands the context of the input sequence. This context awareness allows BERT to generate meaningful sentence embeddings for given tasks. Further, these embeddings may be utilized in downstream tasks corresponding to clustering and classification. Enough about BERT; let’s construct one from scratch.

We generally have BERT(base) and BERT(large). Each have 64 dimensions per head. The large variant comprises 24 encoder layers, while the base variant only has 12. Nevertheless, we will not be limited to those configurations. Surprisingly, we have now complete control over defining the model using Hugging Face Transformers library. All we have now to do is define desired model configurations using the BertConfig class.

I selected 6 heads and 384 total model dimensions to comply with the unique implementation. In this manner, each head has 64 dimensions just like the unique implementation. Let’s initialize our BERT model.

from transformers import BertConfig, BertForMaskedLM
config = BertConfig(
hidden_size = 384,
vocab_size= tokenizer.vocab_size,
num_hidden_layers = 6,
num_attention_heads = 6,
intermediate_size = 1024,
max_position_embeddings = 256
)
model = BertForMaskedLM(config=config)
print(model.num_parameters()) #10457864

Here, I won’t describe how tokenization works under the hood. As an alternative, let’s train one from scratch using Hugging Face tokenizers library. Please note that the tokenizer utilized in the unique BERT implementation is WordPiece tokenizer, one more subword-based tokenization method. You may learn more about this tokenization using the neat HuggingFace resource below.

The dataset used here is the Sinhala-400M dataset (under apache-2.0). You may follow the identical with any dataset you’ve got.

As you could notice, some Sinhalese words have been typed using English as well. Let’s train a tokenizer for these corpora.

Let’s import the crucial modules first. The advantage of training tokenizers using Hugging Face Tokenizers library is that we will use existing tokenizers and replace only vocabulary (and merge where applicable) per our training corpus. This implies tokenization steps corresponding to pre-tokenization and post-tokenization can be preserved. For this, we will use a way, train_new_from_iterator BertTokenizer class.

from tokenizers.implementations import ByteLevelBPETokenizer
from tokenizers.processors import BertProcessing
from transformers import AutoTokenizer
from datasets import Dataset
import pandas as pd
#load base tokenizer to coach on dataset
tokenizer_base = AutoTokenizer.from_pretrained("bert-base-cased")
# convert pandas dataset to HF dataset
dataset = Dataset.from_pandas(df.rename(columns={"comment":'text'}))
# define iterator
training_corpus = (
dataset[i : i + 1000]["text"]
for i in range(0, len(dataset), 1000)
)
#train the brand new tokenizer for dataset
tokenizer = tokenizer_base.train_new_from_iterator(training_corpus, 5000)
#test trained tokenizer for sample text
text = dataset['text'][123]
print(text)

# let's check tokenization process
input_ids = tokenizer(text).input_ids
subword_view = [tokenizer.convert_ids_to_tokens(id) for id in input_ids]
np.array(subword_view)

You may see words like ‘cricketer’ decomposed into cricket and ##er, indicating that the tokenizer has been adequately trained. Nevertheless, check out different vocab sizes; mine is 5000, which is comparatively small but suitable for this toy example.

Finally, we will save the trained tokenizer into our directory.

tokenizer.save_pretrained("tokenizer/sinhala-wordpiece-yt-comments")

Let’s define a collator for MLM tasks. Here, we’ll mask 15% of tokens. Anyway, we will set different masking probabilities.

from transformers import DataCollatorForLanguageModeling
data_collator = DataCollatorForLanguageModeling(
tokenizer=tokenizer, multilevel marketing=True, mlm_probability=0.15
)

Let’s tokenize the dataset using a previously created tokenizer. I’m replacing the unique LineByLineTextDataset using my custom class utilizing Hugging Face speed up.

import torch
from torch.utils.data import Dataset
from speed up import Accelerator, DistributedType
class LineByLineTextDataset(Dataset):
def __init__(self, tokenizer, raw_datasets, max_length: int):
self.padding = "max_length"
self.text_column_name = 'text'
self.max_length = max_length
self.accelerator = Accelerator(gradient_accumulation_steps=1)
self.tokenizer = tokenizer
with self.accelerator.main_process_first():
self.tokenized_datasets = raw_datasets.map(
self.tokenize_function,
batched=True,
num_proc=4,
remove_columns=[self.text_column_name],
desc="Running tokenizer on dataset line_by_line",
)
self.tokenized_datasets.set_format('torch',columns=['input_ids'],dtype=torch.long)
def tokenize_function(self,examples):
examples[self.text_column_name] = [
line for line in examples[self.text_column_name] if len(line[0]) > 0 and never line[0].isspace()
]
return self.tokenizer(
examples[self.text_column_name],
padding=self.padding,
truncation=True,
max_length=self.max_length,
return_special_tokens_mask=True,
)
def __len__(self):
return len(self.tokenized_datasets)
def __getitem__(self, i):
return self.tokenized_datasets[i]

Let’s tokenize the dataset.

tokenized_dataset_train = LineByLineTextDataset(
tokenizer= tokenizer,
raw_datasets = dataset,
max_length=256,
)

Alright, let’s code our training loop.

from transformers import Trainer, TrainingArguments
training_args = TrainingArguments(
output_dir="./model",
overwrite_output_dir=True,
push_to_hub=True,
hub_model_id="Ransaka/sinhala-bert-yt-comments",
num_train_epochs=2,
per_device_train_batch_size=32,
save_steps=5_000,
logging_steps = 1000,
save_total_limit=2,
use_mps_device = True, # disable this for those who're running non-mac env
hub_private_repo = False, # please set true if you would like to save model privetly
save_safetensors= True,
learning_rate = 1e-4,
report_to='wandb'
)
trainer = Trainer(
model=model,
args=training_args,
data_collator=data_collator,
train_dataset=tokenized_dataset_train
)trainer.train()

We will invoke the trainer using its train() method.

trainer.train()

After sufficient training, our model may be used for downstream tasks corresponding to zero-shot classification and clustering. You could find the instance using this Hugging Face space for more details.

With limited resources, pre-trained models may only recognize specific linguistic patterns, but they will still be helpful for particular use cases. It is very really useful to fine-tune when possible.

In this text, all images, unless otherwise noted, are by the creator.

1. An Explorable BERT — https://huggingface.co/spaces/exbert-project/exbert
2. BERT Paper — https://arxiv.org/abs/1810.04805
3. Dataset — https://huggingface.co/datasets/Ransaka/Sinhala-400M