Document ChEMBL3D and GEOM checkpoints with shared inference quickstart

README.md

@@ -5,93 +5,179 @@ tags:
   - molecular-generation
   - flow-matching
   - chembl3d
+  - geom-drugs
+  - 3d
 ---
 
-# Morph ChEMBL3D checkpoints
+# Morph: 3D molecular generation
 
-Six checkpoints for unconditional 3D molecular generation, with matching
-vocabularies, training distributions, portable inference configurations and
-the frozen Morph runtime from each training run.
+Unconditional molecular generation checkpoints trained on **ChEMBL3D** and
+**GEOM-Drugs**, with the configurations, vocabularies and training distributions
+needed for inference. No training dataset is required to load a model or generate
+molecules.
 
-| Variant | Completed training epochs | Directory |
-| --- | --- | --- |
-| Explicit hydrogens | 200, 300, 400 | [chembl3d/explicit_h](chembl3d/explicit_h) |
-| Implicit hydrogens (heavy-atom graphs) | 100, 200, 300 | [chembl3d/implicit_h](chembl3d/implicit_h) |
+## Available checkpoints
 
-The `epoch_200.ckpt` name means **200 completed epochs**. The original Lightning
-checkpoint stores `epoch=199`. Files are byte-for-byte copies of the selected
-training checkpoints, including current weights, EMA weights, and optimizer
-state. Sampling uses EMA weights by default.
+| Dataset | Model name | Hydrogen representation | Completed epochs | Files |
+| --- | --- | --- | --- | --- |
+| ChEMBL3D | `chembl3d-explicit-h` | Explicit H atoms | 200, 300, 400 | [Bundle](https://huggingface.co/maltefranke/morph/tree/main/chembl3d/explicit_h) |
+| ChEMBL3D | `chembl3d-implicit-h` | Heavy atoms; implicit H | 100, 200, 300 | [Bundle](https://huggingface.co/maltefranke/morph/tree/main/chembl3d/implicit_h) |
+| GEOM-Drugs | `geom-explicit-h` | Explicit H atoms | 50 | [Bundle](https://huggingface.co/maltefranke/morph/tree/main/geom/explicit_h) |
 
-Each variant contains:
+Epoch numbers count **completed training epochs**: `epoch_050.ckpt` means 50
+completed epochs and stores the zero-based Lightning counter `epoch=49`.
+Checkpoints are original training files containing current weights, EMA weights
+and optimizer state.
 
-- `epoch_NNN.ckpt`: the three selected checkpoints.
-- `atom_tokens.txt`, `edge_tokens.txt`, `charge_tokens.txt`: ordered vocabularies.
-- `distributions.pt`: the original training statistics, including atom counts,
-  type frequencies, coordinate scale, and geometry histograms.
-- `distributions.json`: a readable export of the same statistics.
-- `config.yaml`: portable sampling config with dataset loading disabled.
-- `training_config.yaml`: the original Hydra training configuration, for provenance.
-- `runtime/`: frozen source/configs and dependency specifications for this run.
-- `metadata.json`: epochs, optimizer steps, original filenames, and SHA-256 hashes.
-- `validation.json`: strict checkpoint loading, current/EMA forward checks and
-  a one-molecule end-to-end sampling check. This is a functionality check, not
-  a validity benchmark.
+The [catalog](https://huggingface.co/maltefranke/morph/blob/main/catalog.json)
+is the source of truth for available model names and epochs. Each entry pins an
+immutable release revision so its weights and supporting files stay together.
+New catalog entries are discovered without updating the Python package.
 
-## Download and generate
+## Install and authenticate
 
-Download the variant you need (authenticate with `hf auth login` if required):
+The runtime bundled with GEOM includes the shared inference interface and catalog
+discovery for all models in the table. Install it in a Python environment:
 
 ```bash
-hf download maltefranke/morph --local-dir morph-models \
-  --include 'chembl3d/implicit_h/*' --include sample.py --include README.md --include manifest.json
-cd morph-models
+python -m pip install 'huggingface_hub>=1.0,<2'
+hf auth login
+hf download maltefranke/morph --local-dir morph-runtime \
+  --include 'geom/explicit_h/runtime/*'
+python -m pip install ./morph-runtime/geom/explicit_h/runtime
+morph models list
+```
+
+Authentication requires an account with access to this model repository.
+Python >=3.11 and a compatible PyTorch 2.9 installation are required; CUDA
+generation needs a CUDA-capable PyTorch build. Dependency constraints and tested
+versions are included in the runtime. An existing Morph environment with catalog
+discovery can use the examples directly. On a cluster, run model loading and
+generation inside a compute allocation.
+
+## Load and generate from Python
+
+```python
+from morph.inference import load_pretrained
+
+model = load_pretrained("geom-explicit-h", epoch=50, device="cuda")
+
+trajectories, validity, vocab, coordinate_std = model.generate(
+    n_mols=64,
+    batch_size=16,
+    seed=42,
+    return_traj=False,
+)
+```
+
+Use `device="cpu"` for CPU inference. The same interface loads ChEMBL3D:
+
+```python
+model = load_pretrained("chembl3d-explicit-h", epoch=400, device="cuda")
+# Or choose the heavy-atom model:
+model = load_pretrained("chembl3d-implicit-h", epoch=300, device="cuda")
 ```
 
-Install that variant's runtime in an environment with Python >=3.11 and a
-compatible PyTorch 2.9 installation. A CUDA-capable PyTorch build is needed to
-use `--device cuda`. The tested package versions are recorded in
-`runtime/requirements-tested.txt`; the runtime's `pyproject.toml` specifies the
-dependency constraints.
+Loading downloads only the selected checkpoint and its six required asset files,
+then reuses the Hugging Face cache. Each checkpoint is approximately 510 MB.
+Omitting `epoch` selects the latest published epoch for the chosen model.
+
+**EMA weights are used by default.** The inference loader constructs one network
+and loads only the selected weights; the unselected network and optimizer tensors
+are not materialized from the memory-mapped checkpoint. To select current
+training weights, pass `use_ema=False` to `load_pretrained`.
+
+Generation defaults to 100 integration steps. Each attempt remains in the returned
+lists, including invalid molecules. `trajectories[i][-1]` is its final graph and
+`validity[i]` is its RDKit validity result. Pass `return_traj=True` to retain all
+integration frames. Returned coordinates are already in **Angstrom**; do not
+multiply them by `coordinate_std` again.
+
+## CLI and downloads
+
+Generate GEOM-Drugs molecules and save valid structures to SDF:
 
 ```bash
-python -m pip install ./chembl3d/implicit_h/runtime
-python sample.py --variant implicit_h --epoch 300 \
-  --num-mols 64 --batch-size 16 --device cuda --output implicit_h.sdf
+morph generate --model geom-explicit-h --epoch 50 \
+  --device cuda --batch 64 --batch-size 16 --output geom-e50.sdf
 ```
 
-For explicit hydrogens, download `chembl3d/explicit_h/*` instead, install its
-runtime, and run:
+`--batch` is the total number of generation attempts. Add `--save-trajs` to save
+trajectories and validity flags beside the SDF. `--no-ema` selects current
+weights, and `--steps` and `--seed` control sampling.
+
+Download a portable bundle without loading the model:
 
 ```bash
-python sample.py --variant explicit_h --epoch 400 \
-  --num-mols 64 --batch-size 16 --device cuda --output explicit_h.sdf
+morph models download geom-explicit-h --epoch 50 --local-dir morph-models
+```
+
+Load that local bundle without network access or catalog lookup:
+
+```python
+model = load_pretrained(
+    weights_dir="morph-models/geom/explicit_h",
+    epoch=50,
+    device="cuda",
+)
 ```
 
-The script imports the matching bundled source automatically. Each invocation
-uses one variant in a fresh process. No ChEMBL3D dataset or cluster paths are
-needed. `--device cpu` is supported, but use a compute worker for substantial
-sampling. Compilation is disabled by default; attention uses PyTorch SDPA.
+Use another model name from the table for ChEMBL3D. The interface also accepts
+direct Hub repository IDs with `subfolder` and `revision`. Pass `cache_dir` to
+choose a cache location, or `local_files_only=True` to load a previously cached
+release without network access. The corresponding CLI options are `--cache-dir`
+and `--offline`.
 
-Generated coordinates are in Angstrom. The SDF contains valid molecules; the
-adjacent JSON records every attempted molecule and its validity result, including
-failures. Generation uses 100 integration steps and seed 42 by default.
+The root `sample.py` remains available as the archived ChEMBL3D sampling entry
+point with its corresponding bundled runtimes. The shared interface above
+supports all three published model variants.
 
 ## Hydrogen handling and training data
 
-Both variants use the same split membership: 1,474,008 training, 184,295 validation,
-and 184,242 test molecules. The dataset was filtered to **at most 100 total atoms,
-including hydrogens**, before removing hydrogens for the implicit variant.
+| Dataset | Training filter | Representation |
+| --- | --- | --- |
+| ChEMBL3D | At most 100 total atoms, including H before hydrogen removal | Explicit-H and implicit-H variants |
+| GEOM-Drugs | At most 72 total atoms, including H | Explicit-H variant |
+
+Both ChEMBL3D variants use the same split membership: 1,474,008 training,
+184,295 validation and 184,242 test molecules. The implicit-H model predicts
+heavy atoms, bonds and formal charges; RDKit infers hydrogens during conversion
+and validity checking. Its geometry statistics and coordinate normalization were
+recomputed for the heavy-atom representation.
 
-The explicit variant predicts H atoms as graph vertices. The implicit variant
-predicts heavy atoms, bonds and formal charges; RDKit infers hydrogens during
-conversion and validity checking. Its bundled runtime preserves that behavior.
-Do not mix vocabularies or distributions between variants. Geometry statistics
-and coordinate normalization were recomputed for the heavy-atom representation.
+The GEOM-Drugs epoch 50 checkpoint was obtained by resuming the best checkpoint
+after 48 completed epochs. The resumed run validates every epoch and retains
+checkpoints every fifth epoch. Only the epochs listed in the catalog have been
+published here.
 
-The 100-atom cap describes the training filter, not a hard limit imposed on
-generated molecules. Sample counts follow the saved training prior settings.
+Explicit-H models generate H atoms as graph vertices. Every model uses its own
+ordered vocabularies, distributions and saved configuration; keep those files
+together. The training atom limits describe dataset filtering and do not impose
+a hard atom-count limit during generation.
 
-`manifest.json` lists file sizes and SHA-256 hashes. The original training configs
-contain historical cluster paths and should be treated as provenance; use
-`sample.py` with the portable `config.yaml` for inference.
+## Bundle contents and verification
+
+Each model directory contains:
+
+- `epoch_NNN.ckpt`: the published original training checkpoints.
+- `config.yaml`: portable inference configuration with dataset loading disabled.
+- `atom_tokens.txt`, `edge_tokens.txt`, `charge_tokens.txt`: ordered vocabularies.
+- `distributions.pt`: original training statistics, including atom counts,
+  coordinate scale, type frequencies and geometry histograms.
+- `distributions.json`: a readable export of those statistics.
+- `metadata.json`: dataset, hydrogen mode, epochs, optimizer steps, original
+  filenames and checkpoint SHA-256 hashes.
+- `training_config.yaml`: the original training configuration for provenance.
+- `runtime/`: archived source, configs and dependency specifications.
+- `validation.json`: results of checkpoint loading and sampling checks.
+
+The published bundles were checked for strict current/EMA weight loading and
+end-to-end sampling. These are functionality checks, not validity benchmarks.
+The GEOM release also passed remote file-size and hash verification.
+
+The [ChEMBL3D release manifest](https://huggingface.co/maltefranke/morph/blob/ab9bea92f499f03450ebc99459fce04461504a41/manifest.json)
+describes the original ChEMBL3D release at its pinned revision, including its
+historical model card. GEOM has its own
+[bundle manifest](https://huggingface.co/maltefranke/morph/blob/b839af108eaade757a5d2d668d05498e8c733c1e/geom/explicit_h/manifest.json).
+Original training configs contain historical cluster paths; use `config.yaml`
+through the inference loader for portable generation.

geom/explicit_h/manifest.json

@@ -598,5 +598,9 @@
   "validation.json": {
     "bytes": 294,
     "sha256": "5acbbb78be426e03abd58c3dc062b30b000e91e815c329dfa7f1422c9866b740"
+  },
+  "runtime/README.md": {
+    "bytes": 1041,
+    "sha256": "7e32e04cce313c7faad9a3170ab009591aa327a900ee67f980889b7d7e4c82ba"
   }
 }

geom/explicit_h/runtime/README.md

@@ -0,0 +1,24 @@
+# Morph inference runtime
+
+This source snapshot accompanies the GEOM-Drugs explicit-H release and includes
+the shared inference loader for the ChEMBL3D and GEOM-Drugs checkpoints published
+in `maltefranke/morph`.
+
+Install this directory with `python -m pip install .`. Python >=3.11 and a
+compatible PyTorch 2.9 installation are required. Dependency constraints are in
+`pyproject.toml`; `requirements-tested.txt` records the versions used to validate
+the GEOM release.
+
+```python
+from morph.inference import load_pretrained
+
+model = load_pretrained("geom-explicit-h", epoch=50, device="cuda")
+trajectories, validity, vocab, coordinate_std = model.generate(n_mols=16)
+```
+
+Use `device="cpu"` for CPU inference. Run `morph models list` to discover all
+published model names and epochs. The loader fetches a checkpoint and its
+matching configuration, vocabularies and distributions; no dataset is required.
+
+See the [model card](https://huggingface.co/maltefranke/morph) for installation,
+sampling, hydrogen handling and checkpoint provenance.