Antibody Design — Core Code

This repository contains the core training and evaluation code accompanying our paper on antigen-conditioned antibody design. It builds on the ESM3 protein language model with a three-stage recipe:

1. CPT — continual pre-training of the ESM3 backbone on antibody structures.
2. SFT — supervised fine-tuning for antigen-conditioned CDR design, with an adapter that injects antigen structure features (GearNet).
3. DPO — Direct Preference Optimization on CDR designs ranked by a preference signal, starting from the SFT checkpoint.

Generation and structure-based evaluation scripts are provided to reproduce the design and metric pipeline (RMSD, AAR, sequence similarity, PHR, pseudo log-likelihood).

Directory layout

src/
  cpt/                          Continual pre-training
    config/CPT.yaml             Training configuration
    training/trainer.py         Entry point  (accelerate launch)
    training/config.py          Config dataclasses + loader
    training/masking.py         CDR / structure-token masking
    data_pipeline/npz_dataset.py  NPZ structure dataset (runtime)
  sft/                          Supervised fine-tuning
    config/SFT.yaml             Training configuration
    training/trainer.py         Entry point  (accelerate launch)
    training/config.py          Config dataclasses + loader
    training/masking.py         Masking strategies
    training/adapter_gear.py    ESM3 wrapper + antigen (GearNet) adapter
    data_pipeline/sabdab_data_imgt.py   SAbDab IMGT dataset
    data_pipeline/rabd_data_test.py     RAbD benchmark dataset (used by eval)
    data_pipeline/sabdab_imgt/          IMGT preprocessing sub-package
  dpo/                          Direct Preference Optimization
    configs/cdr_dpo.yaml        Training configuration
    train.py                    Entry point  (accelerate launch)
    trainer.py                  CDR-DPO trainer (policy + reference model)
    config.py                   Config dataclasses + loader
    model_loader.py             Builds policy / reference ESM3 models
    trainer_utils.py            Checkpoint + trainer state helpers
    preference_dataset.py       Paired / ranked preference datasets
                                (reuses sft.* model & data layers)
  eval/                         Generation + evaluation
    generate.py                 Reference-guided design generation
    generate_zero_interface.py  Variant without antigen-interface conditioning
    evaluate.py                 Structure-based metrics (single heavy chain)
    evaluate_hl.py              Variant for merged heavy+light generated chain
scripts/
  run_cpt.sh                    Launch CPT
  run_sft.sh                    Launch SFT
  run_dpo.sh                    Launch DPO
  run_generate.sh               Generate -> evaluate -> PLL (end to end)
  run_evaluate.sh               Evaluate an existing directory of PDBs
  recompute_pll.py              Pseudo log-likelihood (AntiBERTy)

Requirements

• Python 3.10+, PyTorch (CUDA), accelerate
• ESM3 and its model weights
• biopython, numpy, pandas, tqdm, pyyaml
• AntiBERTy (for PLL)
• US-align executable (for structural alignment)
• GearNet antigen structure features
• Optional: MMseqs2 for sequence clustering, Weights & Biases for logging

Data

Training uses the SAbDab antibody database (IMGT-numbered structures) and evaluation uses the RAbD benchmark. Download these separately and point the config / CLI paths at them.

Configure paths before running

The code is released with placeholder paths. Replace them with your own:

• /path/to/antibody-design — the repository root
• /path/to/USalign/USalign — the US-align executable
• /path/to/esm3/data — ESM3 / structure data root
• In src/cpt/config/CPT.yaml and src/sft/config/SFT.yaml: dataset paths, model.model_path (CPT checkpoint used by SFT), antigen_feature_path, and the wandb section (wandb_api_key is read from the WANDB_API_KEY environment variable; set wandb_enabled: false to skip logging).
• In src/eval/{evaluate,evaluate_hl}.py: the DESIGNED_DIR / ORIGINAL_DIR / ALIGNED_DIR / RESULTS_DIR / USALIGN_EXECUTABLE constants, or pass them as CLI arguments.
• In src/dpo/configs/cdr_dpo.yaml: data.train_csv (ranked preference pairs CSV), model.pretrained_model_path (the SFT checkpoint), and checkpoint.output_dir.

The training/eval entry points add src/ to sys.path, so run all commands from the repository root.

Usage

# 1. Continual pre-training
bash scripts/run_cpt.sh

# 2. Supervised fine-tuning (set model.model_path in SFT.yaml to the CPT checkpoint)
bash scripts/run_sft.sh

# 3. DPO (set model.pretrained_model_path in cdr_dpo.yaml to the SFT checkpoint)
bash scripts/run_dpo.sh

# 4. Generate designs, evaluate, and compute PLL
bash scripts/run_generate.sh 0          # GPU id

# Evaluate an existing directory of generated PDBs
bash scripts/run_evaluate.sh datasets/results/run_demo/full/tmp0.6 full

Both trainers accept --config <path> to override the default YAML (src/cpt/config/CPT.yaml, src/sft/config/SFT.yaml); the DPO entry point requires --config (default src/dpo/configs/cdr_dpo.yaml). generate.py exposes --checkpoint, --cdr-mode {h3,full}, --temperature, --num-targets, --samples-per-target, and --output-dir (see --help for the full list).