Add CheMeleon embedding featurizer and model predict_embedding - #588
Merged
Conversation
Codecov Report❌ Patch coverage is 🚀 New features to boost your workflow:
|
for more information, see https://pre-commit.ci
for more information, see https://pre-commit.ci
smcolby
force-pushed
the
feat/chemeleon-featurizer
branch
from
August 13, 2026 22:46
6a9f345 to
81561ea
Compare
chemeleon-test builds the CheMeleon-compatible MPNN architecture with random weights so tests can exercise the foundation-model code path without a network download. predict_embedding now returns an empty array for an empty input list instead of passing dataset_size=0 through to _safe_inference_batch_size and build_dataloader (which rejects batch_size=0). Also consolidated the module's imports, which had drifted apart with one import block sitting after a function definition.
torch.device() has no "gpu" device type; passing accelerator="gpu" straight through raised at _ensure_model. Add _normalize_accelerator to map the Lightning-style "gpu" value to "cuda" before constructing the device. Also add the missing NumPy-style docstring to __init__ required by ruff's D107 check.
These three tests called .build() with from_foundation="chemeleon", pulling the real checkpoint over the network on every run. Use chemeleon-test instead, which exercises the same foundation-model code path with a random-weight architecture and no network access.
smcolby
force-pushed
the
feat/chemeleon-featurizer
branch
from
August 13, 2026 23:59
81561ea to
cf52b43
Compare
The featurizer module self-registers, but nothing imported it: the registry loader never loaded it, so FeaturizerSpec could not resolve CheMeleonEmbeddingFeaturizer and YAML recipes or the anvil CLI could not name the featurizer. Add the module to the featurizer loader list.
Every registered featurizer declares its registry name as a type ClassVar, but the registry decorator does not set it, so CheMeleonEmbeddingFeaturizer was missing the attribute other code and the feature_blocks block-key convention rely on.
The featurizer hardcoded from_foundation='chemeleon', so every unit test downloaded the multi-GB foundation checkpoint and ran a 2048-width forward on CPU. Route the foundation name through a module constant that tests monkeypatch to the random-weight chemeleon-test architecture the branch already provides, keeping the production name unchanged. The batch-size test, which only repeated the shape assertions of the first test, is replaced with a batch invariance check: embeddings must not depend on the chosen forward batch size, up to float32 reduction-order drift.
predict_embedding raises on the first SMILES that RDKit cannot parse, so the featurizer aborted the whole run when a single input row was invalid, and its unconditional arange indices could never represent a partial success. That breaks the repo featurizer contract, whose index arrays identify the successfully featurized rows and which downstream splits and alignment rely on. featurize now filters unparseable SMILES up front, passes only valid rows down, and returns the valid-row indices; an all-invalid input returns an empty matrix without touching the model.
The zero-row path hardcoded a 2048 width, which only matches the CheMeleon checkpoint: any other foundation width would return an inconsistent shape. The built model already knows its output width in message_hidden_dim, so use it. The not-built guard raised AttributeError with a trained-model message; ValueError with a build() pointer is the honest signal for a construction state.
The method previously read its device from wherever the estimator parameters happened to live, which only worked because the caller pre-places the model, and it ran fingerprint through batch norm in training mode even though it is an inference path that should use running stats. Accept an explicit accelerator argument (None keeps the derive-from-params behavior; gpu is mapped to cuda to match predict), move the estimator to that device with to(), and call eval() before the loop like predict does. Single device only: multi-device routing would require running inference through a Lightning Trainer, which is deferred until there is a real multi-GPU need.
CheMeleonEmbeddingFeaturizer was the only featurizer that checked SMILES parsability up front, and its silent row-skipping contract produced outputs shorter than the input with no error, which is one child's behavior no other featurizer shares. Assume callers supply valid, parsable SMILES and let the toolkit raise on the first unparseable entry, in line with the other SMILES featurizers. Indices are now always a full range, and the empty-input short-circuit is kept so an empty call never triggers the checkpoint download. A shared, fail-loud validation policy is the right scope for this check going forward.
Comments are required to be sentence case per .github/copilot-instructions.md, but the comments added in the preceding two commits opened with lowercase first words. Continuation lines stay lowercase. Pre-existing lowercase openers elsewhere in chemprop.py are left untouched for a later style sweep.
When a from_foundation checkpoint had no usable state_dict, build() silently constructed a randomly initialized estimator and presented it as the loaded foundation. Only the inline chemeleon-test payload was meant to be stateless; a weightless file at an arbitrary path (or from a corrupt download) also slipped through. build() now raises a RuntimeError naming the model when no state_dict is present outside chemeleon-test, whose random-weight behavior is preserved.
The old default inherited whatever device the model parameters happened to occupy, coupling the inference device to incidental prior placement. The default is now auto, resolved with Lightning's own auto accelerator selection so the inference device matches what the repo's Trainer picks on the same machine. Known trainer aliases resolve through a small dict (gpu to cuda, tpu to xla) and every other value passes through verbatim as a torch device name. The CheMeleon featurizer now passes its own accelerator explicitly, so its documented cpu default is independent of this change. The predict_embedding unit tests pin accelerator=cpu; under auto on this hardware the device resolves to MPS, which is not bit-reproducible.
…_name The accelerator to device-name mapping lived inline in predict_embedding, so the alias handling (gpu to cuda, tpu to xla, verbatim passthrough) was only reachable through a built model. Moving it into a small pure helper makes the mapping table directly unit-testable without a GPU or foundation model, and the call site collapses to one line. The new tests cover the table hermetically and verify the auto branch delegates to Lightning's selection by patching it at the lookup site. Lightning's own hardware selection and actual device placement are third-party and hardware behavior, so they stay untested by design.
Comments that mark a new logical chunk need a blank line before them when they follow a code statement at the same indent, so they read as a break rather than a continuation. Exempt where the comment already opens a newly indented block.
ChemPropModel is already imported at module scope; the per-test re-imports were dead weight.
…r logic _normalize_accelerator only handled the gpu alias and left auto to pass through verbatim, which torch.device would reject. _resolve_device_name already covers gpu, tpu, and auto (delegating to Lightning's selection), so drop the featurizer's narrower copy in favor of it.
Sentence-case comment convention was missed in these two spots.
Keeps import placement consistent with the rest of the file.
CheMeleonEmbeddingFeaturizer only checked its device via torch.device() inside _ensure_model, several calls removed from where a bad accelerator was actually supplied. Reproduces the same check in a field_validator so invalid values fail at construction time instead of inside featurize(). Also routes accelerator/batch_size through super().__init__() instead of assigning them directly after construction, since the previous shape bypassed pydantic validation (including the new validator) entirely.
…model predict_embedding already resolves accelerator and moves the model to device on every call, so _ensure_model moving it a second time at build time was dead weight: the device set there was immediately redone by the very next featurize() call.
Matches ChemPropModel.predict_embedding's own default so the featurizer uses an available accelerator automatically instead of forcing CPU unless the caller remembers to override it.
smcolby
commented
Aug 20, 2026
smcolby
commented
Aug 20, 2026
smcolby
commented
Aug 20, 2026
smcolby
commented
Aug 20, 2026
dwwest
approved these changes
Aug 28, 2026
dwwest
left a comment
dwwest
left a comment
Contributor
There was a problem hiding this comment.
Just lots of comments about comments for the most part
| "Using CheMeleon overrides settings for depth, message_hidden_dim, messages, and aggregation" | ||
| ) | ||
| elif self.from_foundation == "chemeleon-test": | ||
| # Build CheMeleon-compatible architecture with random weights, |
Contributor
There was a problem hiding this comment.
Maybe a comment here like ONLY FOR TESTING at the beginning, because this was hard to parse without additional context.
Contributor
Author
There was a problem hiding this comment.
Move the testing-only comment up front, not the second line of the comment.
Contributor
Author
There was a problem hiding this comment.
Also move comments ABOVE if statements, like this:
# Build CheMeleon-compatible architecture with random weights,
# for hermetic tests that need no network access
elif self.from_foundation == "chemeleon-test":
blahAudit all comments against this directive.
Contributor
Author
There was a problem hiding this comment.
Also take a pass at long comments on touched files generally
The 11-key BondMessagePassing argument dict existed in three places: the chemeleon-test branch of build(), the foundation-file fixture in test_chemprop.py, and the 2048 embedding width in the featurizer. Each copy could drift from the CheMeleon checkpoint layout independently. Hoist one _CHEMELEON_MP_HPARAMS constant in chemprop.py and have the other two sites spread-override it, so the test fixture shrinks d_h and depth while inheriting every other key and the featurizer derives its zero-row width from the same dict.
…attr The lazy encoder cache was assigned inside a hand-written __init__ with no class-level declaration, which made it look like incidental instance state rather than a deliberate attribute. Declare it the way MolfeatFeaturizer declares _transformer and document why it exists. The __init__ only forwarded two kwargs that pydantic already defaults, so it is dropped. Construction is now keyword-only, which is what the anvil params path and every existing caller already use.
The featurizer is registered in _registry_loader.py and therefore selectable from an anvil config, but it appeared in neither the featurization API tree nor the anvil featurizer table, so nothing linked to its parameters. Add the automodule page, wire it into the featurization toctree, and add the anvil table row alongside the other featurizers.
The repo standardizes on pytest-mock for patching, and monkeypatch was the odd idiom out in these two new tests. Switch both to mocker so the suite reads consistently.
The method's purpose, letting a built ChemPropModel serve as a featurizer without a trained predictor head, was only discoverable by finding the featurizer that calls it. State it and cross-reference the caller.
The comment described what the dict does without saying where the two spellings come from, leaving a reader to wonder why exactly these keys. Name the source of the mismatch: Lightning's Trainer and torch use different names for the same two devices.
The three-line version buried the point in a semicolon pivot and never said which case it was protecting against. Two lines now name the exception and what an empty state_dict means everywhere else.
Comments describing a branch sat inside it, so the reason for taking a path was only visible after committing to reading the body. Move them above the if, elif, else, and for they describe, keeping inside-the-body comments only where they annotate a specific line. Shorten the comment blocks that had grown to seven and eight lines in build() and validate_optional_params. Each now states the constraint that is not obvious from the code and leaves the rest to the code. Drop the line-by-line narration in freeze_weights, where every statement carried a comment restating it.
The featurizer tests reached the published CheMeleon checkpoint through a production sentinel and an autouse patch fixture that rewrote the featurizer's foundation name. ChemPropModel already accepts an arbitrary checkpoint path via from_foundation, so tests can exercise the real loading path with no patching. A session-scoped conftest fixture writes a checkpoint with the CheMeleon hyperparameters and random weights, shared by both test modules. The featurizer tests take it through a make_featurizer factory that seeds the lazy encoder; test_featurize_empty_input_returns_empty deliberately skips the factory so it can still assert the empty path never builds a model.
build() carried a "chemeleon-test" branch that fabricated a weightless CheMeleon-shaped checkpoint, and the state_dict guard right below it had to exempt that same sentinel to avoid rejecting its empty state_dict. The featurizer mirrored the arrangement with a _FOUNDATION_NAME module constant that existed only for tests to patch. Tests now load a real checkpoint through from_foundation, so both mechanisms are dead. Dropping them makes the weightless-foundation guard unconditional, which is what it was always meant to be, and the featurizer names its foundation inline.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
3 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Add CheMeleon MPNN embedding extraction via a dedicated model method and a reusable featurizer.
ChemPropModel.predict_embedding(smiles_list, batch_size=256) -> np.ndarraythat returns pooled pre-predictor embeddings from the CheMeleon foundation encoder without training. The method builds aMoleculeDataset, applies safe inference batch sizing to avoid chemprop’slen % batch == 1drop, runsestimator.fingerprintundertorch.inference_mode(), and returnsnp.float32embeddings.CheMeleonEmbeddingFeaturizerinopenadmet/models/features/chemeleon_embedding.py, registered as"CheMeleonEmbeddingFeaturizer". Exposesacceleratorandbatch_sizeparameters. Lazily builds aChemPropModel(from_foundation="chemeleon"), moves the estimator to the requested device, and delegates topredict_embedding. Output matches existing featurizer contracts:(embeddings, indices).predict_embedding: unbuilt error, shape/dtype(N,2048)float32, safe batch size no drop, device handling, determinism.Quality Assurance & AI Policy
To maintain project quality and respect maintainer bandwidth, please confirm the following:
Status
Developers Certificate of Origin
Changed files
openadmet/models/architecture/chemprop.pyopenadmet/models/features/chemeleon_embedding.pyopenadmet/models/tests/unit/models/test_chemprop.pyopenadmet/models/tests/unit/features/test_chemeleon_embedding.py