StapleBridge / README.md
pranamanam's picture Jingjie00's picture
Upload Staplebridge files (#1)
bb6d2aa
|
Raw
History Blame Contribute Delete
13.9 kB

StapleBridge

Official training code for StapleBridge, a chemistry-aware framework for optimizing existing peptide leads through hydrocarbon stapling.

StapleBridge constructs a finite set of chemically and geometrically feasible stapling interventions for each peptide, learns to rank these interventions, and executes the selected plan with minimal sequence edits.

This repository contains the main StapleBridge training pipeline and one representative pretrained checkpoint.

Framework

StapleBridge framework: feasible stapling plans, exact finite-support control target, and plan execution


1. What is included

The code required to train the main StapleBridge model, and one representative pretrained checkpoint.

release/staplebridge_training/
β”œβ”€β”€ README.md
β”œβ”€β”€ THIRD_PARTY_NOTICES.md
β”œβ”€β”€ requirements.txt
β”œβ”€β”€ .gitignore
β”œβ”€β”€ configs/
β”‚   └── staplebridge_main.yaml          # configuration used for the checkpoint
β”œβ”€β”€ scripts/
β”‚   └──train.py                        # canonical training entry point
β”œβ”€β”€ staplebridge/
β”‚   β”œβ”€β”€ chemistry/                      # design state, actions, edit distance
β”‚   β”œβ”€β”€ data/                           # split loading, schemas, catalog, vocab
β”‚   β”œβ”€β”€ hydrocarbon/                    # plans, catalog, q_ref, q*, q_theta, geometry
β”‚   β”œβ”€β”€ models/                         # policy / value nets, controlled kernel
β”‚   β”œβ”€β”€ oracles/                        # ESM2 + anchor/block reference priors
β”‚   β”œβ”€β”€ reference/                      # reference energy and kernel
β”‚   β”œβ”€β”€ training/                       # stack construction, main loop, losses
β”‚   β”œβ”€β”€ integrations/                   # PeptiVerse wrapper
β”‚   └── utils/                          # paths, profiling
β”œβ”€β”€ data/
β”‚   └── README.md                       # expected input schema and file layout
└── checkpoints/
    └── staplebridge_seed42_best.pt     # representative seed=42 model

The training code in staplebridge/training/ is numerically identical to the run that produced the shipped checkpoint.

2. Environment setup

The released checkpoint was trained under:

Version
Python 3.10.20
PyTorch 2.12.1+cu130 (CUDA 13.0)
NumPy 2.0.2
PyYAML 6.0.3
RDKit 2026.03.5
transformers 4.46.0
python -m venv .venv && source .venv/bin/activate
# Install torch first, matched to your CUDA build: https://pytorch.org
pip install -r requirements.txt

3. External model dependencies

No third-party model weights are bundled. Four external resources must be provided and referenced from the config: the ESM2-650M snapshot, the PeptiVerse distribution, and the two SMILES encoders PeptiVerse depends on. Paths may be absolute, or relative to the package root.

ESM2-650M (frozen sequence context)

facebook/esm2_t33_650M_UR50D, used frozen β€” never fine-tuned. It supplies the reference-process peptide prior and the V2 plan head's anchor/local-context features. It is also the feature source for the plan head, which refuses to build without it.

huggingface-cli download facebook/esm2_t33_650M_UR50D \
    --local-dir models/esm2_t33_650M_UR50D

Then set, in configs/staplebridge_main.yaml:

reference_priors:
  peptide:
    model_name_or_path: models/esm2_t33_650M_UR50D
property_predictor:
  esm_model_name_or_path: models/esm2_t33_650M_UR50D

The config runs the prior with offline: true and strict_runtime: true, so the snapshot must already be on disk; training fails fast rather than downloading or silently substituting a fallback. Licensed by Meta under the ESM2 terms.

PeptiVerse (property oracles)

The main training objective optimises the PeptiVerse permeability-penetrance E/Z product mean. Obtain the PeptiVerse checkout and its classifier weights separately, then set:

property_predictor:
  peptiverse_root: external/PeptiVerse
  classifier_weight_root: external/PeptiVerse
  manifest_path: external/PeptiVerse/basic_models.txt

Scoring is strict: strict: true, enable_fallback: false, allow_wt_token_fallback: false. If the oracle stack cannot load, training aborts β€” it never degrades to a heuristic.

Toxicity, hemolysis and half-life are monitored only. Solubility and binding affinity are excluded from the objective.

PeptideCLM-23M and ChemBERTa-77M (required)

The basic_models.txt manifest selects predictors embedded with PeptideCLM and ChemBERTa, so both are required β€” not optional. The permeability-penetrance predictor that defines the objective is itself ChemBERTa-embedded. Loading fails fast without them.

huggingface-cli download aaronfeller/PeptideCLM-23M-all \
    --local-dir models/PeptideCLM-23M-all
huggingface-cli download DeepChem/ChemBERTa-77M-MLM \
    --local-dir models/ChemBERTa-77M-MLM
property_predictor:
  peptideclm_model_name_or_path: models/PeptideCLM-23M-all
  chemberta_model_name_or_path: models/ChemBERTa-77M-MLM

scripts/check_config.py verifies all four before training starts.

4. Data

The processed training and validation data are not included in this release and will be handled separately. No preprocessing, download or reconstruction utilities are provided.

Training reads two JSON Lines files, resolved from the config:

data/real/                 # data.root
β”œβ”€β”€ train.jsonl            # data.train_file
└── valid.jsonl            # data.valid_file

See data/README.md for the expected input schema β€” in particular the required per-residue CΞ± coordinates, which staple-geometry feasibility depends on.

The reference protocol uses 4020 training and 111 validation leads; scripts/check_config.py asserts those counts, so substituting a differently sized dataset requires relaxing the check.

5. Training command


python scripts/train.py \
    --config configs/staplebridge_main.yaml \
    --out-dir outputs/main_seed42