bi0s commited on
Commit
371bcd8
·
verified ·
1 Parent(s): c8f7597

Add validated STAMP infill and baseline checkpoints with ordered vocabularies

Browse files

Two completed 10% training runs, exact vocabulary-prefix and control-token checks, strict state loading, finite CPU forward validation, and SHA-256 manifest.

README.md ADDED
@@ -0,0 +1,73 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ ---
2
+ tags:
3
+ - chemistry
4
+ - molecular-generation
5
+ - diffusion
6
+ - pytorch
7
+ ---
8
+
9
+ # STAMP diffusion checkpoints
10
+
11
+ Inference checkpoints and matching hybrid STAMP vocabularies for [YuX-Ren/stamp_diff](https://github.com/YuX-Ren/stamp_diff), validated against source commit `8381d722fe57e6515666ea0523dd78a448790da4`.
12
+
13
+ | Model | Checkpoint | Matching vocabulary | Objective |
14
+ | --- | --- | --- | --- |
15
+ | Current model | `infill_10pct/best.pt` | `infill_10pct/vocab.json` | 50% whole-sequence and 50% interior-span augmentation |
16
+ | Baseline | `baseline_10pct/best.pt` | `baseline_10pct/vocab.json` | Whole-sequence augmentation |
17
+
18
+ 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.
19
+
20
+ ## Exact vocabulary contract
21
+
22
+ Each `vocab.json` is a JSON object containing **`vocab`** and **`atomic`**. The original file bytes and token ordering are preserved:
23
+
24
+ - `vocab`: 12,705 ordered tokens, with `<PAD>`, `<MASK>`, `<EOS>` at indices 0, 1, 2.
25
+ - `atomic`: 12,000 motif strings used by the hybrid tokenizer.
26
+ - The checkpoint has 12,707 tokens: the JSON vocabulary followed by `<EXPAND>` at index 12,705 and `<DELETE>` at index 12,706.
27
+
28
+ The exact check is:
29
+
30
+ ```python
31
+ checkpoint["vocab"] == vocab_json["vocab"] + ["<EXPAND>", "<DELETE>"]
32
+ ```
33
+
34
+ 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.
35
+
36
+ ## Download and verify
37
+
38
+ Authenticate with an account that can access this private repository, then download the selected model:
39
+
40
+ ```python
41
+ import json
42
+ import torch
43
+ from huggingface_hub import hf_hub_download
44
+ from stamp.diffusion.dreamon import DreamOnDiffusion
45
+ from stamp.diffusion.model import Config
46
+
47
+ repo = "bi0s/stamp-diff"
48
+ variant = "infill_10pct" # or "baseline_10pct"
49
+ checkpoint_path = hf_hub_download(repo, f"{variant}/best.pt")
50
+ vocab_path = hf_hub_download(repo, f"{variant}/vocab.json")
51
+ checkpoint = torch.load(checkpoint_path, map_location="cpu", weights_only=True)
52
+ with open(vocab_path) as handle:
53
+ vocabulary = json.load(handle)
54
+
55
+ assert checkpoint["vocab"] == vocabulary["vocab"] + ["<EXPAND>", "<DELETE>"]
56
+ assert checkpoint["config"]["vocab_size"] == len(checkpoint["vocab"])
57
+ assert set(vocabulary["atomic"]).issubset(vocabulary["vocab"])
58
+ model = DreamOnDiffusion(
59
+ Config(**checkpoint["config"]), **checkpoint["dreamon_options"]
60
+ )
61
+ model.load_state_dict(checkpoint["state"], strict=True)
62
+ model.eval()
63
+ ```
64
+
65
+ Install the linked source repository to import `stamp`. These checkpoints contain the custom model state and are loaded with the project's `DreamOnDiffusion` class.
66
+
67
+ `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`.
68
+
69
+ ## Validation and provenance
70
+
71
+ `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.
72
+
73
+ 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.
baseline_10pct/best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:39761af0325c0b5e9dea1f630d41549bbf0b25b486de5d6b123f82eb80e3a7b1
3
+ size 212344188
baseline_10pct/vocab.json ADDED
The diff for this file is too large to render. See raw diff
 
infill_10pct/best.pt ADDED
@@ -0,0 +1,3 @@
 
 
 
 
1
+ version https://git-lfs.github.com/spec/v1
2
+ oid sha256:227143d0e32129a81dcc45771676012d9d818eff7777d2928f6951aa062b0ce1
3
+ size 212344252
infill_10pct/vocab.json ADDED
The diff for this file is too large to render. See raw diff
 
manifest.json ADDED
@@ -0,0 +1,139 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ {
2
+ "format_version": 1,
3
+ "source_repository": "https://github.com/YuX-Ren/stamp_diff",
4
+ "source_commit": "8381d722fe57e6515666ea0523dd78a448790da4",
5
+ "recommended_model": "infill_10pct",
6
+ "vocabulary_rule": "checkpoint['vocab'] == vocab_json['vocab'] + ['<EXPAND>', '<DELETE>']",
7
+ "models": [
8
+ {
9
+ "name": "infill_10pct",
10
+ "source_run": "large_pretrain_infill_20261015",
11
+ "checkpoint_step": 39000,
12
+ "training_completion": {
13
+ "step": 41377,
14
+ "nominal_steps": 41559,
15
+ "best_validation": 1.5793923139572144,
16
+ "discarded_tail_microbatches": 3,
17
+ "segment_start_consumed_batches": 0
18
+ },
19
+ "config": {
20
+ "vocab_size": 12707,
21
+ "max_length": 256,
22
+ "width": 768,
23
+ "heads": 12,
24
+ "layers": 12,
25
+ "dropout": 0.1,
26
+ "pad_weight": 0.1,
27
+ "condition_dropout": 0.15,
28
+ "boundary_weight": 1.0
29
+ },
30
+ "dreamon_options": {
31
+ "merge_probability": 0.5,
32
+ "max_delete": 16,
33
+ "infill_probability": 0.5
34
+ },
35
+ "tokenization": "hybrid-v1",
36
+ "vocab_tokens": 12705,
37
+ "checkpoint_tokens": 12707,
38
+ "atomic_motifs": 12000,
39
+ "controls": [
40
+ {
41
+ "token": "<EXPAND>",
42
+ "index": 12705
43
+ },
44
+ {
45
+ "token": "<DELETE>",
46
+ "index": 12706
47
+ }
48
+ ],
49
+ "validation": {
50
+ "ordered_vocab_matches": true,
51
+ "atomic_tokens_in_vocab": true,
52
+ "state_dict_strict_load": true,
53
+ "token_embedding_rows_match": true,
54
+ "output_rows_match": true,
55
+ "cpu_forward_finite": true
56
+ },
57
+ "files": [
58
+ {
59
+ "path": "infill_10pct/best.pt",
60
+ "size_bytes": 212344252,
61
+ "sha256": "227143d0e32129a81dcc45771676012d9d818eff7777d2928f6951aa062b0ce1"
62
+ },
63
+ {
64
+ "path": "infill_10pct/vocab.json",
65
+ "size_bytes": 664430,
66
+ "sha256": "e7a6dea85773f86977a5daef527de30ee711a35978c69c1f0f689576fc199349"
67
+ }
68
+ ]
69
+ },
70
+ {
71
+ "name": "baseline_10pct",
72
+ "source_run": "large_pretrain_20261013",
73
+ "checkpoint_step": 40000,
74
+ "training_completion": {
75
+ "step": 41377,
76
+ "nominal_steps": 41559,
77
+ "best_validation": 2.1470673084259033,
78
+ "discarded_tail_microbatches": 3,
79
+ "segment_start_consumed_batches": 0
80
+ },
81
+ "config": {
82
+ "vocab_size": 12707,
83
+ "max_length": 256,
84
+ "width": 768,
85
+ "heads": 12,
86
+ "layers": 12,
87
+ "dropout": 0.1,
88
+ "pad_weight": 0.1,
89
+ "condition_dropout": 0.15,
90
+ "boundary_weight": 1.0
91
+ },
92
+ "dreamon_options": {
93
+ "merge_probability": 0.5,
94
+ "max_delete": 16
95
+ },
96
+ "tokenization": "hybrid-v1",
97
+ "vocab_tokens": 12705,
98
+ "checkpoint_tokens": 12707,
99
+ "atomic_motifs": 12000,
100
+ "controls": [
101
+ {
102
+ "token": "<EXPAND>",
103
+ "index": 12705
104
+ },
105
+ {
106
+ "token": "<DELETE>",
107
+ "index": 12706
108
+ }
109
+ ],
110
+ "validation": {
111
+ "ordered_vocab_matches": true,
112
+ "atomic_tokens_in_vocab": true,
113
+ "state_dict_strict_load": true,
114
+ "token_embedding_rows_match": true,
115
+ "output_rows_match": true,
116
+ "cpu_forward_finite": true
117
+ },
118
+ "files": [
119
+ {
120
+ "path": "baseline_10pct/best.pt",
121
+ "size_bytes": 212344188,
122
+ "sha256": "39761af0325c0b5e9dea1f630d41549bbf0b25b486de5d6b123f82eb80e3a7b1"
123
+ },
124
+ {
125
+ "path": "baseline_10pct/vocab.json",
126
+ "size_bytes": 664430,
127
+ "sha256": "e7a6dea85773f86977a5daef527de30ee711a35978c69c1f0f689576fc199349"
128
+ }
129
+ ]
130
+ }
131
+ ],
132
+ "additional_files": [
133
+ {
134
+ "path": "motif_vocab.txt",
135
+ "size_bytes": 2034995,
136
+ "sha256": "c9bc4f175df775a7b6603a26780a3033eebefd705dee0a1ffc4410d6373b3e32"
137
+ }
138
+ ]
139
+ }
motif_vocab.txt ADDED
The diff for this file is too large to render. See raw diff