Extending a pretrained transformer vocabulary

[epwalsh]

This short post shows you how to extend the vocabulary of a pretrained transformer model with AllenNLP.

We assume you're using the PretrainedTransformerTokenizer with either the
PretrainedTransformerIndexer + PretrainedTransformerEmbedder or PretrainedTransformerMismatchedIndexer + PretrainedTransformerMismatchedEmbedder, in which case your configuration file would look something like this:

local model_name = "t5-base";

local transformer_spec = {
  "type": "pretrained_transformer",
  "model_name": model_name,
};

{
  "dataset_reader": {
    "tokenizer": transformer_spec,
    "token_indexers": { "tokens": transformer_spec },
    // additional dataset reader specs...
  },
  "model": {
    "text_field_embedder": {
      "type": "basic",
      "token_embedders": { "tokens": transformer_spec },
    },
    // additional model specs...
  },
  // additional training specs...
}

Then all you need to do is specify the tokens you want to add to the additional_special_tokens parameter in tokenizer_kwargs:

 local model_name = "t5-base";
+local extra_tokens = ["FOO", "BAR"];
 
 local transformer_spec = {
   "type": "pretrained_transformer",
   "model_name": model_name,
+  "tokenizer_kwargs": {
+    "additional_special_tokens": extra_tokens,
+  },
 };
 
 {
   "dataset_reader": {
     "tokenizer": transformer_spec,
     "token_indexers": { "tokens": transformer_spec },
     // additional dataset reader specs...
   },
   "model": {
     "text_field_embedder": {
       "type": "basic",
       "token_embedders": { "tokens": transformer_spec },
     },
     // additional model specs...
   },
   // additional training specs...
 }

And that's it! The PretrainedTransformerEmbedder or PretrainedTransformerMismatchedEmbedder will then automatically resize the embeddings by adding randomly initialized rows.