Data problem with fine-tuning ESM2 BioNeMo

Healthcare & Life Sciences BioNeMo
1

Hello everyone,

We have encountered some difficulties while trying to fine-tune the ESM2 model.

  • We have tried to run the pretrain.py scripts from the /examples/protein/esm2nv with example data as follows:

python pretrain.py ++model.data.dataset.train=/workspace/bionemo/examples/tests/test_data/protein/train/x000.csv ++model.data.dataset.val=/workspace/bionemo/examples/tests/test_data/protein/val/x000.csv ++model.data.dataset.test=/workspace/bionemo/examples/tests/test_data/protein/test/x000.csv ++trainer.devices=1

  • with the esm1 model it worked
  • with the esm2 model, which we would actually like to use, the problem was that some files were not found of which we do not understand why they would be needed: AssertionError: Following files do not exist /workspace/bionemo/data/uniref202104_esm2/uf50/train/x000.csv , …

Why is this file needed for the ESM2 model, is there a bug somewhere in the configuration maybe?
Also, in the future we would prefer to use a data preprocessing script from bionemo instead of our own, but we noticed there are multiple preprocessing scripts available. Is there one you would specifically recommend for our purposes?

Every input is very much appreciated!

2

Hi, Is this issue resolved?

3

Hi @irits.mathew

Please use the latest version of BioNeMo Framework: BioNeMo Framework

There are also examples of performing ESM2 model pretraining: ESM-2 Pretraining - BioNeMo Framework

Thank you!

4

Hi, I am getting issues with the UniRef paths. I do have a csv files with data on which i want to finetune my ESM2 model on. I am downloading the pretrained weights. But I am unable to understand why it is looking for uf90 and uf50 path. I have set do_processing = False. I am trying to run the training on AWS sagemaker.

Is it mandatoy to have proprocessing? Can you help here?

5

Could you please share which BioNeMo Framework version you are using for ESM2 fine-tuning?

Additionally, if you are running on AWS SageMaker (with BioNeMo Framework v1. x), we have an example notebook available here for ESM2 fine-tuning: amazon-sagemaker-with-nvidia-bionemo/train-ESM.ipynb at main · aws-samples/amazon-sagemaker-with-nvidia-bionemo · GitHub

If using BioNeMo Framework v2.5, you can follow this tutorial: ESM-2 Fine-tuning - BioNeMo Framework

6

I am using v1.10 and referring to same github that you provided. Train-ESM refers to fine tuning ESM1 on UnitProt right? I was looking to implement the same for ESM2. I modified the config file to match esm2650 values, but Iam getting the uniref not found in dict error.
AlgorithmError: ExecuteUserScriptError: ExitCode 1 ErrorMessage "[rank0]: omegaconf.errors.ConfigAttributeError: Key ‘uniref90_path’ is not in struct [rank0]: full_key: data.train.uf90.uniref90_path [rank0]: object_type=dict [rank0]:

7

Looks like it is missing some parameters in the config file.

There are some guidance provided here about how to prepare your own data for ESM2: GitHub - NVIDIA/bionemo-framework at release-1.11

Please make sure you have those data paths provided in your config.

8

Yes, but do we need to provide uniref paths too? If yes, were to obtain that? I have my csv files already right? Then why do we need uniref path? and also preprocessing =True. In the github link you had shared for sagemaker its mentioned not to modify training script. Then how can I do preprocessing?

Would you be able to explain the flow?

9

In your case, Uniref files are not required. You will need to define the required paths for your input files, preprocess them (if they haven’t already been prepped), and use those for continued training using the ESM2 650m checkpoint.

The flow should look similar to what is explained in this example: Pretrain from Scratch, Continue Training from an Existing Checkpoint, and Fine-tune ESM-2nv on Custom Data — NVIDIA BioNeMo Framework

Instead of using the SAbDab_Liberis dataset, you’d use your own data. Then, make sure you have the ESM2 model checkpoint file esm2nv_650M_converted.nemo downloaded. Do preprocessing if required.

Then, provide the required config parameter/values, for example:

    ++model.data.train.custom_pretraining_fasta_path={train_fasta} \
    ++model.data.val.custom_pretraining_fasta_path={val_fasta} \
    ++model.data.test.custom_pretraining_fasta_path={test_fasta} \
    ++model.data.dataset_path={paratope_dir} \
    ++model.data.train.dataset_path={paratope_dir} \

This should resolve the issue.

10

Does it mean the input files for ESM2 should be in fasta format and not csv?