amelie-schreiber/esm2_loras

Trying to train LoRAs for ESM-2

esm2_loras

This is an attempt at training a Low Rank Adaptation (LoRA) for the protein language model ESM-2 for a token classification task. In particular, we attempt to train an RNA binding site predictor. There are still some issues to work out and any feedback or advice would be much appreciated. This code is for a small model so it should perform wandb sweeps for hyperparameter search in a reasonable amount of time on almost any GPU. You can easily swap out for larger models though if you want.

Model Weights and Config

The model itself

"AmelieSchreiber/esm2_t6_8M_UR50D_lora_rna_binding_sites"

can be found on Hugging Face here.

Setting up this repo

To set up the the conda environment, clone the repo and run:

conda env create -f environment.yml

Then run:

conda activate lora_esm_2

To train the model run:

from lora_esm2_script import train_protein_model

train_protein_model()

To use, try running:

from transformers import AutoModelForTokenClassification, AutoTokenizer
from peft import PeftModel
import torch
import numpy as np
import random



# Path to the saved LoRA model
model_path = "esm2_t6_8M-finetuned-lora_2023-08-03_18-32-25"
# ESM2 base model
base_model_path = "facebook/esm2_t6_8M_UR50D"

# Load the model
base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
loaded_model = PeftModel.from_pretrained(base_model, model_path)

# Load the tokenizer
loaded_tokenizer = AutoTokenizer.from_pretrained(model_path)

# New unseen protein sequence
new_protein_sequence = "FDLNDFLEQKVLVRMEAIINSMTMKERAKPEIIKGSRKRRIAAGSGMQVQDVNRLLKQFDDMQRMMKKM"

# Tokenize the new sequence
inputs = loaded_tokenizer(new_protein_sequence, truncation=True, padding='max_length', max_length=512, return_tensors="pt")

# Make predictions
with torch.no_grad():
    outputs = loaded_model(**inputs)
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=2)

# Print logits for debugging
print("Logits:", logits)

# Convert predictions to a list
predicted_labels = predictions.squeeze().tolist()

# Get input IDs to identify padding and special tokens
input_ids = inputs['input_ids'].squeeze().tolist()

# Define a set of token IDs that correspond to special tokens
special_tokens_ids = {loaded_tokenizer.cls_token_id, loaded_tokenizer.pad_token_id, loaded_tokenizer.eos_token_id}

# Filter the predicted labels using the special_tokens_ids to remove predictions for special tokens
binding_sites = [label for label, token_id in zip(predicted_labels, input_ids) if token_id not in special_tokens_ids]

print("Predicted binding sites:", binding_sites)

esm2_loras's Issues

ValueError: EsmForTokenClassification does not support gradient checkpointing

Hi Amelie, I am looking to finetune ESM model on my small downstream task. I am following your article https://huggingface.co/blog/AmelieSchreiber/esm2-ptm.

I am getting the following error on gradient_checkpointing_enable() function that you might already have a solution/suggestions.

    model.gradient_checkpointing_enable() # SP commented
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/sureshp/anaconda3/envs/qlora/lib/python3.11/site-packages/transformers/modeling_utils.py", line 1631, in gradient_checkpointing_enable
    raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
ValueError: EsmForTokenClassification does not support gradient checkpointing.

Thanks.

Incorrectly loading trained weights

Hi, this is Kiarash.

So you are incorrectly loading your already trained weights. LoRA weights are not the same as other HF weights, you need to load them on top of your base model. This was already described in the link I sent you, but for your convenience, I am sharing the exact lines of code you would need to run.

from transformers import AutoModelForTokenClassification, AutoTokenizer
from peft import PeftModel

# Path to the saved LoRA model
model_path = "esm2_t6_8M-finetuned-lora_2023-08-03_18-32-25"
# ESM2 base model
base_model_path = "facebook/esm2_t6_8M_UR50D"

# Load the model
base_model = AutoModelForTokenClassification.from_pretrained(base_model_path)
loaded_model = PeftModel.from_pretrained(base_model, model_path)

# Load the tokenizer
loaded_tokenizer = AutoTokenizer.from_pretrained(model_path)