title pretraining bert with hugging face transformers sreyan ghosh date created 20220701 last modified 20220827 description pretraining bert using hugging face transformers on nsp and mlm accelerator gpu introduction bert bidirectional encoder representations from transformers in the field of computer vision researchers have repeatedly shown the value of transfer learning pretraining a neural network model on a known taskdataset for instance imagenet classification and then performing finetuning using the trained neural network as the basis of a new specificpurpose model in recent years researchers have shown that a similar technique can be useful in many natural language tasks bert makes use of transformer an attention mechanism that learns contextual relations between words or subwords in a text in its vanilla form transformer includes two separate mechanisms an encoder that reads the text input and a decoder that produces a prediction for the task since berts goal is to generate a language model only the encoder mechanism is necessary the detailed workings of transformer are described in a paper by google as opposed to directional models which read the text input sequentially lefttoright or righttoleft the transformer encoder reads the entire sequence of words at once therefore it is considered bidirectional though it would be more accurate to say that its nondirectional this characteristic allows the model to learn the context of a word based on all of its surroundings left and right of the word when training language models a challenge is defining a prediction goal many models predict the next word in a sequence e g the child came home from a directional approach which inherently limits context learning to overcome this challenge bert uses two training strategies masked language modeling mlm before feeding word sequences into bert 15 of the words in each sequence are replaced with a mask token the model then attempts to predict the original value of the masked words based on the context provided by the other nonmasked words in the sequence next sentence prediction nsp in the bert training process the model receives pairs of sentences as input and learns to predict if the second sentence in the pair is the subsequent sentence in the original document during training 50 of the inputs are a pair in which the second sentence is the subsequent sentence in the original document while in the other 50 a random sentence from the corpus is chosen as the second sentence the assumption is that the random sentence will represent a disconnect from the first sentence though google provides a pretrained bert checkpoint for english you may often need to either pretrain the model from scratch for a different language or do a continuedpretraining to fit the model to a new domain in this notebook we pretrain bert from scratch optimizing both mlm and nsp objectves using transformers on the wikitext english dataset loaded from datasets setup installing the requirements shell pip install githttps github comhuggingfacetransformers git pip install datasets pip install huggingfacehub pip install nltk importing the necessary libraries set random seed define certain variables load the wikitext dataset we now download the wikitext language modeling dataset it is a collection of over 100 million tokens extracted from the set of verified good and featured articles on wikipedia we load the dataset from datasetshttps github comhuggingfacedatasets for the purpose of demonstration in this notebook we work with only the train split of the dataset this can be easily done with the loaddataset function the dataset just has one column which is the raw text and this is all we need for pretraining bert training a new tokenizer first we train our own tokenizer from scratch on our corpus so that can we can use it to train our language model from scratch but why would you need to train a tokenizer that s because transformer models very often use subword tokenization algorithms and they need to be trained to identify the parts of words that are often present in the corpus you are using the transformers tokenizer as the name indicates will tokenize the inputs including converting the tokens to their corresponding ids in the pretrained vocabulary and put it in a format the model expects as well as generate the other inputs that model requires first we make a list of all the raw documents from the wikitext corpus next we make a batchiterator function that will aid us to train our tokenizer in this notebook we train a tokenizer with the exact same algorithms and parameters as an existing one for instance we train a new version of the bertcased tokenzier on wikitext2 using the same tokenization algorithm first we need to load the tokenizer we want to use as a model now we train our tokenizer using the entire train split of the wikitext2 dataset so now we our done training our new tokenizer next we move on to the data preprocessing steps data preprocessing for the sake of demonstrating the workflow in this notebook we only take small subsets of the entire wikitext train and test splits before we can feed those texts to our model we need to preprocess them and get them ready for the task as mentioned earlier the bert pretraining task includes two tasks in total the nsp task and the mlm task transformers have an easy to implement collator called the datacollatorforlanguagemodeling however we need to get the data ready for nsp manually next we write a simple function called the preparetrainfeatures that helps us in the preprocessing and is compatible with datasets to summarize our preprocessing function should get the dataset ready for the nsp task by creating pairs of sentences a b where b either actually follows a or b is randomly sampled from somewhere else in the corpus it should also generate a corresponding label for each pair which is 1 if b actually follows a and 0 if not tokenize the text dataset into it s corresponding token ids that will be used for embedding lookup in bert create additional inputs for the model like tokentypeids attentionmask etc we define the maximum number of tokens after tokenization that each training sample will have function to prepare features for nsp task arguments examples a dictionary with 1 key text text list of raw documents str returns examples a dictionary with 4 keys inputids list of tokenized concatnated and batched sentences from the individual raw documents int tokentypeids list of integers 0 or 1 corresponding to 0 for senetence no 1 and padding 1 for sentence no 2 attentionmask list of integers 0 or 1 corresponding to 1 for nonpadded tokens 0 for padded nextsentencelabel list of integers 0 or 1 corresponding to 1 if the second sentence actually follows the first 0 if the senetence is sampled from somewhere else in the corpus remove unwanted samples from the training set split the documents from the dataset into it s individual sentences convert the tokens into ids using the trained tokenizer define the outputs we usually want to fill up the entire sequence since we are padding to blocksize anyways so short sequences are generally wasted computation however we sometimes i e shortseqprob 0 1 10 of the time want to use shorter sequences to minimize the mismatch between pretraining and finetuning the targetseqlength is just a rough target however whereas blocksize is a hard limit aend is how many segments from currentchunk go into the a first sentence this should rarely go for more than one iteration for large corpora however just to be careful we try to make sure that the random document is not the same as the document we re processing we didn t actually use these segments so we put them back so they don t go to waste add token type ids 0 for sentence a 1 for sentence b we delete all the unnecessary columns from our dataset for mlm we are going to use the same preprocessing as before for our dataset with one additional step we randomly mask some tokens by replacing them by mask and the labels will be adjusted to only include the masked tokens we don t have to predict the nonmasked tokens if you use a tokenizer you trained yourself make sure the mask token is among the special tokens you passed during training to get the data ready for mlm we simply use the collator called the datacollatorforlanguagemodeling provided by the transformers library on our dataset that is already ready for the nsp task the collator expects certain parameters we use the default ones from the original bert paper in this notebook the returntensors tf ensures that we get tf tensor objects back next we define our training set with which we train our model again datasets provides us with the totfdataset method which will help us integrate our dataset with the collator defined above the method expects certain parameters columns the columns which will serve as our independant variables labelcols the columns which will serve as our labels or dependant variables batchsize our batch size for training shuffle whether we want to shuffle our training dataset collatefn our collator function defining the model to define our model first we need to define a config which will help us define certain parameters of our model architecture this includes parameters like number of transformer layers number of attention heads hidden dimension etc for this notebook we try to define the exact config defined in the original bert paper we can easily achieve this using the bertconfig class from the transformers library the frompretrained method expects the name of a model here we define the simplest model with which we also trained our model i e bertbasecased for defining our model we use the tfbertforpretraining class from the transformers library this class internally handles everything starting from defining our model to unpacking our inputs and calculating the loss so we need not do anything ourselves except defining the model with the correct config we want now we define our optimizer and compile the model the loss calculation is handled internally and so we need not worry about that finally all steps are done and now we can start training our model our model has now been trained we suggest to please train the model on the complete dataset for atleast 50 epochs for decent performance the pretrained model now acts as a language model and is meant to be finetuned on a downstream task thus it can now be finetuned on any downstream task like question answering text classification etc now you can push this model to model hub and also share it with with all your friends family favorite pets they can all load it with the identifier yourusernamethenameyoupicked so for instance python model pushtohubpretrainedbert organizationkerasio tokenizer pushtohubpretrainedbert organizationkerasio and after you push your model this is how you can load it in the future python from transformers import tfbertforpretraining model tfbertforpretraining frompretrainedyourusernamemyawesomemodel or since it s a pretrained model and you would generally use it for finetuning on a downstream task you can also load it for some other task like python from transformers import tfbertforsequenceclassification model tfbertforsequenceclassification frompretrainedyourusernamemyawesomemodel in this case the pretraining head will be dropped and the model will just be initialized with the transformer layers a new taskspecific head will be added with random weights introduction bert bidirectional encoder representations from transformers in the field of computer vision researchers have repeatedly shown the value of transfer learning pretraining a neural network model on a known task dataset for instance imagenet classification and then performing fine tuning using the trained neural network as the basis of a new specific purpose model in recent years researchers have shown that a similar technique can be useful in many natural language tasks bert makes use of transformer an attention mechanism that learns contextual relations between words or subwords in a text in its vanilla form transformer includes two separate mechanisms an encoder that reads the text input and a decoder that produces a prediction for the task since bert s goal is to generate a language model only the encoder mechanism is necessary the detailed workings of transformer are described in a paper by google as opposed to directional models which read the text input sequentially left to right or right to left the transformer encoder reads the entire sequence of words at once therefore it is considered bidirectional though it would be more accurate to say that it s non directional this characteristic allows the model to learn the context of a word based on all of its surroundings left and right of the word when training language models a challenge is defining a prediction goal many models predict the next word in a sequence e g the child came home from _ a directional approach which inherently limits context learning to overcome this challenge bert uses two training strategies masked language modeling mlm before feeding word sequences into bert 15 of the words in each sequence are replaced with a mask token the model then attempts to predict the original value of the masked words based on the context provided by the other non masked words in the sequence next sentence prediction nsp in the bert training process the model receives pairs of sentences as input and learns to predict if the second sentence in the pair is the subsequent sentence in the original document during training 50 of the inputs are a pair in which the second sentence is the subsequent sentence in the original document while in the other 50 a random sentence from the corpus is chosen as the second sentence the assumption is that the random sentence will represent a disconnect from the first sentence though google provides a pretrained bert checkpoint for english you may often need to either pretrain the model from scratch for a different language or do a continued pretraining to fit the model to a new domain in this notebook we pretrain bert from scratch optimizing both mlm and nsp objectves using transformers on the wikitext english dataset loaded from datasets setup installing the requirements shell pip install git https github com huggingface transformers git pip install datasets pip install huggingface hub pip install nltk importing the necessary libraries set random seed define certain variables batch size to train the tokenizer on total number of unique subwords the tokenizer can have maximum number of tokens in an input sample probability that the next sentence is the actual next sentence in nsp probability of generating shorter sequences to minimize the mismatch between pretraining and fine tuning maximum number of tokens in an input sample after padding probability with which tokens are masked in mlm batch size for pretraining the model on maximum number of epochs to train the model for learning rate for training the model name of pretrained model from model hub load the wikitext dataset we now download the wikitext language modeling dataset it is a collection of over 100 million tokens extracted from the set of verified good and featured articles on wikipedia we load the dataset from datasets https github com huggingface datasets for the purpose of demonstration in this notebook we work with only the train split of the dataset this can be easily done with the load_dataset function the dataset just has one column which is the raw text and this is all we need for pretraining bert training a new tokenizer first we train our own tokenizer from scratch on our corpus so that can we can use it to train our language model from scratch but why would you need to train a tokenizer that s because transformer models very often use subword tokenization algorithms and they need to be trained to identify the parts of words that are often present in the corpus you are using the transformers tokenizer as the name indicates will tokenize the inputs including converting the tokens to their corresponding ids in the pretrained vocabulary and put it in a format the model expects as well as generate the other inputs that model requires first we make a list of all the raw documents from the wikitext corpus next we make a batch_iterator function that will aid us to train our tokenizer in this notebook we train a tokenizer with the exact same algorithms and parameters as an existing one for instance we train a new version of the bert cased tokenzier on wikitext 2 using the same tokenization algorithm first we need to load the tokenizer we want to use as a model now we train our tokenizer using the entire train split of the wikitext 2 dataset so now we our done training our new tokenizer next we move on to the data pre processing steps data pre processing for the sake of demonstrating the workflow in this notebook we only take small subsets of the entire wikitext train and test splits before we can feed those texts to our model we need to pre process them and get them ready for the task as mentioned earlier the bert pretraining task includes two tasks in total the nsp task and the mlm task transformers have an easy to implement collator called the datacollatorforlanguagemodeling however we need to get the data ready for nsp manually next we write a simple function called the prepare_train_features that helps us in the pre processing and is compatible with datasets to summarize our pre processing function should get the dataset ready for the nsp task by creating pairs of sentences a b where b either actually follows a or b is randomly sampled from somewhere else in the corpus it should also generate a corresponding label for each pair which is 1 if b actually follows a and 0 if not tokenize the text dataset into it s corresponding token ids that will be used for embedding look up in bert create additional inputs for the model like token_type_ids attention_mask etc we define the maximum number of tokens after tokenization that each training sample will have function to prepare features for nsp task arguments examples a dictionary with 1 key text text list of raw documents str returns examples a dictionary with 4 keys input_ids list of tokenized concatnated and batched sentences from the individual raw documents int token_type_ids list of integers 0 or 1 corresponding to 0 for senetence no 1 and padding 1 for sentence no 2 attention_mask list of integers 0 or 1 corresponding to 1 for non padded tokens 0 for padded next_sentence_label list of integers 0 or 1 corresponding to 1 if the second sentence actually follows the first 0 if the senetence is sampled from somewhere else in the corpus remove un wanted samples from the training set split the documents from the dataset into it s individual sentences convert the tokens into ids using the trained tokenizer define the outputs a buffer stored current working segments we usually want to fill up the entire sequence since we are padding to block_size anyways so short sequences are generally wasted computation however we sometimes i e short_seq_prob 0 1 10 of the time want to use shorter sequences to minimize the mismatch between pretraining and fine tuning the target_seq_length is just a rough target however whereas block_size is a hard limit a_end is how many segments from current_chunk go into the a first sentence this should rarely go for more than one iteration for large corpora however just to be careful we try to make sure that the random document is not the same as the document we re processing we didn t actually use these segments so we put them back so they don t go to waste add token type ids 0 for sentence a 1 for sentence b we delete all the un necessary columns from our dataset for mlm we are going to use the same preprocessing as before for our dataset with one additional step we randomly mask some tokens by replacing them by mask and the labels will be adjusted to only include the masked tokens we don t have to predict the non masked tokens if you use a tokenizer you trained yourself make sure the mask token is among the special tokens you passed during training to get the data ready for mlm we simply use the collator called the datacollatorforlanguagemodeling provided by the transformers library on our dataset that is already ready for the nsp task the collator expects certain parameters we use the default ones from the original bert paper in this notebook the return_tensors tf ensures that we get tf tensor objects back next we define our training set with which we train our model again datasets provides us with the to_tf_dataset method which will help us integrate our dataset with the collator defined above the method expects certain parameters columns the columns which will serve as our independant variables label_cols the columns which will serve as our labels or dependant variables batch_size our batch size for training shuffle whether we want to shuffle our training dataset collate_fn our collator function defining the model to define our model first we need to define a config which will help us define certain parameters of our model architecture this includes parameters like number of transformer layers number of attention heads hidden dimension etc for this notebook we try to define the exact config defined in the original bert paper we can easily achieve this using the bertconfig class from the transformers library the from_pretrained method expects the name of a model here we define the simplest model with which we also trained our model i e bert base cased for defining our model we use the tfbertforpretraining class from the transformers library this class internally handles everything starting from defining our model to unpacking our inputs and calculating the loss so we need not do anything ourselves except defining the model with the correct config we want now we define our optimizer and compile the model the loss calculation is handled internally and so we need not worry about that finally all steps are done and now we can start training our model our model has now been trained we suggest to please train the model on the complete dataset for atleast 50 epochs for decent performance the pretrained model now acts as a language model and is meant to be fine tuned on a downstream task thus it can now be fine tuned on any downstream task like question answering text classification etc now you can push this model to model hub and also share it with with all your friends family favorite pets they can all load it with the identifier your username the name you picked so for instance python model push_to_hub pretrained bert organization keras io tokenizer push_to_hub pretrained bert organization keras io and after you push your model this is how you can load it in the future python from transformers import tfbertforpretraining model tfbertforpretraining from_pretrained your username my awesome model or since it s a pretrained model and you would generally use it for fine tuning on a downstream task you can also load it for some other task like python from transformers import tfbertforsequenceclassification model tfbertforsequenceclassification from_pretrained your username my awesome model in this case the pretraining head will be dropped and the model will just be initialized with the transformer layers a new task specific head will be added with random weights
import nltk
import random
import logging

