How to adapt a multilingual T5 model for a single language | by David Dale | Towards Data Science

T5 is an encoder-decoder transformer from Google that once was SOTA on several NLU and NLG problems and is still very useful as a base for seq2seq tasks such as text summarization. The first T5 model was for English only, and then the massively multilingual version followed. This model covers 101 languages and is massive indeed.
This post shows how to extract a single-language model from the multilingual one by pruning its redundant embeddings. This reduces the number of parameters more than twice without significant loss in quality. Our result is for Russian, but you can try it with any other of the 101 languages that mT5 features.
The idea is similar to one in the paper Load What You Need: Smaller Versions of Multilingual BERT. We use the original tokenizer to process a Russian corpus, count the frequencies of different tokens, and preserve only the tokens that were used frequently enough, pruning all the others.
We also preserve a small number of English tokens in the model to make it bilingual. We need this to enable the model to transfer knowledge from English to Russian downstream tasks, and also because English words and phrases often occur within modern Russian texts.
We start by loading the existing multilingual model.
import torchfrom transformers import T5ForConditionalGeneration, T5Tokenizertokenizer = T5Tokenizer.from_pretrained("google/mt5-base")model = T5ForConditionalGeneration.from_pretrained('google/mt5-base')
The model consists mostly of embeddings: 33% of its parameters are input embeddings (shared between its encoder and decoder) and 33% are output embeddings.
def msize(m):    return sum(p.numel() for p in m.parameters())print(msize(model.shared) / msize(model))   # 0.3298print(msize(model.lm_head) / msize(model))  # 0.3298
To estimate the frequency of different tokens, we take a Russian and an English sentence corpora from the Leipzig corpora collection. We use these two languages because we want our model to be bilingual in the end.
import pandas as pdimport csvfrom collections import Counterfrom tqdm.auto import tqdm, trangedf_ru = pd.read_csv('rus-ru_web-public_2019_1M-sentences.txt', sep='\t', header=None, quoting=csv.QUOTE_NONE)df_ru.columns = ['idx', 'text']cnt_ru = Counter()for text in tqdm(df_ru.text):    cnt_ru.update(tokenizer.encode(text))print(len(cnt_ru), len(cnt_ru)/tokenizer.vocab_size)  # 58438 0.2336
After counting the tokens in the Russian corpus we discover that only 23% of the model vocabulary was used. Moreover, the top 20K tokens constitute more than 99% of the Russian corpus. For English, the statistics are similar.
for top in 10_000, 20_000, 30_000:    print(top, sum(v for k, v in cnt_ru.most_common(top)) / sum(cnt_ru.values()))# 10000 0.9645# 20000 0.9940# 30000 0.9982
We decide to use the following composition of vocabulary:
1K of top tokens of the original tokenizer (just in case)
Top 10K of the English vocabulary
Top 20K of the Russian vocabulary
The 100 special tokens that T5 uses
This gives us a vocabulary of 30K tokens, 12% of the 250K tokens in the multilingual version.
new_tokens = set(range(1000))for i, (k, v) in enumerate(cnt_en.most_common(10_000)):    if k not in new_tokens:        new_tokens.add(k)for i, (k, v) in enumerate(cnt_ru.most_common(25_000)):    if len(new_tokens) == 29_900:        print(i, 'Russan tokens are included')        break    if k not in new_tokens:        new_tokens.add(k)for t in range(tokenizer.vocab_size - 100, tokenizer.vocab_size):    new_tokens.add(t)print(len(new_tokens))kept_ids = sorted(new_tokens)
Updating the neural networks is easy: just replace the parameters of its input and output embeddings. This reduces the model size by 58% (from 2.2GB to 0.9GB).
new_size = len(kept_ids)new_emb = torch.nn.Embedding(new_size, model.shared.embedding_dim)new_head = torch.nn.Linear(in_features=model.lm_head.in_features, out_features=new_size, bias=False)for new_id, old_id in enumerate(kept_ids):    new_emb.weight.data[new_id] = model.shared.weight.data[old_id]    new_head.weight.data[new_id] = model.lm_head.weight.data[old_id]model.shared.weight = new_emb.weightmodel.lm_head.weight = new_head.weightmodel.config.__dict__['vocab_size'] = new_sizemodel.config.__dict__['_name_or_path'] = 'cointegrated/rut5-base'
Updating the tokenizer is surprisingly more tricky. T5 uses Sentencepiece tokenizer, which is implemented in C and is opaque to Python. Fortunately, we can download its model and deploy it into Python using its Protobuf representation.
! wget https://raw.githubusercontent.com/google/sentencepiece/master/src/sentencepiece_model.proto! protoc --python_out=. sentencepiece_model.protoimport sentencepiece_model_pb2 as spmpsmp = tokenizer.sp_model.serialized_model_proto()m = spmp.ModelProto()m.ParseFromString(smp)print('the loaded model has pieces:', len(m.pieces))new_pieces = [m.pieces[idx] for idx in kept_ids]print('the new pieces:', len(new_pieces))# replace the content of the first 30K piecesfor i, p in enumerate(new_pieces):    m.pieces[i].piece = p.piece    m.pieces[i].score = p.score    m.pieces[i].type = p.type# drop the remaining piecesn = len(new_pieces)for i in trange(len(m.pieces) - n):    m.pieces.pop(len(m.pieces) - 1)print(len(m.pieces))with open('new_sp.model', 'wb') as f:    f.write(m.SerializeToString())new_tokenizer = T5Tokenizer('new_sp.model', extra_ids=0)
Now we can save the new model and the new tokenizer.
new_tokenizer.save_pretrained('rut5-base')model.save_pretrained('rut5-base')
All the code for creating models up to this stage is available on Github. The Russian T5 model is available in the Huggingface repository.
Frankly, this model is pretty useless by itself, because mT5 was trained only on the unsupervised task of predicting missing words. However, this model can be fine-tuned for many other tasks: text summarization, translation, dialogue response generation, paraphrasing, etc. In the next post, we will show how to perform such fine-tuning. Subscribe to stay tuned!
The post was written by David Dale (https://daviddale.ru/en), a research scientist in NLP and developer of chatbots.