Language models have token the world by storm recently and, given the already explored analogies between protein primary sequence and text, there’s been a lot of interest in applying these models to protein sequences. Interest is not only coming from academia and the pharmaceutical industry, but also some very unlikely suspects such as ByteDance – yes the same ByteDance of TikTok fame. So if you also fancy trying your hand at building a protein language model then read on, it’s surprisingly easy. <\/p>\n\n\n\n
Training your own protein language model from scratch is made remarkably easy by the HuggingFace Transformers library, which allows you to specify a model architecture, tokenise your training data, and train a model in only a few lines of code. Under the hood, the Transformers library uses PyTorch (or optionally Tensorflow) models, allowing you to dig deeper into customising training or model architecture, or simply leave it to the highly abstracted Transformers library to handle it all for you.<\/p>\n\n\n\n
For this article, I’ll assume you already understand how language models work, and are now looking to implement one yourself, trained from scratch. <\/p>\n\n\n\n\n\n\n\n
We’ll be using both the Transformers and Datasets library from HuggingFace, as well as PyTorch. Installing transformers from conda can be problematic sometimes, so I’d recommend using pip, which can still be done under a conda virtual environment. <\/p>\n\n\n\n
pip install transformers\npip install datasets<\/pre>\n\n\n\nYou’ll also need to install PyTorch, you can check system-dependent guides here<\/a>.<\/p>\n\n\n\n
2. Prepare Your Dataset<\/h2>\n\n\n\n
With your dataset of protein sequences separated into train, validation, and test sets, we’ll use the Datasets library from HuggingFace for easy and efficient tokenisation.<\/p>\n\n\n\n
Protein language models typically tokenize on a character level (i.e., residue), so you won’t need to train a custom tokenizer like in most text-based language models. Download this example tokenizer, which includes a token for each of the 20 amino acids, as well as padding, and N and C-term tokens (represented by ‘1’ and ‘2’ respectively). <\/p>\n\n\n\n
protein_tokeniser<\/a>Download<\/a><\/div>\n\n\n\n<\/p>\n\n\n\n
tokeniser = AutoTokenizer.from_pretrained('path\/to\/tokeniser', use_fast=True)<\/pre>\n\n\n\nNext, we’ll load the train, val, and test sets. The Datasets library has some handy features such as caching any data preparation steps (e.g. tokenisation), meaning this step won’t be repeated between training runs. This saves a lot of compute, especially with the large datasets typically used in language models.<\/p>\n\n\n\n
data_set_paths = {\"train\": \"\/path\/to\/train.csv\", \"test\": \"\/path\/to\/test.csv\", \"val\": \"\/path\/to\/val.csv\"}\ndatasets = load_dataset('csv', data_files=data_sets, cache_dir='\/where\/to\/store\/the\/cache')<\/pre>\n\n\n\n3. Define the model architecture<\/strong><\/h2>\n\n\n\n
Choose a model architecture from the Transformers library (check their documentation for a full list). Here, we’ll use GPT2. Customize any model architecture hyperparameters as needed, since we won’t keep any pretrained model weights:<\/p>\n\n\n\n
transformer_config = GPT2Config(vocab_size=tokenizer.vocab_size, n_layer=12, n_embd=512, n_head=12, n_inner=2048)\n\nmodel = GPT2LMHeadModel(config=transformer_config)<\/pre>\n\n\n\nRemember to move the model to the GPU using
model.to()<\/code> if using CUDA.<\/p>\n\n\n\n4. Train your language model<\/strong><\/h2>\n\n\n\n
Now, it’s time to train our model! The Trainer class in HuggingFace makes this easy with training parameters from a TrainingArguments instance. Check the TrainingArguments documentation here<\/a> for customization options, or write your own training loop using PyTorch<\/a>.<\/p>\n\n\n\n
training_args = TrainingArguments(\n\t\t\t\t\t\t\t\toutput_dir=\".\/output\",\n\t\t\t\t\t\t\t\toverwrite_output_dir=True,\n\t\t\t\t\t\t\t\tnum_train_epochs=2,\n\t\t\t\t\t\t\t\tper_device_train_batch_size=16, \n\t\t\t\t\t\t\t\tper_device_eval_batch_size=16, \n\t\t\t\t\t\t\t\tsave_steps=10_000, \n\t\t\t\t\t\t\t\tsave_total_limit=2) \ntrainer = Trainer(\n\t\t\t\t model=model,\n\t\t\t\t args=training_args,\n\t\t\t\t train_dataset=tokenized_datasets[\"train\"],\n\t\t\t\t eval_dataset=tokenized_datasets[\"validation\"]\n\t\t\t\t) \ntrainer.train()\n<\/pre>\n\n\n\n5. Generate samples<\/strong><\/h2>\n\n\n\n