import keras

nltk.download("punkt")

keras.utils.set_random_seed(42)



TOKENIZER_BATCH_SIZE = 256  
TOKENIZER_VOCABULARY = 25000  

BLOCK_SIZE = 128  
NSP_PROB = 0.50  
SHORT_SEQ_PROB = 0.1  
MAX_LENGTH = 512  

MLM_PROB = 0.2  

TRAIN_BATCH_SIZE = 2  
MAX_EPOCHS = 1  
LEARNING_RATE = 1e-4  

MODEL_CHECKPOINT = "bert-base-cased"  





from datasets import load_dataset

dataset = load_dataset("wikitext", "wikitext-2-raw-v1")



print(dataset)





all_texts = [
    doc for doc in dataset["train"]["text"] if len(doc) > 0 and not doc.startswith(" =")
]




def batch_iterator():
    for i in range(0, len(all_texts), TOKENIZER_BATCH_SIZE):
        yield all_texts[i : i + TOKENIZER_BATCH_SIZE]




from transformers import AutoTokenizer

tokenizer = AutoTokenizer.from_pretrained(MODEL_CHECKPOINT)



tokenizer = tokenizer.train_new_from_iterator(
    batch_iterator(), vocab_size=TOKENIZER_VOCABULARY
)







dataset["train"] = dataset["train"].select([i for i in range(1000)])
dataset["validation"] = dataset["validation"].select([i for i in range(1000)])





