File size: 3,892 Bytes
371bcd8 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 | ---
tags:
- chemistry
- molecular-generation
- diffusion
- pytorch
---
# STAMP diffusion checkpoints
Inference checkpoints and matching hybrid STAMP vocabularies for [YuX-Ren/stamp_diff](https://github.com/YuX-Ren/stamp_diff), validated against source commit `8381d722fe57e6515666ea0523dd78a448790da4`.
| Model | Checkpoint | Matching vocabulary | Objective |
| --- | --- | --- | --- |
| Current model | `infill_10pct/best.pt` | `infill_10pct/vocab.json` | 50% whole-sequence and 50% interior-span augmentation |
| Baseline | `baseline_10pct/best.pt` | `baseline_10pct/vocab.json` | Whole-sequence augmentation |
Both models were trained from scratch on the same approximately 10% corpus subset. These are the best validation checkpoints from completed training runs, not optimizer resume files. The infill checkpoint is from step 39,000 and the baseline from step 40,000. Both training runs completed at step 41,377. Their validation losses use different objectives and should not be compared directly.
## Exact vocabulary contract
Each `vocab.json` is a JSON object containing **`vocab`** and **`atomic`**. The original file bytes and token ordering are preserved:
- `vocab`: 12,705 ordered tokens, with `<PAD>`, `<MASK>`, `<EOS>` at indices 0, 1, 2.
- `atomic`: 12,000 motif strings used by the hybrid tokenizer.
- The checkpoint has 12,707 tokens: the JSON vocabulary followed by `<EXPAND>` at index 12,705 and `<DELETE>` at index 12,706.
The exact check is:
```python
checkpoint["vocab"] == vocab_json["vocab"] + ["<EXPAND>", "<DELETE>"]
```
Keep the JSON vocabulary unchanged. Appending these control tokens to `vocab.json` would violate the fragment generation loader's contract. Both models use byte-identical `vocab.json` files, but each model directory includes its matching copy.
## Download and verify
Authenticate with an account that can access this private repository, then download the selected model:
```python
import json
import torch
from huggingface_hub import hf_hub_download
from stamp.diffusion.dreamon import DreamOnDiffusion
from stamp.diffusion.model import Config
repo = "bi0s/stamp-diff"
variant = "infill_10pct" # or "baseline_10pct"
checkpoint_path = hf_hub_download(repo, f"{variant}/best.pt")
vocab_path = hf_hub_download(repo, f"{variant}/vocab.json")
checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
with open(vocab_path) as handle:
vocabulary = json.load(handle)
assert checkpoint["vocab"] == vocabulary["vocab"] + ["<EXPAND>", "<DELETE>"]
assert checkpoint["config"]["vocab_size"] == len(checkpoint["vocab"])
assert set(vocabulary["atomic"]).issubset(vocabulary["vocab"])
model = DreamOnDiffusion(
Config(**checkpoint["config"]), **checkpoint["dreamon_options"]
)
model.load_state_dict(checkpoint["state"], strict=True)
model.eval()
```
Install the linked source repository to import `stamp`. These checkpoints contain the custom model state and are loaded with the project's `DreamOnDiffusion` class.
`motif_vocab.txt` is the additional principal-subgraph vocabulary used by the existing native-linker context compiler. Supply it with `--motif-vocab` when using `--context-mode native_linker`; the matching model `vocab.json` is still required with `--vocab`.
## Validation and provenance
`manifest.json` records the source revision, model configuration, objective options, file sizes and SHA-256 hashes. Both checkpoints passed exact ordered vocabulary matching, atomic-token membership checks, strict state-dict loading, embedding/output vocabulary-size checks, and finite CPU forward passes under null and high-quality conditions.
The checkpoints store BF16 weights. Model loading above creates the normal PyTorch parameter dtype and loads those weights. Data, optimizer state, smoke-test weights, incomplete full-corpus training checkpoints, logs and historical ZINC ablations are excluded.
|