Refactor spectrum_vector code, fix imports, and upgrade dependencies - #29
Open
rukubrakov wants to merge 17 commits into
Open
Refactor spectrum_vector code, fix imports, and upgrade dependencies#29rukubrakov wants to merge 17 commits into
rukubrakov wants to merge 17 commits into
Conversation
Eliminated mz_array, intensity_array, spectrum_vector, max_peak fields and all methods that depended on them; the transformer approach fully superseded this code.
Added missing @staticmethod decorators, fixed bare except, SIM118/SIM401/C401 simplifications, and renamed unused loop variables.
…le dep. Renamed precursor_hook → global_token_hook to match the renamed API in depthcharge 0.4.x.
Replace float adduct-mass vectors with categorical nn.Embedding (metabo-depthcharge vocab) across training and inference pipelines; fix HiGHS Not Solved errors crashing preprocessing by setting catch_errors=True.
Replaces manual nn.Embedding + FloatEncoder with MetadataEncoder, which properly masks CE=0 as missing via a Linear+mask rather than producing a non-zero sinusoidal encoding.
…ding code. CE is now normalized (/100) and handles missing/stepped values via metabo-depthcharge's parser. Removes encode_adduct_mass and OneHotEncoding which were no longer used.
Tries collision_energy, collision_energy_1, and ce key variants to cover MSG, Spectraverse, and legacy formats.
…n dep to GitHub. Replaces CLS-token extraction with AttnAggregator, removes FloatEncoder import, delegates to super().forward() via global_token_hook override.
Contributor
📝 WalkthroughWalkthroughThis PR integrates the external ChangesMetadata Encoding & Data Pipeline Refactor
SLURM Script Enhancements
Estimated code review effort🎯 4 (Complex) | ⏱️ ~60 minutes Possibly related PRs
Suggested reviewers
Poem
🚥 Pre-merge checks | ✅ 4 | ❌ 1❌ Failed checks (1 warning)
✅ Passed checks (4 passed)
✏️ Tip: You can configure your own custom pre-merge checks in the settings. ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
There was a problem hiding this comment.
Pull request overview
This PR modernizes SIMBA’s spectrum embedding stack by migrating away from legacy spectrum-vector code and depthcharge-specific encoder assumptions, adds categorical metadata encoding (adduct, collision energy, ion mode), and updates dependencies/scripts to support the new pipeline while improving preprocessing robustness.
Changes:
- Rewrites the spectrum encoder to subclass
metabo-depthcharge’sSpectrumEncoderand updates similarity models to consume pooled(B, d_model)embeddings directly. - Switches adduct handling from large one-hot vectors to categorical indices (
encode_adduct) and normalizes collision energy viaencode_collision_energy, including broader MGF CE key support. - Removes dead spectrum-vector preprocessing/utilities and hardens ILP edit-distance preprocessing by mapping solver failures to
np.nan.
Reviewed changes
Copilot reviewed 16 out of 16 changed files in this pull request and generated 6 comments.
Show a summary per file
| File | Description |
|---|---|
| tools/slurm/train_joint_metadata.slurm.sh | New SLURM training entrypoint enabling adduct/CE/ion_mode features. |
| tools/slurm/inference_joint_metadata.slurm.sh | New SLURM inference script for the attention-pooled metadata model. |
| tools/slurm/inference_joint_metadata_cls.slurm.sh | New SLURM inference script for older CLS-pool checkpoints. |
| simba/core/training/train_utils.py | Removes dead spectrum-vector helper method. |
| simba/core/models/spectrum_encoder.py | Replaces depthcharge-based encoder with metabo-depthcharge encoder + metadata injection via hook. |
| simba/core/models/similarity_models.py | Updates model forward paths for categorical adducts and pooled encoder output. |
| simba/core/data/weighted_sampling.py | Cleans up duplicate import ordering. |
| simba/core/data/spectrum.py | Removes legacy spectrum_vector/max_peak fields from SpectrumExt serialization/state. |
| simba/core/data/preprocessor.py | Removes unused spectrum-vector preprocessing logic; refines imports/helpers. |
| simba/core/data/loaders/loaders.py | Improves MGF parsing (CE fallback keys, cleanup) and minor iterator/robustness tweaks. |
| simba/core/data/encoding.py | Removes adduct encoding utilities / pandas dependency; keeps ion activation/method helpers. |
| simba/core/data/datasets/multitask_dataset.py | Switches adduct storage from one-hot float arrays to int64 categorical indices. |
| simba/core/data/datasets/multitask_dataset_builder.py | Uses encode_adduct + encode_collision_energy and adjusts metadata dtypes. |
| simba/core/data/datasets/encoder_dataset_builder.py | Uses categorical adduct indices + normalized CE encoding for encoder dataset creation. |
| simba/core/chemistry/edit_distance/edit_distance.py | Treats ILP solver failures as np.nan instead of crashing preprocessing. |
| pyproject.toml | Upgrades depthcharge dependency spec and adds metabo-depthcharge sourcing via uv. |
Comments suppressed due to low confidence (1)
simba/core/data/spectrum.py:55
- This refactor removes the
spectrum_vector/max_peakfields and their setters. The repository’s unit tests (e.g.tests/unit/test_spectrum_ext.py) still callset_spectrum_vector()andset_max_peak(), so the test suite will fail unless those tests are updated or minimal backward-compatible stubs are kept here.
If the intent is to fully remove these fields, please update the tests accordingly; otherwise, consider retaining them as legacy attributes/setters without using them elsewhere.
# extra variables
self.params = params
self.retention_time = retention_time
self.smiles = smiles
self.library = library
self.inchi = inchi
self.ionmode = ionmode
self.adduct = adduct
self.ce = ce
self.ion_activation = ion_activation
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
| im = params["ionization_method"] if "ionization_method" in params else None | ||
| # Try common MGF key variants: COLLISION_ENERGY (MSG/MassSpecGym), | ||
| # COLLISION_ENERGY_1 (Spectraverse stepped CE), CE (legacy) | ||
| ce = params.get("collision_energy") or params.get("collision_energy_1") or params.get("ce") |
Comment on lines
+95
to
+110
| self._extra_kwargs = kwargs | ||
|
|
||
| precursor_charge = ( | ||
| kwargs["precursor_charge"].float().to(device).view(batch_size) | ||
| ) | ||
| # skip the use of the precursor charge field | ||
| if self.use_ion_mode: | ||
| placeholder[:, 1] = precursor_charge | ||
| metadata = {} | ||
| if self.use_adduct and "adduct" in kwargs: | ||
| metadata["adduct"] = kwargs["adduct"].long().to(device).view(batch_size) | ||
| if self.use_ce and "ce" in kwargs: | ||
| metadata["collision_energy"] = kwargs["ce"].float().to(device).view(batch_size) | ||
|
|
||
| current_idx = 2 # keep track of where to insert metadata | ||
| precursor_mz = kwargs["precursor_mass"].float().to(device).view(batch_size) | ||
|
|
||
| ionmode = kwargs["ionmode"].float().to(device).view(batch_size) | ||
| if self.use_ion_mode: | ||
| placeholder[:, current_idx] = ionmode | ||
| current_idx += 1 | ||
| return super().forward( | ||
| mz=mz_array, | ||
| intensity=intensity_array, | ||
| precursor_mz=precursor_mz, | ||
| metadata=metadata if metadata else None, | ||
| ) |
Comment on lines
134
to
139
| def set_murcko_scaffold(self, murcko_scaffold): | ||
| self.murcko_scaffold = murcko_scaffold | ||
|
|
||
| def set_smiles(self, smiles): | ||
| self.smiles = smiles | ||
|
|
Comment on lines
472
to
476
| kwargs_0["ionmode"] = batch["ionmode_0"].float() | ||
| kwargs_1["ionmode"] = batch["ionmode_1"].float() | ||
| batch["adduct_0"] = torch.nan_to_num( | ||
| batch["adduct_0"], nan=0.0, posinf=0.0, neginf=0.0 | ||
| ) | ||
| batch["adduct_1"] = torch.nan_to_num( | ||
| batch["adduct_1"], nan=0.0, posinf=0.0, neginf=0.0 | ||
| ) | ||
|
|
||
| kwargs_0["adduct"] = batch["adduct_0"].float() | ||
| kwargs_1["adduct"] = batch["adduct_1"].float() | ||
| kwargs_0["adduct"] = batch["adduct_0"].long() | ||
| kwargs_1["adduct"] = batch["adduct_1"].long() | ||
|
|
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Actionable comments posted: 4
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (6)
simba/core/chemistry/edit_distance/edit_distance.py (1)
774-816:⚠️ Potential issue | 🟠 Major | ⚡ Quick winRun Ruff formatter on this file before merge.
CI is currently blocked because
ruff format --checkreports thatsimba/core/chemistry/edit_distance/edit_distance.pywould be reformatted. Please apply the formatter so this change passes the existing quality gate.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/chemistry/edit_distance/edit_distance.py` around lines 774 - 816, Run the Ruff formatter on the edited file to satisfy the linter: format simba/core/chemistry/edit_distance/edit_distance.py (which contains the simba_solve_pair_mces function and related constants like VERY_HIGH_DISTANCE) using the ruff formatter (e.g., ruff format or your repo's pre-configured command), update the commit with the formatted changes, and push so CI's ruff format --check passes.Source: Pipeline failures
simba/core/data/loaders/loaders.py (1)
442-446:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDon't coerce missing SMILES to
""here.In the Janssen/default validation paths, spectra with no SMILES already pass validation, so this change turns a bad record into
SpectrumExt.smiles=""instead of rejecting it. Downstream training code canonicalizesSpectrumExt.smilesdirectly; empty strings will collapse unrelated spectra into the same "molecule" and contaminate grouping/fingerprint steps. Please fail closed here and tighten the upstream validation instead of materializing a missing structure as an empty string.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/data/loaders/loaders.py` around lines 442 - 446, The code currently coerces missing SMILES to an empty string which causes absent structures to be treated as valid molecules; change the handling in the loader so you do not default to "" — use smiles = params.get("smiles") (no default) and if smiles is None then fail the record (e.g., return None) so missing SMILES are not materialized; locate the assignment to smiles and the nearby precursor_mz check in LoadData.get_precursor_mz / the loader function and ensure SpectrumExt.smiles is never set to "" (tighten upstream validation instead of converting missing values to empty strings).simba/core/data/preprocessor.py (1)
132-158:⚠️ Potential issue | 🟠 MajorHarden the Classyfire HTTP call and fix
min_intensitybeing ignored in preprocessing.
_get_classbuilds the Classyfire URL by interpolating rawmol_valinto the query string (...classyfire?{mol_type}={mol_val}) without URL-encoding, and it has notimeout/requests.RequestExceptionhandling (onlyJSONDecodeErroris caught), so preprocessing can hang or silently fail.preprocess_all_spectraoverwrites themin_intensityparameter inside the loop (min_intensity = 0.00/0.01based ontraining), so callers cannot control the intensity threshold.Suggested fix
- r = requests.get( - f"https://gnps-structure.ucsd.edu/classyfire?{mol_type}={mol_val}" - ) - if r.status_code != 200: + try: + r = requests.get( + "https://gnps-structure.ucsd.edu/classyfire", + params={mol_type: mol_val}, + timeout=10, + ) + r.raise_for_status() + except requests.RequestException: return None🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/data/preprocessor.py` around lines 132 - 158, Harden _get_class by URL-encoding mol_val when building the query string (use urllib.parse.quote or equivalent) and add a requests timeout and broad exception handling: catch requests.exceptions.RequestException around requests.get, return None on error, and keep the existing JSONDecodeError handling; ensure you still check response.status_code before parsing. In preprocess_all_spectra, stop overwriting the incoming min_intensity inside the loop—use the parameter value (or set a single default before the loop only if min_intensity is None) so callers can control the threshold; update any local variables like min_intensity = 0.00/0.01 to only apply when no parameter was provided.simba/core/models/similarity_models.py (2)
466-499:⚠️ Potential issue | 🟠 Major | ⚡ Quick winGuard the optional metadata reads in
SimilarityModelMultitask.forward.
multitask_dataset.pyonly materializesionmode_*/adduct_*when adduct metadata is enabled andce_*when CE is enabled, but this block now reads all metadata tensors unconditionally. Any config with one ofuse_adduct,use_ce,use_ion_activation,use_ion_method, oruse_ion_modeturned off will fail withKeyErrorbefore the encoder is called.Suggested fix
- batch["ionmode_0"] = torch.nan_to_num( - batch["ionmode_0"], nan=0.0, posinf=0.0, neginf=0.0 - ) - batch["ionmode_1"] = torch.nan_to_num( - batch["ionmode_1"], nan=0.0, posinf=0.0, neginf=0.0 - ) - kwargs_0["ionmode"] = batch["ionmode_0"].float() - kwargs_1["ionmode"] = batch["ionmode_1"].float() - kwargs_0["adduct"] = batch["adduct_0"].long() - kwargs_1["adduct"] = batch["adduct_1"].long() - - batch["ce_0"] = torch.nan_to_num(batch["ce_0"], nan=0.0, posinf=0.0, neginf=0.0) - batch["ce_1"] = torch.nan_to_num(batch["ce_1"], nan=0.0, posinf=0.0, neginf=0.0) - kwargs_0["ce"] = batch["ce_0"].float() - kwargs_1["ce"] = batch["ce_1"].float() - - batch["ion_activation_0"] = torch.nan_to_num( - batch["ion_activation_0"], nan=0.0, posinf=0.0, neginf=0.0 - ) - batch["ion_activation_1"] = torch.nan_to_num( - batch["ion_activation_1"], nan=0.0, posinf=0.0, neginf=0.0 - ) - kwargs_0["ion_activation"] = batch["ion_activation_0"].float() - kwargs_1["ion_activation"] = batch["ion_activation_1"].float() - - batch["ion_method_0"] = torch.nan_to_num( - batch["ion_method_0"], nan=0.0, posinf=0.0, neginf=0.0 - ) - batch["ion_method_1"] = torch.nan_to_num( - batch["ion_method_1"], nan=0.0, posinf=0.0, neginf=0.0 - ) - - kwargs_0["ion_method"] = batch["ion_method_0"].float() - kwargs_1["ion_method"] = batch["ion_method_1"].float() + if self.use_ion_mode: + batch["ionmode_0"] = torch.nan_to_num( + batch["ionmode_0"], nan=0.0, posinf=0.0, neginf=0.0 + ) + batch["ionmode_1"] = torch.nan_to_num( + batch["ionmode_1"], nan=0.0, posinf=0.0, neginf=0.0 + ) + kwargs_0["ionmode"] = batch["ionmode_0"].float() + kwargs_1["ionmode"] = batch["ionmode_1"].float() + + if self.use_adduct: + kwargs_0["adduct"] = batch["adduct_0"].long() + kwargs_1["adduct"] = batch["adduct_1"].long() + + if self.use_ce: + batch["ce_0"] = torch.nan_to_num( + batch["ce_0"], nan=0.0, posinf=0.0, neginf=0.0 + ) + batch["ce_1"] = torch.nan_to_num( + batch["ce_1"], nan=0.0, posinf=0.0, neginf=0.0 + ) + kwargs_0["ce"] = batch["ce_0"].float() + kwargs_1["ce"] = batch["ce_1"].float() + + if self.use_ion_activation: + batch["ion_activation_0"] = torch.nan_to_num( + batch["ion_activation_0"], nan=0.0, posinf=0.0, neginf=0.0 + ) + batch["ion_activation_1"] = torch.nan_to_num( + batch["ion_activation_1"], nan=0.0, posinf=0.0, neginf=0.0 + ) + kwargs_0["ion_activation"] = batch["ion_activation_0"].float() + kwargs_1["ion_activation"] = batch["ion_activation_1"].float() + + if self.use_ion_method: + batch["ion_method_0"] = torch.nan_to_num( + batch["ion_method_0"], nan=0.0, posinf=0.0, neginf=0.0 + ) + batch["ion_method_1"] = torch.nan_to_num( + batch["ion_method_1"], nan=0.0, posinf=0.0, neginf=0.0 + ) + kwargs_0["ion_method"] = batch["ion_method_0"].float() + kwargs_1["ion_method"] = batch["ion_method_1"].float()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/similarity_models.py` around lines 466 - 499, The forward method in SimilarityModelMultitask is reading optional metadata keys (ionmode_*, adduct_*, ce_*, ion_activation_*, ion_method_*) unconditionally which causes KeyError when those features are disabled; update SimilarityModelMultitask.forward to guard each read by checking either the model config flags (e.g., self.use_adduct, self.use_ce, self.use_ion_activation, self.use_ion_method, self.use_ion_mode) or the presence of the key in the incoming batch dict before calling torch.nan_to_num and assigning to kwargs_0/kwargs_1, and only add those entries to kwargs when present; ensure any default handling (e.g., skip adding the key or provide a sensible tensor default) is consistent with how multitask_dataset.py materializes those tensors so the encoder receives only the enabled metadata.
774-803:⚠️ Potential issue | 🟠 MajorCheckpoint reloads can drop metadata-enabled encoder weights
EmbeddingExtractor.load_twin_network()andSimba.load_model()both callSimilarityModel(Multitask).load_from_checkpoint(..., strict=False|strict)without passinguse_adduct,use_ce, or anyuse_ion_*flags. SinceSimilarityModeldefaults these flags toFalseand builds different encoder submodules accordingly—and the repo contains nosave_hyperparameters()/hparamsusage to restore the trained flag values from the checkpoint—strict=Falsecan silently ignore/detach the metadata-specific weights.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/similarity_models.py` around lines 774 - 803, The checkpoint reloads omit encoder metadata flags causing metadata-specific weights to be skipped; fix EmbeddingExtractor.load_twin_network and Simba.load_model so calls to SimilarityModel.load_from_checkpoint and SimilarityModelMultitask.load_from_checkpoint explicitly pass the encoder flags (e.g. use_adduct, use_ce, use_ion_pos, use_ion_neg or whatever use_ion_* names are defined) by reading them from the current config (self.config.model.encoder.*) or, if present, from the checkpoint hyperparameters (checkpoint["hyper_parameters"] / checkpoint.get("hparams")), and forward those boolean flags into the load_from_checkpoint call (keep strict as-is) so the loaded model builds the same encoder submodules and does not drop metadata-enabled weights.simba/core/models/spectrum_encoder.py (1)
1-137:⚠️ Potential issue | 🟡 Minor | ⚡ Quick winRun
ruff formaton this file before merge.The Code Quality pipeline is already red because
ruff format --checkwants to reformat this file.🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/spectrum_encoder.py` around lines 1 - 137, This file fails ruff formatting; run the formatter and commit the changes. Fix by running `ruff format simba/core/models/spectrum_encoder.py` (or `ruff format .` for the repo), update the diff, and re-run `ruff format --check`; ensure the cleaned-up file containing SpectrumTransformerEncoderCustom, its __init__, forward, and global_token_hook methods is committed so CI passes. Optionally add/verify ruff in pre-commit to prevent future formatting failures.Source: Pipeline failures
🧹 Nitpick comments (1)
simba/core/data/preprocessor.py (1)
173-187: ⚡ Quick winDon't overwrite the caller's
min_intensityinside the loop.
preprocess_all_spectra(..., min_intensity=...)currently ignores the supplied value and hard-codes0.0/0.01on every iteration. If the training/eval default needs to vary, compute aneffective_min_intensityonce before the loop or make the parameterNoneby default instead of exposing an argument that never takes effect.Suggested fix
- random.seed(random_seed) - for i, spectrum in tqdm(enumerate(spectrums)): - if training: - min_intensity = 0.00 - else: - min_intensity = 0.01 + random.seed(random_seed) + effective_min_intensity = 0.00 if training else min_intensity + for i, spectrum in tqdm(enumerate(spectrums), total=len(spectrums)): @@ - min_intensity=min_intensity, + min_intensity=effective_min_intensity,🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/data/preprocessor.py` around lines 173 - 187, The loop in preprocess_all_spectra overwrites the caller's min_intensity by setting min_intensity=0.00/0.01 for each spectrum; instead compute an effective_min_intensity once before the loop (e.g., if min_intensity is None then set based on training flag) and pass that to preprocess_spectrum, or make the function default min_intensity=None and only substitute a computed default once; update references in preprocess_all_spectra and calls to preprocess_spectrum so the user-supplied min_intensity is respected.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@pyproject.toml`:
- Line 72: The dependency "metabo-depthcharge" must be moved out of
tool.uv.sources and added to project.dependencies as a PEP 508 direct URL pinned
to an immutable tag or commit (e.g., "metabo-depthcharge @
git+https://...@<tag-or-commit>") so the package metadata carries the origin and
installs are reproducible; update project.dependencies to include that pinned
direct URL, remove it from tool.uv.sources, and (optionally) ensure uv.lock is
generated/checked in if you intend to keep uv workflows.
In `@simba/core/data/datasets/multitask_dataset.py`:
- Around line 108-112: The code incorrectly allocates dictionary["ionmode_0"]
and dictionary["ionmode_1"] under the use_adduct branch; change the allocation
so ionmode arrays are created when self.use_ion_mode is true and adduct arrays
when self.use_adduct is true. Specifically, in the method that builds the
initial dictionary (where dictionary["ionmode_0"/"ionmode_1"] and
dictionary["adduct_0"/"adduct_1"] are created), replace the current if
self.use_adduct block with two checks: if self.use_ion_mode -> allocate ionmode
arrays (np.zeros(..., dtype=np.float32)), and if self.use_adduct -> allocate
adduct arrays (np.zeros(..., dtype=np.int64)); this aligns with
get_original_dictionary which expects ionmode keys when use_ion_mode is enabled.
In `@tools/slurm/inference_joint_metadata_cls.slurm.sh`:
- Around line 14-16: The script currently runs inference with pool="attention"
which will randomly initialize AttnAggregator weights because the loaded
checkpoint was trained with pool=None; update the script to detect and enforce
matching pooling config before running inference: read the checkpoint's saved
pooling setting (e.g., a "pool" or encoder/pooling field in the checkpoint
metadata) and compare it to the runtime pool variable (pool="attention"), and if
they differ either (a) override the runtime pool to match the checkpoint (set
pool=None) or (b) fail fast with a clear error and exit; apply this check early
in the script (before the inference/aggregation steps and also cover the related
code paths around lines 39-54) so mismatched AttnAggregator/encoder configs are
never used for production-style inference.
In `@tools/slurm/inference_joint_metadata.slurm.sh`:
- Around line 19-33: Update the script startup to fail fast by changing the
shell options from "set -uo pipefail" to "set -euo pipefail" and make the
working-directory change robust by replacing the unguarded "cd
/home/nkubrakov/simba-integration" with "cd /home/nkubrakov/simba-integration ||
exit 1"; apply the same edits in the three affected scripts
(inference_joint_metadata.slurm.sh, inference_joint_metadata_cls.slurm.sh,
train_joint_metadata.slurm.sh) so the job stops immediately on errors or if the
cd fails.
---
Outside diff comments:
In `@simba/core/chemistry/edit_distance/edit_distance.py`:
- Around line 774-816: Run the Ruff formatter on the edited file to satisfy the
linter: format simba/core/chemistry/edit_distance/edit_distance.py (which
contains the simba_solve_pair_mces function and related constants like
VERY_HIGH_DISTANCE) using the ruff formatter (e.g., ruff format or your repo's
pre-configured command), update the commit with the formatted changes, and push
so CI's ruff format --check passes.
In `@simba/core/data/loaders/loaders.py`:
- Around line 442-446: The code currently coerces missing SMILES to an empty
string which causes absent structures to be treated as valid molecules; change
the handling in the loader so you do not default to "" — use smiles =
params.get("smiles") (no default) and if smiles is None then fail the record
(e.g., return None) so missing SMILES are not materialized; locate the
assignment to smiles and the nearby precursor_mz check in
LoadData.get_precursor_mz / the loader function and ensure SpectrumExt.smiles is
never set to "" (tighten upstream validation instead of converting missing
values to empty strings).
In `@simba/core/data/preprocessor.py`:
- Around line 132-158: Harden _get_class by URL-encoding mol_val when building
the query string (use urllib.parse.quote or equivalent) and add a requests
timeout and broad exception handling: catch requests.exceptions.RequestException
around requests.get, return None on error, and keep the existing JSONDecodeError
handling; ensure you still check response.status_code before parsing. In
preprocess_all_spectra, stop overwriting the incoming min_intensity inside the
loop—use the parameter value (or set a single default before the loop only if
min_intensity is None) so callers can control the threshold; update any local
variables like min_intensity = 0.00/0.01 to only apply when no parameter was
provided.
In `@simba/core/models/similarity_models.py`:
- Around line 466-499: The forward method in SimilarityModelMultitask is reading
optional metadata keys (ionmode_*, adduct_*, ce_*, ion_activation_*,
ion_method_*) unconditionally which causes KeyError when those features are
disabled; update SimilarityModelMultitask.forward to guard each read by checking
either the model config flags (e.g., self.use_adduct, self.use_ce,
self.use_ion_activation, self.use_ion_method, self.use_ion_mode) or the presence
of the key in the incoming batch dict before calling torch.nan_to_num and
assigning to kwargs_0/kwargs_1, and only add those entries to kwargs when
present; ensure any default handling (e.g., skip adding the key or provide a
sensible tensor default) is consistent with how multitask_dataset.py
materializes those tensors so the encoder receives only the enabled metadata.
- Around line 774-803: The checkpoint reloads omit encoder metadata flags
causing metadata-specific weights to be skipped; fix
EmbeddingExtractor.load_twin_network and Simba.load_model so calls to
SimilarityModel.load_from_checkpoint and
SimilarityModelMultitask.load_from_checkpoint explicitly pass the encoder flags
(e.g. use_adduct, use_ce, use_ion_pos, use_ion_neg or whatever use_ion_* names
are defined) by reading them from the current config
(self.config.model.encoder.*) or, if present, from the checkpoint
hyperparameters (checkpoint["hyper_parameters"] / checkpoint.get("hparams")),
and forward those boolean flags into the load_from_checkpoint call (keep strict
as-is) so the loaded model builds the same encoder submodules and does not drop
metadata-enabled weights.
In `@simba/core/models/spectrum_encoder.py`:
- Around line 1-137: This file fails ruff formatting; run the formatter and
commit the changes. Fix by running `ruff format
simba/core/models/spectrum_encoder.py` (or `ruff format .` for the repo), update
the diff, and re-run `ruff format --check`; ensure the cleaned-up file
containing SpectrumTransformerEncoderCustom, its __init__, forward, and
global_token_hook methods is committed so CI passes. Optionally add/verify ruff
in pre-commit to prevent future formatting failures.
---
Nitpick comments:
In `@simba/core/data/preprocessor.py`:
- Around line 173-187: The loop in preprocess_all_spectra overwrites the
caller's min_intensity by setting min_intensity=0.00/0.01 for each spectrum;
instead compute an effective_min_intensity once before the loop (e.g., if
min_intensity is None then set based on training flag) and pass that to
preprocess_spectrum, or make the function default min_intensity=None and only
substitute a computed default once; update references in preprocess_all_spectra
and calls to preprocess_spectrum so the user-supplied min_intensity is
respected.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: 84cab01c-4740-4f8a-8d8f-fc32e4c2c2a4
📒 Files selected for processing (16)
pyproject.tomlsimba/core/chemistry/edit_distance/edit_distance.pysimba/core/data/datasets/encoder_dataset_builder.pysimba/core/data/datasets/multitask_dataset.pysimba/core/data/datasets/multitask_dataset_builder.pysimba/core/data/encoding.pysimba/core/data/loaders/loaders.pysimba/core/data/preprocessor.pysimba/core/data/spectrum.pysimba/core/data/weighted_sampling.pysimba/core/models/similarity_models.pysimba/core/models/spectrum_encoder.pysimba/core/training/train_utils.pytools/slurm/inference_joint_metadata.slurm.shtools/slurm/inference_joint_metadata_cls.slurm.shtools/slurm/train_joint_metadata.slurm.sh
💤 Files with no reviewable changes (2)
- simba/core/training/train_utils.py
- simba/core/data/encoding.py
gdewael
reviewed
Jun 15, 2026
|
|
||
| # Ion Activation | ||
| if hasattr(spectrum, 'ion_activation') and spectrum.ion_activation is not None and spectrum.ion_activation != "None": | ||
| ia[i] = encode_ion_activation(spectrum.ion_activation) |
There was a problem hiding this comment.
We could also consider adding these to metabo-depthcharge to provide a unified interface to encoding metadata. Similar for encode_ionization_method at L81
gdewael
reviewed
Jun 15, 2026
| im = params["ionization_method"] if "ionization_method" in params else None | ||
| # Try common MGF key variants: COLLISION_ENERGY (MSG/MassSpecGym), | ||
| # COLLISION_ENERGY_1 (Spectraverse stepped CE), CE (legacy) | ||
| ce = params.get("collision_energy") or params.get("collision_energy_1") or params.get("ce") |
There was a problem hiding this comment.
Could also use https://metabo-depthcharge.readthedocs.io/en/latest/api/generated/metabo_depthcharge.spec.preprocessing.CollapseSteppedCE.html But this somewhat requires to use the Spectrum class object defined by metabo-depthcharge.
Another note here: In massspecgym, normalized collision energies are stored under the "COLLISION_ENERGY" MGF field, whereas in SpectraVerse both NORMALIZED_COLLISION_ENERGY_{1,2,3} and COLLISION_ENERGY_{1,2,3} are provided. For parity, normalized collision energies should be used. (the above linked class allows this seamlessly).
gdewael
reviewed
Jun 15, 2026
| ) | ||
|
|
||
| super().__init__( | ||
| *args, pool="attention", metadata_encoder=metadata_enc, **kwargs |
There was a problem hiding this comment.
Note that this may have unforeseen consequences on how easy it is for the model to make its predictions conditional on metadata: pool="cls" will take the embedding of the global_token at output, where all metadata was injected.
It is not crazy to think that it easier for the model to use metadata information when explicitly selecting this token as the "summary" embedding, instead of using attention pooling (even after all the transformer layers).
Fine for now, but mental note in future experiments on metadata inclusion, relevant to discuss with Janne and Sebastian, as they have been heavily experimenting with this.
gdewael
reviewed
Jun 15, 2026
gdewael
left a comment
gdewael
left a comment
There was a problem hiding this comment.
I've left some comments/thoughts here and there. I dont have other major notes on how the code should be restructured to utilize metabo-depthcharge.
…N_ENERGY fields over absolute COLLISION_ENERGY
…, add set -euo pipefail and cd guards to SLURM scripts
…tion Switch ia/im from one-hot float to int64 index throughout; add configurable pool (attention/cls) wired through config, training, inference, and EmbeddingExtractor.
Contributor
There was a problem hiding this comment.
Actionable comments posted: 2
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (3)
simba/core/models/similarity_models.py (3)
828-831:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDtype mismatch:
ion_activationandion_methodshould be.long(), not.float().Same issue as in
SimilarityModel.forward— the spectrum encoder expects these aslongtensors for embedding lookup.Proposed fix
if "ion_activation" in batch: - kwargs["ion_activation"] = batch["ion_activation"].float() + kwargs["ion_activation"] = batch["ion_activation"].long() if "ion_method" in batch: - kwargs["ion_method"] = batch["ion_method"].float() + kwargs["ion_method"] = batch["ion_method"].long()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/similarity_models.py` around lines 828 - 831, The ion_activation and ion_method tensors are being converted to float type when they should be converted to long type, since the spectrum encoder expects these as long tensors for embedding lookup operations. Change the `.float()` method calls to `.long()` for both the ion_activation and ion_method batch dictionary entries to match the expected input types for the spectrum encoder.
133-139:⚠️ Potential issue | 🟠 Major | ⚡ Quick winDtype mismatch:
ion_activationandion_methodshould be.long(), not.float().The spectrum encoder expects
ion_activationandion_methodaslongtensors (seespectrum_encoder.pylines 109-116 where.long()is called). However, here they are cast to.float(), which is inconsistent withSimilarityModelMultitask.forward(lines 493, 497) that correctly uses.long().Proposed fix
if self.use_ion_activation: - kwargs_0["ion_activation"] = batch["ion_activation_0"].float() - kwargs_1["ion_activation"] = batch["ion_activation_1"].float() + kwargs_0["ion_activation"] = batch["ion_activation_0"].long() + kwargs_1["ion_activation"] = batch["ion_activation_1"].long() if self.use_ion_method: - kwargs_0["ion_method"] = batch["ion_method_0"].float() - kwargs_1["ion_method"] = batch["ion_method_1"].float() + kwargs_0["ion_method"] = batch["ion_method_0"].long() + kwargs_1["ion_method"] = batch["ion_method_1"].long()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/similarity_models.py` around lines 133 - 139, In the code block where use_ion_activation and use_ion_method conditions are checked, change the dtype casting from .float() to .long() for both ion_activation and ion_method. Specifically, in the if self.use_ion_activation block, change batch["ion_activation_0"].float() and batch["ion_activation_1"].float() to use .long() instead. Similarly, in the if self.use_ion_method block, change batch["ion_method_0"].float() and batch["ion_method_1"].float() to use .long() instead. This ensures consistency with the spectrum encoder's expected input types and matches the correct implementation in SimilarityModelMultitask.forward.
344-383:⚠️ Potential issue | 🟠 Major | ⚡ Quick winMissing
poolparameter inSimilarityModelMultitask.__init__.
SimilarityModel.__init__now accepts apoolparameter (line 51), butSimilarityModelMultitaskdoesn't include it in its signature or pass it tosuper().__init__(). This means thepoolconfiguration passed viaload_from_checkpoint(ininference.pyline 243 andtraining.pyline 467) won't actually reach the encoder.Proposed fix
def __init__( self, d_model, n_layers, n_classes, use_gumbel, dropout=0.1, weights=None, lr=None, use_element_wise=True, use_cosine_distance=True, weights_sim2=None, use_edit_distance_regresion=False, use_mces20_log_loss=True, use_fingerprints=False, use_precursor_mz_for_model=True, tau_gumbel_softmax=10, gumbel_reg_weight=0.1, USE_LEARNABLE_MULTITASK=True, use_adduct=False, use_ce=False, use_ion_activation=False, use_ion_method=False, use_ion_mode=False, + pool: str = "attention", ): """Initialize the CCSPredictor""" super().__init__( d_model=d_model, n_layers=n_layers, dropout=dropout, weights=weights, lr=lr, use_element_wise=use_element_wise, use_cosine_distance=use_cosine_distance, use_adduct=use_adduct, use_ce=use_ce, use_ion_activation=use_ion_activation, use_ion_method=use_ion_method, use_ion_mode=use_ion_mode, + pool=pool, )🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/similarity_models.py` around lines 344 - 383, The SimilarityModelMultitask.__init__ method is missing the pool parameter that is now required by its parent class SimilarityModel.__init__. Add the pool parameter to the __init__ signature with an appropriate default value and pass it to the super().__init__() call along with the other parameters currently being passed (d_model, n_layers, dropout, weights, lr, use_element_wise, use_cosine_distance, use_adduct, use_ce, use_ion_activation, use_ion_method, use_ion_mode).
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@simba/core/data/datasets/multitask_dataset_builder.py`:
- Around line 5-9: The `encode_adduct` function is used in the
multitask_dataset_builder.py file at line 133 but is not included in the import
statement at the beginning of the file. Add `encode_adduct` to the existing
import list from `metabo_depthcharge.spec.metadata_parsers` along with the other
imported functions like `encode_collision_energy`, `encode_ion_activation`, and
`encode_ionization_method`. This will ensure the function is available when
`use_adduct=True` is set, preventing a NameError at runtime.
- Around line 138-142: The code directly accesses spec.ion_activation and
spec.ionization_method attributes which will raise AttributeError if these
attributes don't exist on the spec object. Replace these direct attribute
accesses with getattr calls using a None fallback pattern (e.g., getattr(spec,
"ion_activation", None)) to match the safer pattern already used in
encoder_dataset_builder.py. This should be applied to both the
encode_ion_activation call and the encode_ionization_method call in the
conditional blocks.
---
Outside diff comments:
In `@simba/core/models/similarity_models.py`:
- Around line 828-831: The ion_activation and ion_method tensors are being
converted to float type when they should be converted to long type, since the
spectrum encoder expects these as long tensors for embedding lookup operations.
Change the `.float()` method calls to `.long()` for both the ion_activation and
ion_method batch dictionary entries to match the expected input types for the
spectrum encoder.
- Around line 133-139: In the code block where use_ion_activation and
use_ion_method conditions are checked, change the dtype casting from .float() to
.long() for both ion_activation and ion_method. Specifically, in the if
self.use_ion_activation block, change batch["ion_activation_0"].float() and
batch["ion_activation_1"].float() to use .long() instead. Similarly, in the if
self.use_ion_method block, change batch["ion_method_0"].float() and
batch["ion_method_1"].float() to use .long() instead. This ensures consistency
with the spectrum encoder's expected input types and matches the correct
implementation in SimilarityModelMultitask.forward.
- Around line 344-383: The SimilarityModelMultitask.__init__ method is missing
the pool parameter that is now required by its parent class
SimilarityModel.__init__. Add the pool parameter to the __init__ signature with
an appropriate default value and pass it to the super().__init__() call along
with the other parameters currently being passed (d_model, n_layers, dropout,
weights, lr, use_element_wise, use_cosine_distance, use_adduct, use_ce,
use_ion_activation, use_ion_method, use_ion_mode).
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e4c8b61b-4961-467d-9c35-35071bc6b7a0
📒 Files selected for processing (23)
pyproject.tomlsimba/configs/model/simba_default.yamlsimba/core/chemistry/edit_distance/edit_distance.pysimba/core/data/datasets/encoder_dataset_builder.pysimba/core/data/datasets/multitask_dataset.pysimba/core/data/datasets/multitask_dataset_builder.pysimba/core/data/loaders/loaders.pysimba/core/models/similarity_models.pysimba/core/models/spectrum_encoder.pysimba/workflows/inference.pysimba/workflows/training.pytests/unit/test_embedder_multitask.pytests/unit/test_spectrum_ext.pytools/slurm/hyperparam_search.slurm.shtools/slurm/inference.slurm.shtools/slurm/inference_joint_metadata.slurm.shtools/slurm/inference_joint_metadata_cls.slurm.shtools/slurm/preprocessing.slurm.shtools/slurm/preprocessing_nist20.slurm.shtools/slurm/preprocessing_scaffold_v2.slurm.shtools/slurm/preprocessing_spectraverse.slurm.shtools/slurm/train.slurm.shtools/slurm/train_joint_metadata.slurm.sh
💤 Files with no reviewable changes (1)
- tests/unit/test_spectrum_ext.py
✅ Files skipped from review due to trivial changes (1)
- tools/slurm/preprocessing.slurm.sh
🚧 Files skipped from review as they are similar to previous changes (5)
- tools/slurm/inference_joint_metadata_cls.slurm.sh
- tools/slurm/train_joint_metadata.slurm.sh
- tools/slurm/inference_joint_metadata.slurm.sh
- simba/core/chemistry/edit_distance/edit_distance.py
- simba/core/data/loaders/loaders.py
Contributor
There was a problem hiding this comment.
Caution
Some comments are outside the diff and can’t be posted inline due to platform limitations.
⚠️ Outside diff range comments (1)
simba/core/models/similarity_models.py (1)
133-139:⚠️ Potential issue | 🟠 Major | ⚡ Quick winInconsistent dtype casting:
ion_activationandion_methodshould use.long()instead of.float().The base class
SimilarityModel.forwardcasts these categorical metadata fields to.float(), but:
- The encoder expects
.long()for categorical embedding lookups (perspectrum_encoder.pylines 102-108)- The dataset provides
int64(permultitask_dataset.py)- The subclass
SimilarityModelMultitask.forwardcorrectly uses.long()at lines 495-500While the encoder internally calls
.long(), this creates unnecessary dtype conversion and breaks consistency with the refactoring pattern applied toadductand in the subclass.Proposed fix
if self.use_ion_activation: - kwargs_0["ion_activation"] = batch["ion_activation_0"].float() - kwargs_1["ion_activation"] = batch["ion_activation_1"].float() + kwargs_0["ion_activation"] = batch["ion_activation_0"].long() + kwargs_1["ion_activation"] = batch["ion_activation_1"].long() if self.use_ion_method: - kwargs_0["ion_method"] = batch["ion_method_0"].float() - kwargs_1["ion_method"] = batch["ion_method_1"].float() + kwargs_0["ion_method"] = batch["ion_method_0"].long() + kwargs_1["ion_method"] = batch["ion_method_1"].long()🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the rest with a brief reason, keep changes minimal, and validate. In `@simba/core/models/similarity_models.py` around lines 133 - 139, The ion_activation and ion_method fields are being incorrectly cast to .float() when they should be cast to .long() for categorical embedding lookups, as these are categorical metadata fields provided as int64 from the dataset. Update the four assignments in the use_ion_activation and use_ion_method conditional blocks to use .long() instead of .float() for ion_activation_0, ion_activation_1, ion_method_0, and ion_method_1 to match the pattern used in the SimilarityModelMultitask.forward method and ensure consistency with the encoder's expectations.
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Outside diff comments:
In `@simba/core/models/similarity_models.py`:
- Around line 133-139: The ion_activation and ion_method fields are being
incorrectly cast to .float() when they should be cast to .long() for categorical
embedding lookups, as these are categorical metadata fields provided as int64
from the dataset. Update the four assignments in the use_ion_activation and
use_ion_method conditional blocks to use .long() instead of .float() for
ion_activation_0, ion_activation_1, ion_method_0, and ion_method_1 to match the
pattern used in the SimilarityModelMultitask.forward method and ensure
consistency with the encoder's expectations.
ℹ️ Review info
⚙️ Run configuration
Configuration used: defaults
Review profile: CHILL
Plan: Pro
Run ID: e466c874-2f77-454e-a8a6-1047d3df9e09
📒 Files selected for processing (3)
simba/core/data/datasets/multitask_dataset_builder.pysimba/core/models/similarity_models.pysimba/core/models/spectrum_encoder.py
🚧 Files skipped from review as they are similar to previous changes (2)
- simba/core/models/spectrum_encoder.py
- simba/core/data/datasets/multitask_dataset_builder.py
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
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.
Dependencies (pyproject.toml)
Added depthcharge-ms ≥0.4.9 and metabo-depthcharge (pinned to GitHub main)
Replaced local editable path for metabo-depthcharge with git+https://github.com/bittremieuxlab/metabo-depthcharge.git
Dead code removal
spectrum.py: stripped leftover spectrum_vector, mz_array, intensity_array, max_peak fields and the messy old getstate/setstate
encoding.py: deleted encode_adduct_mass, OneHotEncoding, ADDUCT_TO_MASS, import pandas — only ion activation/method helpers remain
preprocessor.py: removed large block of unused spectrum-vector preprocessing code
train_utils.py: removed ~12 lines of dead utility code
weighted_sampling.py: minor cleanup
ILP crash fix (edit_distance.py)
catch_errors=False → True; distance == -1 (solver failure) now maps to np.nan instead of crashing the preprocessing run
Adduct encoding (encoder_dataset_builder.py, multitask_dataset_builder.py, multitask_dataset.py)
Replaced one-hot adduct encoding over 30+ adducts with categorical integer index via encode_adduct() from metabo-depthcharge (9 known adducts + 0=unknown)
Adduct arrays changed from float32 to int64; similarity_models.py uses .long() instead of .float()
CE encoding (encoder_dataset_builder.py, multitask_dataset_builder.py)
Replaced raw float passthrough with encode_collision_energy() from metabo-depthcharge, which normalises CE÷100 and handles None/NaN/stepped strings
MGF loader (loaders.py)
CE field now read with fallback chain: collision_energy → collision_energy_1 → ce, covering MSG, Spectraverse and legacy formats
Spectrum encoder — full rewrite (spectrum_encoder.py, similarity_models.py)
Old: subclassed depthcharge's SpectrumTransformerEncoder, concatenate+project peak encoding, raw scalar global token, one-hot adduct via FloatEncoder
New: subclasses metabo-depthcharge's SpectrumEncoder — additive sinusoidal peak encoding, learned CLS + sinusoidal precursor_mz global token, MetadataEncoder for adduct (embedding) and CE (linear + zero-mask), pool="attention" (AttnAggregator over all tokens)
Zero imports from depthcharge in the encoder; ion_mode/activation/method use nn.Linear and are injected via global_token_hook override
similarity_models.py: removed 3 tuple-unpack (emb, _) sites and 3 [:, 0, :] slices — encoder now returns (B, d_model) directly
SLURM scripts (new)
train_joint_metadata.slurm.sh: trains on joint dataset with use_adduct, use_ce, use_ion_mode
inference_joint_metadata.slurm.sh: inference with the attn-pool model, proper paths.output_dir and model feature flags
inference_joint_metadata_cls.slurm.sh: inference script for the older CLS-pool checkpoint
Summary by CodeRabbit