# 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](figure/framework.png)](figure/framework.pdf)

---

## 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. What is intentionally not included

This release covers **training only**. Not included:

- Baseline and ablation implementations, and their checkpoints.
- Alternative stapling chemistries.
- Downstream and post-hoc evaluation pipelines, including the test-set
  evaluator, structural and docking analyses, and permeability benchmarking.
- Analysis, figure-generation and manuscript material.
- Training outputs: logs, cached predictions, intermediate checkpoints.
- **Processed training and validation data** — see §5; these will be handled
  separately.

One checkpoint is shipped: the representative seed=42 model.

Reproducing the paper's *evaluation* numbers requires the evaluation pipelines,
which are outside this release. What is here reproduces *training*. -->

## 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 |

```bash
python -m venv .venv && source .venv/bin/activate
# Install torch first, matched to your CUDA build: https://pytorch.org
pip install -r requirements.txt
```

<!-- StapleBridge's own model is small (~142k parameters total) and runs on CPU. A
GPU is needed in practice because the frozen ESM2-650M prior and the PeptiVerse
oracles are invoked per candidate. Set the device in one place — the config
requires `training.device`, `property_predictor.device` and
`reference_priors.peptide.device` to agree, and `check_config.py` enforces it. -->

## 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.

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

Then set, in `configs/staplebridge_main.yaml`:

```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:

```yaml
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.

```bash
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
```