max_num_tokens = BLOCK_SIZE - tokenizer.num_special_tokens_to_add(pair=True)


def prepare_train_features(examples):
    

    
    examples["document"] = [
        d.strip() for d in examples["text"] if len(d) > 0 and not d.startswith(" =")
    ]
    
    examples["sentences"] = [
        nltk.tokenize.sent_tokenize(document) for document in examples["document"]
    ]
    
    examples["tokenized_sentences"] = [
        [tokenizer.convert_tokens_to_ids(tokenizer.tokenize(sent)) for sent in doc]
        for doc in examples["sentences"]
    ]

    
    examples["input_ids"] = []
    examples["token_type_ids"] = []
    examples["attention_mask"] = []
    examples["next_sentence_label"] = []

    for doc_index, document in enumerate(examples["tokenized_sentences"]):
        current_chunk = []  
        current_length = 0
        i = 0

        
        
        
        
        
        
        
        target_seq_length = max_num_tokens

        if random.random() < SHORT_SEQ_PROB:
            target_seq_length = random.randint(2, max_num_tokens)

        while i < len(document):
            segment = document[i]
            current_chunk.append(segment)
            current_length += len(segment)
            if i == len(document) - 1 or current_length >= target_seq_length:
                if current_chunk:
                    
                    
                    a_end = 1
                    if len(current_chunk) >= 2:
                        a_end = random.randint(1, len(current_chunk) - 1)

                    tokens_a = []
                    for j in range(a_end):
                        tokens_a.extend(current_chunk[j])

                    tokens_b = []

                    if len(current_chunk) == 1 or random.random() < NSP_PROB:
                        is_random_next = True
                        target_b_length = target_seq_length - len(tokens_a)

                        
                        
                        
                        
                        for _ in range(10):
                            random_document_index = random.randint(
                                0, len(examples["tokenized_sentences"]) - 1
                            )
                            if random_document_index != doc_index:
                                break

                        random_document = examples["tokenized_sentences"][
                            random_document_index
                        ]
                        random_start = random.randint(0, len(random_document) - 1)
                        for j in range(random_start, len(random_document)):
                            tokens_b.extend(random_document[j])
                            if len(tokens_b) >= target_b_length:
                                break
                        
                        
                        num_unused_segments = len(current_chunk) - a_end
                        i -= num_unused_segments
                    else:
                        is_random_next = False
                        for j in range(a_end, len(current_chunk)):
                            tokens_b.extend(current_chunk[j])

                    input_ids = tokenizer.build_inputs_with_special_tokens(
                        tokens_a, tokens_b
                    )
                    
                    token_type_ids = tokenizer.create_token_type_ids_from_sequences(
                        tokens_a, tokens_b
                    )

                    padded = tokenizer.pad(
                        {"input_ids": input_ids, "token_type_ids": token_type_ids},
                        padding="max_length",
                        max_length=MAX_LENGTH,
                    )

                    examples["input_ids"].append(padded["input_ids"])
                    examples["token_type_ids"].append(padded["token_type_ids"])
                    examples["attention_mask"].append(padded["attention_mask"])
                    examples["next_sentence_label"].append(1 if is_random_next else 0)
                    current_chunk = []
                    current_length = 0
            i += 1

    
    del examples["document"]
    del examples["sentences"]
    del examples["text"]
    del examples["tokenized_sentences"]

    return examples


tokenized_dataset = dataset.map(
    prepare_train_features,
    batched=True,
    remove_columns=["text"],
    num_proc=1,
)



from transformers import DataCollatorForLanguageModeling

collater = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=True, mlm_probability=MLM_PROB, return_tensors="tf"
)



train = tokenized_dataset["train"].to_tf_dataset(
    columns=["input_ids", "token_type_ids", "attention_mask"],
    label_cols=["labels", "next_sentence_label"],
    batch_size=TRAIN_BATCH_SIZE,
    shuffle=True,
    collate_fn=collater,
)

validation = tokenized_dataset["validation"].to_tf_dataset(
    columns=["input_ids", "token_type_ids", "attention_mask"],
    label_cols=["labels", "next_sentence_label"],
    batch_size=TRAIN_BATCH_SIZE,
    shuffle=True,
    collate_fn=collater,
)





from transformers import BertConfig

config = BertConfig.from_pretrained(MODEL_CHECKPOINT)



from transformers import TFBertForPreTraining

model = TFBertForPreTraining(config)



from keras.optimizers import Adam

model.compile(optimizer=Adam(learning_rate=LEARNING_RATE))



model.fit(train, validation_data=validation, epochs=MAX_EPOCHS)