Train GPT-2 in your own language. A step-by-step guide to train your own... | by Arshabhi Kayal | Towards Data Science

We all know modern day Natural Language Processing (NLP) has progressed by leaps and bounds in the past couple of years following the development of attention networks and transformers. It paved the way for a plethora of new algorithms achieving State-Of-The-Art (SOTA) for the different tasks of NLP.
OpenAI has been one of the leaders in providing their own language model (now released GPT-3) which is trained on a huge corpus of internet data. Since, GPT-3 is a recent phenomenon and in English at the moment, and is only accessible through API given by OpenAI, we shift our focus on the earlier version of it i.e. GPT-2. To know about the internal nuts and bolts of GPT-2, I’d suggest you to go through this link. For more depths into Attention and Transformers, here are some excellent links:
The illustrated Transformer by Jay Alammar
The Annotated Transformer by Harvard NLP
GPT-2 was also released for English, which makes it difficult for someone trying to generate text in a different language.
So why not train your own GPT-2 model on your favourite language for text generation? That is exactly what we are going to do. So, without further ado, let us jump in.
For the demo, I have considered a non-Latin alphabet script (Bengali here), because why not!! I have used Huggingface’s implementation for the model.
Gathering good quality data is one of the most important stages as all Data Scientists would agree. So, we are going to assume that you already have a folder containing .txt files having all the data cleaned and stored. For ease, you can use the Wikipedia article data, which is available and can be downloaded with the following code.
python wikipedia_download.py --lang bn
This will create a folder containing all Wikipedia files looking like:
Note: due to resource constraint, and since it is for demo purpose, I have trained the model in a small subset of books by Satyajit Ray, especially his detective Feluda series.
Now, the second step will be to tokenize the data. For that, we use the following class.
We use BPE (Byte Pair Encoding), which is a sub word encoding, this generally takes care of not treating different forms of word as different. (e.g. greatest will be treated as two tokens: ‘great’ and ‘est’ which is advantageous since it retains the similarity between great and greatest, while ‘greatest’ has another token ‘est’ added which makes it different). Also, it is not as low level as character-level encoding, which doesn’t retain any value of a particular word.
Another small but subtle point is NFKC (Normalization Form Compatibility Composition) in line 13 of code. It is one of the standard Unicode compatibility form. It would not matter much if the language is English, but since we are using Bengali, which contains a different form of character, we are using this specific one. More on it can be found in this link
So what we do here is tokenize our data and save it in a folder. Two files will be created (merges.txt and vocab.json) in a specified directory. To run the file, use the following code:
from tokenise import BPE_tokenfrom pathlib import Pathimport os# the folder 'text' contains all the filespaths = [str(x) for x in Path("./text/").glob("**/*.txt")]tokenizer = BPE_token()# train the tokenizer modeltokenizer.bpe_train(paths)# saving the tokenized data in our specified folder save_path = 'tokenized_data'tokenizer.save_tokenizer(save_path)
Before the real magic begins, we need to make sure the artilleries are ready. Let us start with some initializations.
import tensorflow as tffrom transformers import GPT2Config, TFGPT2LMHeadModel, GPT2Tokenizer# loading tokenizer from the saved model pathtokenizer = GPT2Tokenizer.from_pretrained(save_path)tokenizer.add_special_tokens({  "eos_token": "
",  "bos_token": "
",  "unk_token": "<unk>",  "pad_token": "<pad>",  "mask_token": "<mask>"})# creating the configurations from which the model can be madeconfig = GPT2Config(  vocab_size=tokenizer.vocab_size,  bos_token_id=tokenizer.bos_token_id,  eos_token_id=tokenizer.eos_token_id)# creating the modelmodel = TFGPT2LMHeadModel(config)
We also create a single string from all our documents and tokenize it.
single_string = ''for filename in paths:  with open(filename, "r", encoding='utf-8') as f:   x = f.read()  single_string += x + tokenizer.eos_tokenstring_tokenized = tokenizer.encode(single_string)
After we have encoded the whole string, we now move on to make a TensorFlow dataset, slicing the data into equal intervals, so that our model can learn. Here we use a block size of 100 (length of token in each example) and a batch size of 16. This is kept low else we can run it with ease on a RTX 2060 GPU.
examples = []block_size = 100BATCH_SIZE = 12BUFFER_SIZE = 1000for i in range(0, len(string_tokenized) - block_size + 1, block_size):  examples.append(string_tokenized[i:i + block_size])inputs, labels = [], []for ex in examples:  inputs.append(ex[:-1])  labels.append(ex[1:])dataset = tf.data.Dataset.from_tensor_slices((inputs, labels))dataset = dataset.shuffle(BUFFER_SIZE).batch(BATCH_SIZE, drop_remainder=True)
Now comes the part we’ve been waiting for, making the model and training. So we define our optimizer, loss functions and the metrics, and start training.
# defining our optimizeroptimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)# definining our loss functionloss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)# defining our metric which we want to observemetric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')# compiling the modelmodel.compile(optimizer=optimizer, loss=[loss, *[None] * model.config.n_layer], metrics=[metric])
Now, let’s train the model
num_epoch = 10history = model.fit(dataset, epochs=num_epoch)
To predict, we just need to simply encode the input text and pass it to the model
text = "লালমোহনবাবু "# encoding the input textinput_ids = tokenizer.encode(text, return_tensors='tf')# getting out outputbeam_output = model.generate(  input_ids,  max_length = 50,  num_beams = 5,  temperature = 0.7,  no_repeat_ngram_size=2,  num_return_sequences=5)
Now, if you are a Bengali, then you can point it out that the output although the sentence is syntactically correct, it doesn’t look cohesive. True, but for this demo, I have kept this demo a minimal as possible.
Well, after long training time, what good will it do if we close our session and all our trained model is just lost and we again need to train it from scratch. So, let’s save the model and the tokenizer so that we can retrain from where we left off
from transformers import WEIGHTS_NAME, CONFIG_NAMEimport osoutput_dir = './model_bn_custom/'# creating directory if it is not presentif not os.path.exists(output_dir):  os.mkdir(output_dir)model_to_save = model.module if hasattr(model, 'module') else modeloutput_model_file = os.path.join(output_dir, WEIGHTS_NAME)output_config_file = os.path.join(output_dir, CONFIG_NAME)# save model and model configsmodel.save_pretrained(output_dir)model_to_save.config.to_json_file(output_config_file)# save tokenizertokenizer.save_pretrained(output_dir)
We have already done all the hard work, so to load the saved model and tokenizer, we only need to execute two lines of code and we’re all set.
tokenizer = GPT2Tokenizer.from_pretrained(output_dir)model = TFGPT2LMHeadModel.from_pretrained(output_dir)
Voila! Now you can train your own model in your own language. And create content which can race with some of the best literary works in any language.
This blog gives a framework of how can one train GPT-2 model in any language. This is not at par with some of the pre-trained model available, but to reach that state, we need a lot of training data and computational power.