```yaml
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`](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

```bash

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

<!-- `check_config.py` is read-only: it verifies the configuration and the presence of
the external assets without loading a model or training. Run it first — it turns
a misconfiguration into an immediate error rather than a failure hours in.

The CLI is intentionally small — **seed (42) and device come from the config**,
not from flags, so a run cannot silently diverge from the recorded protocol.
`--resume <checkpoint>` restores model, optimizer and RNG state so the remaining
epochs match an uninterrupted run. `train.py` refuses to start in a non-empty
`--out-dir`.

Written to `--out-dir`: `resolved_config.yaml`, `metrics.jsonl` (per-epoch),
`run_summary.json`, per-epoch validation summaries, and
`checkpoints/{best_kl,best_pv,latest,epoch_NNN}.pt`. **`best_kl.pt` is the
selected model** — see §7. -->

<!-- ## 7. Main reproducibility settings

Read from `configs/staplebridge_main.yaml` and asserted by
`scripts/check_config.py`:

| Setting | Value |
| --- | --- |
| Splits | `train.jsonl` (4020) / `valid.jsonl` (111), full validation every epoch |
| Seed | 42 |
| Epochs | 10, no early stopping (all 10 always run) |
| Horizon | 8 (training and validation must match) |
| Max neighboring actions | 128 |
| Chunk size | 32 |
| Committed plans per lead | 4 |
| `beta` (Exact-SB) | 1.0 |
| Plan-loss weight | 1.0 |
| Plan head | hidden dim 128, frozen ESM2-650M features |
| Objective | `permeability_penetrance`, E/Z product mean, neutral-canonical SMILES |
| Monitored only | toxicity, hemolysis, half-life |
| Excluded | solubility, binding affinity |
| PeptiVerse | strict SMILES mode, all fallbacks disabled |
| Catalog | hydrocarbon only (`include_optional: false`): S5-S5/i,i+4 and R8-S5/i,i+7 |
| Hard constraints | chemistry + geometry + edit budget 6.0, min sequence identity 0.6, protected edits forbidden, exact committed-plan completion |
| Decoding | strict hierarchical plan-first (`hierarchical_plan_ranking: true`) |
| **Checkpoint selection** | **minimum validation `q_star_vs_q_theta_kl`**, guarded by `require_both_topologies` |

Optimizer: Adam, lr 1e-3, grad-norm clip 1.0.

### Checkpoint selection

The rule is: the epoch minimising `KL(q* ‖ q_theta)` on the full 111-lead
validation split, among epochs where both staple topologies appear in the
selected candidates. Validation permeability and the test split play no part in
selecting it.

For the shipped checkpoint that rule chose **epoch 9**, validation
`q_star_vs_q_theta_kl = 0.081568` — the minimum over all 10 epochs (epoch 10 rose
to 0.082585). `train.py` also maintains `best_pv.pt` by validation
delta-penetrance for monitoring; it is **not** the selected model and is not
shipped.

## 8. Checkpoint loading

`checkpoints/staplebridge_seed42_best.pt` is the `best_kl.pt` of the seed=42
training run, copied byte-for-byte with weights unmodified
(sha256 `22d273164a3132d43617c51649947028f11859f90ca0bee28540edc9ad62a298`).

Verify it loads against this package's model definitions: -->

<!-- ```bash
python scripts/smoke_test.py
```

This checks imports, config, model construction and a strict checkpoint load
(zero missing / unexpected keys). It is not an evaluation suite.

Minimal manual load:

```python
import torch, yaml
from staplebridge.hydrocarbon.plan_control import build_hydrocarbon_plan_head
from staplebridge.models.policy_net import PolicyNet
from staplebridge.models.value_net import ValueNet

config = yaml.safe_load(open("configs/staplebridge_main.yaml"))
emb_dim = int(config["model"]["emb_dim"])

ckpt = torch.load("checkpoints/staplebridge_seed42_best.pt",
                  map_location="cpu", weights_only=False)

policy = PolicyNet(emb_dim=emb_dim)
value = ValueNet(emb_dim=emb_dim)
policy.load_state_dict(ckpt["policy_state_dict"])   # strict, exact match
value.load_state_dict(ckpt["value_state_dict"])
print("epoch:", ckpt["epoch"])                      # -> 9
```

The plan head is the one component that cannot be built without the frozen
ESM2-650M prior — it refuses a stand-in by design, so training can never
silently substitute a different context model:

```python
from staplebridge.training.stack import build_stack
stack = build_stack(config, seed=42)                # requires ESM2 on disk
head = build_hydrocarbon_plan_head(config, emb_dim, "cpu",
                                   esm2_prior=stack["reference_priors"].peptide)
head.load_state_dict(ckpt["plan_head_state_dict"])
```

Checkpoint payload: `policy_state_dict` (20 tensors), `value_state_dict` (20),
`plan_head_state_dict` (6), plus non-model metadata — `epoch`, the full resolved
`config`, `plan_control_enabled`, `optimizer_state_dict`, and RNG state
(`plan_rng_state`, `python_random_state`, `numpy_random_state`,
`torch_rng_state`, `cuda_rng_state_all`). The optimizer and RNG entries exist so
`--resume` can continue a run bit-identically; they are not model weights.

## 9. Reproducibility note

- **Seed.** Fixed at 42 in the config and applied to Python, NumPy and Torch
  (including all CUDA devices). Plan sampling uses its own seeded
  `random.Random(42)`, checkpointed so `--resume` continues the same stream.
- **Determinism.** `torch.use_deterministic_algorithms(False)`, matching the
  reference run. Results are therefore **device- and version-dependent**: a
  different GPU model, CUDA version, or PyTorch build can shift metrics
  slightly. Exact bit-level reproduction requires the environment in §3.
- **External models.** The frozen ESM2-650M snapshot and the PeptiVerse
  classifier stack are part of the training objective. Different versions of
  either change the optimisation target and will not reproduce these numbers.
- **Caches.** Training maintains a persistent ESM2 prior cache and an exact-SB
  target cache (SQLite, under `outputs/cache/`). These are pure speedups —
  deterministic values keyed by content — and do not change results.
- **Cost.** 10 epochs over 4020 leads with full 111-lead validation each epoch.
  The reference run totalled about 3.9 h of epoch time on one GPU (~16 GiB peak).
  Cache warming dominates the first epoch — 4964 s, against 760-890 s for later
  epochs once the ESM2 and exact-SB caches are populated.
- **Scale guard.** `train.py` hard-asserts `train_n == 4020` and `epochs == 10`,
  so the protocol cannot be shrunk by editing the config alone.

## License

See [`THIRD_PARTY_NOTICES.md`](THIRD_PARTY_NOTICES.md) for third-party
components. A license for the StapleBridge code itself has not yet been
selected; add a `LICENSE` file before publishing. -->