STAMP diffusion checkpoints
Inference checkpoints and matching hybrid STAMP vocabularies for 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:
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:
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.