From LLaMA to ProLLaMA, GPT to ProtGPT: Revolutionizing Drug Discovery with Protein Language Models

Source PDB

In the rapidly evolving landscape of artificial intelligence, Natural Language Processing (NLP) has been a cornerstone of advancements, leading to the development of sophisticated models like OpenAI’s GPT series and Meta’s LLaMA. These models, primarily designed for human language understanding and generation, have showcased remarkable capabilities, setting the stage for their application in various fields. One of the most promising and transformative applications lies at the intersection of NLP and protein science, giving rise to specialized models like ProLLaMA and ProtGPT. These protein language models (PLMs) are not only enhancing our understanding of protein sequences but also accelerating the discovery of novel proteins and drugs. In this blog, we will explore how Gen AI is revolutionizing protein engineering and provide a detailed guide on how to finetune GPT or LLaMa models to generate Novel protein Sequences, including an explanation of the tokenizer used in PLMs.

Understanding Protein Language Models

Protein language models are designed to understand and generate sequences of amino acids in proteins, much like how traditional NLP models handle words and sentences. By leveraging large datasets of protein sequences and structures, PLMs can learn the intricate patterns and relationships inherent in biological data. This capability allows them to predict protein functions, interactions, and even the effects of mutations with high accuracy.

The Journey from Natural Language Processing to Protein Language Models

Natural Language Processing (NLP) has witnessed tremendous growth with models like Llama and GPT, which have demonstrated unprecedented capabilities in understanding and generating human language. These models are built on architectures like Transformers, which excel at capturing long-range dependencies and contextual information. Inspired by the success of NLP, researchers have adapted these architectures to the field of protein engineering. Proteins, like natural languages, have their own “grammar” encoded in the sequence of amino acids. By treating protein sequences as a language, PLMs like Prollama and ProtGPT have been developed, allowing for powerful predictions and generation of novel protein sequences.

Meta’s LLaMA (Large Language Model Meta AI) has been instrumental in advancing NLP, demonstrating the power of transformer architectures in understanding and generating human language. The transition to ProLLaMA involves tailoring these architectures to the specific nuances of protein sequences. ProLLaMA utilizes the same underlying principles but is trained on vast datasets of protein sequences, enabling it to capture the complexities of protein folding, function, and interaction. Similarly, OpenAI’s GPT (Generative Pre-trained Transformer) series has set benchmarks in NLP, with models like GPT-3 showcasing human-like text generation and comprehension. ProtGPT represents a leap towards applying these advancements to protein science. By fine-tuning GPT models on protein data, ProtGPT can predict protein sequences, interactions, and functions with remarkable accuracy.

Building Your Own Protein Language Model

Since proteins can be thought of like sequences of text, there’s been a lot of excitement about using these models to study proteins. If you’re curious about creating your own protein language model, it’s easier than you might think. With the HuggingFace Transformers library, you can quickly build a protein language model from scratch. The library lets you choose a model type, prepare your protein data, and train the model with just a few lines of code. It runs on PyTorch (or Tensorflow, if you prefer), giving you the option to dive deeper into customizing the model or let the library handle everything for you.

1. Install the Required Libraries

Before diving into the code, you need to set up your environment with the necessary libraries. We’ll be using the Transformers and Datasets libraries from HuggingFace, along with PyTorch.

It’s recommended to install these libraries using pip, especially if you’re working within a Conda environment:

pip install transformers
pip install datasets

You’ll also need to install PyTorch. You can find system-specific installation instructions here.

2. Prepare Your Dataset

With your dataset of protein sequences ready, you’ll need to split it into training, validation, and test sets. We’ll use HuggingFace’s Datasets library for efficient tokenization and data handling.

Tokenization

In protein language models, tokenization typically happens at the character level (i.e., amino acid residue). You won’t need to train a custom tokenizer like in traditional text-based language models. Instead, you can use a pre-built tokenizer that includes tokens for all 20 amino acids, along with padding tokens and tokens for the N-terminus and C-terminus.

from transformers import AutoTokenizer
# Load the tokenizer
tokenizer = AutoTokenizer.from_pretrained('path/to/tokenizer', use_fast=True)

Load the Dataset

Next, you’ll load your train, validation, and test datasets. The Datasets library offers convenient features like caching, which avoids repeating data preparation steps between training runs.

from datasets import load_dataset
# Define dataset paths
data_set_paths = {
"train": "/path/to/train.csv",
"test": "/path/to/test.csv",
"val": "/path/to/val.csv"
}
# Load datasets
datasets = load_dataset('csv', data_files=data_set_paths, cache_dir='/where/to/store/cache')

3. Define the Model Architecture

Now that your data is ready, it’s time to define the model architecture. We’ll use GPT-2 for this example, but you can choose any model architecture from the Transformers library. Since we’re training the model from scratch, we won’t load any pre-trained weights.

from transformers import GPT2Config, GPT2LMHeadModel
# Configure the model
transformer_config = GPT2Config(
vocab_size = tokenizer.vocab_size,
  n_layer=12,
  n_embd=512,
  n_head=12,
  n_inner=2048
)
# Initialize the model
model = GPT2LMHeadModel(config=transformer_config)

If you’re using a GPU, remember to move the model to CUDA:

model.to('cuda')

4. Train Your Language Model

With the model architecture in place, you’re ready to start training. HuggingFace’s Trainer class simplifies the process, allowing you to specify training parameters through a TrainingArguments instance.

from transformers import Trainer, TrainingArguments
# Define training arguments
training_args = TrainingArguments(
    output_dir="./output",
    overwrite_output_dir=True,
    num_train_epochs=2,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    save_steps=3000,
    save_total_limit=2
)
# Initialize the trainer
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=datasets["train"],
    eval_dataset=datasets["val"]
)
# Start training
trainer.train()

5. Generate Protein Sequences

After training, you can generate new protein sequences using the trained model. There are several methods for generating sequences, and here we’ll use top-p sampling.

#load the Checkpoint
model = GPT2LMHeadModel.from_pretrained("pth/to/ckpt")
model.to(torch_device)
#Give some start tokens
model_inputs = tokeniser(f'M', return_tensors='pt').to(torch_device)
greedy_output = model.generate(**model_inputs, max_new_tokens=250,num_beams=30, num_return_sequences=20)
decode = tokeniser.decode(greedy)

And that’s it! You’ve successfully trained a protein language model from scratch and generated new protein sequences.

Conclusion

Generative AI, through the development of PLMs like Prollama and ProtGPT, is revolutionizing the field of protein engineering. These models enable the rapid design, prediction, and understanding of proteins, opening new avenues for drug discovery and therapeutic development. By following the steps outlined above, one can build their own PLM and contribute to this exciting frontier of science. With the continuous advancements in AI and computational biology, the future of protein engineering looks incredibly promising, with endless possibilities for innovation and discovery.

Feel free to experiment with different architectures, hyperparameters, and datasets to further enhance your PLM. The journey from Llama to Prollama and GPT to ProtGPT is just the beginning of a new era in protein engineering.