EpiTune

EpiTune is a state-of-the-art b-cell linear epitope prediction tool created by fine-tuning ESM-2 with LoRA on IEDB epitope data, created as part of the "EpiTune: An Accurate Epitope Prediction Model with Mechanistic Insights" publication.

The 'linear' and 'conformational' directories contain the fine-tuning, validation, and inference scripts for the respective linear and conformational versions of EpiTune. Additionally, the linear folder contains the complete EpiTune-Linear-35m model, while the conformational folder contains the EpiTune-Conf-35m model.

This repo features the code and data to fully reproduce all linear and conformational versions of EpiTune 35m, 650m, and 15B. If you wish to access the model weights of the larger models for inference, please reach out to the following points of contact: chiril.calin@pfizer.com, bradley.perrinjr@pfizer.com

SETUP

Setup the environment file within a conda environment.

You can either create a new environment and install from requirements.txt,

conda create -n epitune python=3.11

conda activate epitune

Or, run only this last step if in an existing environment.

pip install -r requirements.txt

Setup Notes:

-The requirements are designed for CUDA 13.0 and may not work for other CUDA versions.

-For some versions of Linux, building from requirements.txt will not properly build the wheel for gmpy2. This can usually be remedied by running

conda install -c conda-forge gmpy2

before building the requirements file.

-If you do not have CUDA toolkit, such as if building requirements.txt with a new environment, you may receive an error when trying to run with EpiTune with CUDA. If this occurs, the easiest remedy is by running

conda install nvidia:cuda-toolkit

in your environment, although manual installation of the CUDA toolkit can also be done here: https://developer.nvidia.com/cuda/toolkit for best practice.

-This library will most likely not work on Windows due to several package dependencies not supported for the OS, at least not without significant manual environment setup. WSL is a confirmed working alternative in this case.

INFERENCE

Demo inference can be launched as so from either the conformational or linear folders:

python linear_inference.py

FINE-TUNING

Additional setup for fine-tuning:

The fine-tuning scripts ('tune_epitune.py') in each model directory should be run with the "accelerate" library. While the accelerate library is included, you will need to configure accelerate to your system post-environment setup by running:

accelerate config

Depending on your configuration, you may see different model results. Reference for config used for EpiTune located at bottom of readme.

Furthermore, you will need to download the ESM-2 base model weights.

This should happen automatically with the running of the fine-tuning script, where the transformers library will download the base model from huggingface. (These base models are pretty big, so you might want to delete them if you are no longer using EpiTune! Huggingface's model downloads are by default stored to .cache/huggingface/hub)

To achieve full reproducibility, removing non-determinism requires a CUBLAS_WORKSPACE_CONFIG prefix (to allow torch.use_deterministic_algorithms):

CUBLAS_WORKSPACE_CONFIG=:4096:8 accelerate launch tune_epitune.py

Otherwise, simply launch as

accelerate launch tune_epitune.py

Note: Running the algorithms with determinism enabled may have performance impacts (see PyTorch documentation on using deterministic algorithms).

Finding Best Checkpoint

Depending on setup, the HuggingFace trainer may report different metrics during training, and metrics should be calculated after training to find the best checkpoint. This can be done with

python find_best_checkpoint.py

---

---

---

Accelerate config options chosen for EpiTune:

(These should be modified to fit your system if not attempting to reproduce exact results):

In which compute environment are you running?

This Machine

Which type of machine are you using?

multi-GPU

How many different machines will you use?

1

Check distributed operations?

no

Torch dynamo?

no

DeepSpeed?

yes

Specify json?

no

Deepspeed ZeRO optimization stage?

3

Where to offload optimizer states?

None

Where to offload parameters?

cpu

How many grad accumulation steps?

1

Gradient clipping?

no

Save 16-bit model weights w/ ZeRO Stage-3?

yes

Enable deepspeed.zero.Init w/ ZeRO Stage-3?

yes

MoE training?

no

How many GPU(s) for distributed training?

4

Mixed precision?

fp16

---