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