Pretraining & Finetuning of Protein Language Model with BioNeMO & WandB JP

docker login nvcr.io
Username: $oauthtoken
Password <insert NGC API token here>

docker pull nvcr.io/nvidia/clara/bionemo-framework:1.5

docker run -it --rm --gpus all nvcr.io/nvidia/clara/bionemo-framework:1.5 bash

wget --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/ngc-apps/ngc_cli/versions/3.42.0/files/ngccli_linux.zip -O ngccli_linux.zip && unzip ngccli_linux.zip
find ngc-cli/ -type f -exec md5sum {} + | LC_ALL=C sort | md5sum -c ngc-cli.md5
sha256sum ngccli_linux.zip
chmod u+x ngc-cli/ngc
echo "export PATH=\"\$PATH:$(pwd)/ngc-cli\"" >> ~/.bash_profile && source ~/.bash_profile
ngc config set

export BIONEMO_HOME="/workspace/bionemo"
export WANDB_API_KEY=<your wandb API key> # SaaSの場合は、https://wandb.ai/authorizeより取得可能です 

# もしdedicated cloud版またはオンプレミス版のwandbを使用している場合は、"WANDB_BASE_URL"を適切に設定してください。
# あなたの会社のwandb管理者が"WANDB_BASE_URL"をご存知です。
[optional] export WANDB_BASE_URL=<your wandb base url>

import os
os.environ["WANDB_ENTITY"]="<your team where you want to log>" 
os.environ["WANDB_PROJECT"]="BioNeMo_protein_LLM_pretraining"

!python examples/protein/esm2nv/pretrain.py\
 --config-path=conf\
 --config-name=pretrain_esm2_650M\
 ++do_training=False\
 ++exp_manager.wandb_logger_kwargs.name='preproceed_data_upload'\
 ++exp_manager.wandb_logger_kwargs.job_type='data_upload'\
 ++wandb_artifacts.wandb_use_artifact_path='${oc.env:WANDB_ENTITY}/${oc.env:WANDB_PROJECT}/uniref202104_esm2_qc_test200_val200:v0'\
 ++wandb_artifacts.wandb_log_artifact_name='uniref202104_esm2_qc_test200_val200_preprocessed'\
 ++model.data.val_size=500\
 ++model.data.test_size=100\
 ++model.data.uf50_datapath=/uniref50_train_filt.fasta\
 ++model.data.uf90_datapath=/ur90_ur50_sampler.fasta\
 ++model.data.cluster_mapping_tsv=/mapping.tsv\
 ++model.data.dataset_path=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf50\
 ++model.data.uf90.uniref90_path=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uf90\
 ++model.data.train.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta\
 ++model.data.train.uf90_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/ur90_ur50_sampler.fasta\
 ++model.data.train.cluster_mapping_tsv=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/mapping.tsv\
 ++model.data.val.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta\
 ++model.data.test.uf50_datapath=/workspace/bionemo/examples/tests/test_data/uniref202104_esm2_qc_test200_val200/uniref50_train_filt.fasta

!python examples/protein/esm2nv/pretrain.py \
 --config-path=conf \
 --config-name=pretrain_esm2_650M \
 ++do_training=True \
 ++do_testing=True \
 ++wandb_artifacts.wandb_use_artifact_path='${oc.env:WANDB_ENTITY}/${oc.env:WANDB_PROJECT}/uniref202104_esm2_qc_test200_val200_preprocessed:v0'\
 ++model.data.dataset_path=/ \
 ++model.data.uf90.uniref90_path=/uf90 \
 ++trainer.devices=1 \
 ++model.tensor_model_parallel_size=1 \
 ++model.micro_batch_size=4 \
 ++trainer.max_steps=1 \
 ++trainer.val_check_interval=1 \
 ++exp_manager.create_wandb_logger=True \
 ++exp_manager.checkpoint_callback_params.save_top_k=10

import os
os.environ["WANDB_ENTITY"]="<your team where you want to log>" 
os.environ["WANDB_PROJECT"]="BioNeMo_protein_LLM_finetuning"

!wget -q -O /tmp/ngccli_linux.zip --content-disposition https://api.ngc.nvidia.com/v2/resources/nvidia/ngc-apps/ngc_cli/versions/3.38.0/files/ngccli_linux.zip && unzip -o /tmp/ngccli_linux.zip -d /tmp && chmod u+x /tmp/ngc-cli/ngc && rm /tmp/ngccli_linux.zip

!python download_models.py --download_dir /workspace/bionemo/models esm2nv_650m

!python download_models.py --download_dir /workspace/bionemo/models esm2nv_650m
# save model
import wandb
with wandb.init(name="model_upload") as run:
    artifact = wandb.Artifact(
        name="esm2nv_650m",
        type="model",
        description="original esm2nv_650m",
        metadata={"path":"bionemo/models/esm2nv_650M_converted.nemo"},
    )
    artifact.add_file("/workspace/bionemo/models/esm2nv_650M_converted.nemo")
    run.log_artifact(artifact)

# save dataset 
with wandb.init(name="data_upload") as run:
    artifact = wandb.Artifact(
        name="downstream_taskdataset",
        type="dataset",
        description="bionemo/examples/tests/test_data/protein/downstream",
        metadata={"path":"/workspace/bionemo/examples/tests/test_data/protein/downstream"},
    )
    artifact.add_dir("/workspace/bionemo/examples/tests/test_data/protein/downstream")
    run.log_artifact(artifact)

!python examples/protein/downstream/downstream_flip.py\
 --config-path="../esm2nv/conf"\
 --config-name=downstream_sec_str_LORA\
 ++trainer.max_epochs=15\
 ++wandb_artifacts.wandb_use_artifact_data_path='${oc.env:WANDB_ENTITY}/${oc.env:WANDB_PROJECT}/downstream_taskdataset:v0'\
 ++wandb_artifacts.wandb_use_artifact_model_path='${oc.env:WANDB_ENTITY}/${oc.env:WANDB_PROJECT}/esm2nv_650m:v0'\
 ++model.data.dataset_path=/\
 ++model.restore_encoder_path=/esm2nv_650M_converted.nemo