msfiddle 2.1.0__tar.gz → 2.2.0__tar.gz

msfiddle-2.2.0/PKG-INFO

@@ -0,0 +1,169 @@
+Metadata-Version: 2.4
+Name: msfiddle
+Version: 2.2.0
+Summary: A package for predicting chemical formulas from tandem mass spectra
+Home-page: https://github.com/JosieHong/FIDDLE
+Author: Yuhui Hong
+Author-email: josieexception@outlook.com
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: OS Independent
+Classifier: Topic :: Scientific/Engineering :: Chemistry
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
+Classifier: Topic :: Scientific/Engineering :: Medical Science Apps.
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy<2.0.0,>=1.20.0
+Requires-Dist: pandas>=2.0.0
+Requires-Dist: tqdm>=4.60.0
+Requires-Dist: pyyaml>=6.0
+Requires-Dist: scikit-learn>=1.0.0
+Requires-Dist: scipy>=1.8.0
+Requires-Dist: pyarrow>=10.0.0
+Requires-Dist: rdkit>=2022.03.5
+Requires-Dist: molmass
+Requires-Dist: pyteomics
+Provides-Extra: inference
+Requires-Dist: torch; extra == "inference"
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: license-file
+Dynamic: provides-extra
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# FIDDLE
+
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![DOI](https://img.shields.io/badge/DOI-10.5281%2Fzenodo.19181279-blue)](https://doi.org/10.5281/zenodo.19181279)
+[![Release](https://img.shields.io/github/v/release/JosieHong/FIDDLE?label=Release)](https://github.com/JosieHong/FIDDLE/releases)
+
+[<img src="https://raw.githubusercontent.com/JosieHong/FIDDLE/main/img/fiddle_logo.png" align="right" width="220">](https://github.com/JosieHong/FIDDLE)
+
+FIDDLE is a deep learning method for predicting molecular formulas from MS/MS spectra. This repository contains the full research codebase for model training, evaluation, and paper reproduction, plus the [`msfiddle`](https://pypi.org/project/msfiddle/) PyPI package (in [`./msfiddle/`](https://github.com/JosieHong/FIDDLE/tree/main/msfiddle)) for inference-only use.
+
+- **Paper:** [Nature Communications (2025)](https://www.nature.com/articles/s41467-025-66060-9)
+- **CLI and Python API:** `pip install "msfiddle[inference]"` — see the [package docs](https://msfiddle.readthedocs.io)
+- **Try this demo!** [FIDDLE on Hugging Face](https://huggingface.co/spaces/J0siee/FIDDLE)
+
+<br clear="right"/>
+
+> **Breaking change (v2.0.0):** The rescore model has been redesigned (Siamese architecture), see details in [CHANGELOG.md](https://github.com/JosieHong/FIDDLE/blob/main/CHANGELOG.md).
+
+## Set up
+
+### Requirements
+
+1. Install [Anaconda](https://docs.anaconda.com/anaconda/install/), if not already installed.
+
+2. Create the environment with the necessary packages:
+
+```bash
+conda env create -f environment.yml
+```
+
+On Apple Silicon (macOS ARM64), use `environment_mac.yml` instead.
+
+3. (optional) Install [BUDDY](https://github.com/Philipbear/msbuddy) and [SIRIUS](https://v6.docs.sirius-ms.io/) following the respective installation instructions provided in each tool's documentation. 
+
+### Pre-trained Model Weights
+
+To use the pre-trained models, please use the following scripts to download the weights from the [release page](https://github.com/JosieHong/FIDDLE/releases/tag/v2.2.0) and place them in the `./check_point/` directory:
+
+- **Orbitrap models**:
+  - `fiddle_tcn_orbitrap.pt`: formula prediction model on Orbitrap spectra
+  - `fiddle_rescore_orbitrap.pt`: rescore model on Orbitrap spectra
+- **Q-TOF models**:
+  - `fiddle_tcn_qtof.pt`: formula prediction model on Q-TOF spectra
+  - `fiddle_rescore_qtof.pt`: rescore model on Q-TOF spectra
+
+```bash
+bash ./running_scripts/download_models.sh
+```
+
+## Usage
+
+The input format is `mgf`, where `title`, `precursor_mz`, `precursor_type`, `collision_energy` fields are required. Here, we sampled 21 spectra from the EMBL-MCF 2.0 dataset as an example.
+
+```mgf
+BEGIN IONS
+TITLE=EMBL_MCF_2_0_HRMS_Library000531
+PEPMASS=129.01941
+CHARGE=1-
+PRECURSOR_TYPE=[M-H]-
+PRECURSOR_MZ=129.01941
+COLLISION_ENERGY=50.0
+SMILES=[H]OC(=O)C([H])=C(C(=O)O[H])C([H])([H])[H]
+FORMULA=C5H6O4
+THEORETICAL_PRECURSOR_MZ=129.018785
+PPM=4.844255818912111
+SIMULATED_PRECURSOR_MZ=129.02032113281717
+41.2041 0.410228
+55.7698 0.503672
+56.8647 0.461943
+85.0296 100.0
+129.0196 8.036902
+END IONS
+```
+
+**Run FIDDLE!**
+
+```bash
+python scripts/run_fiddle.py --test_data ./msfiddle/demo/input_msms.mgf \
+                    --config_path ./msfiddle/config/fiddle_tcn_orbitrap.yml \
+                    --resume_path ./check_point/fiddle_tcn_orbitrap.pt \
+                    --rescore_resume_path ./check_point/fiddle_rescore_orbitrap.pt \
+                    --result_path ./msfiddle/demo/output_fiddle.csv --device 0
+```
+
+If you'd like to integrate the results from SIRIUS and BUDDY, please organize the results in the format shown in `./msfiddle/demo/buddy_output.csv` and `./msfiddle/demo/sirius_output.csv`, and provide them to run FIDDLE:
+
+```bash
+python scripts/run_fiddle.py --test_data ./msfiddle/demo/input_msms.mgf \
+                    --config_path ./msfiddle/config/fiddle_tcn_orbitrap.yml \
+                    --resume_path ./check_point/fiddle_tcn_orbitrap.pt \
+                    --rescore_resume_path ./check_point/fiddle_rescore_orbitrap.pt \
+                    --buddy_path ./msfiddle/demo/output_buddy.csv \
+                    --sirius_path ./msfiddle/demo/output_sirius.csv \
+                    --result_path ./msfiddle/demo/output_fiddle_all.csv --device 0
+```
+
+See [`scripts/test_caffeine.py`](https://github.com/JosieHong/FIDDLE/blob/main/scripts/test_caffeine.py) for a worked example running FIDDLE on a caffeine Orbitrap spectrum fetched live from GNPS.
+
+## Reproduce paper results
+
+All scripts should be run from the repository root (`FIDDLE/`).
+
+| Script | Description |
+|---|---|
+| `running_scripts/experiments_test_benchmark.sh` | Evaluate on external benchmarks (CASMI 2016, CASMI 2017, EMBL-MCF 2.0) |
+| `running_scripts/experiments_test_nist23.sh` | Evaluate on NIST23 |
+| `running_scripts/experiments_test_chimeric.sh` | Evaluate on chimeric spectra |
+| `running_scripts/experiments_test_noised.sh` | Evaluate under noise conditions |
+| `running_scripts/experiments_ablation_study.sh` | Run ablation study |
+| `running_scripts/experiments_demo.sh` | Run demo experiment |
+| `running_scripts/train_released_models.sh` | Train TCN and rescore models for both Orbitrap and Q-TOF |
+
+For training from scratch, see the train scripts (`scripts/train_tcn_gpus.py`, `scripts/train_tcn_gpus_cl.py`, `scripts/train_rescore.py`) and the corresponding config files in `./msfiddle/config/`.
+
+## Citation
+
+```
+@article{hong2025fiddle,
+  title={FIDDLE: a deep learning method for chemical formulas prediction from tandem mass spectra},
+  author={Hong, Yuhui and Li, Sujun and Ye, Yuzhen and Tang, Haixu},
+  journal={Nature Communications},
+  volume={16},
+  number={1},
+  pages={11102},
+  year={2025},
+  publisher={Nature Publishing Group UK London}
+}
+```

msfiddle-2.2.0/README.md

@@ -0,0 +1,128 @@
+# FIDDLE
+
+[![License](https://img.shields.io/badge/License-Apache%202.0-blue.svg)](https://opensource.org/licenses/Apache-2.0)
+[![DOI](https://img.shields.io/badge/DOI-10.5281%2Fzenodo.19181279-blue)](https://doi.org/10.5281/zenodo.19181279)
+[![Release](https://img.shields.io/github/v/release/JosieHong/FIDDLE?label=Release)](https://github.com/JosieHong/FIDDLE/releases)
+
+[<img src="https://raw.githubusercontent.com/JosieHong/FIDDLE/main/img/fiddle_logo.png" align="right" width="220">](https://github.com/JosieHong/FIDDLE)
+
+FIDDLE is a deep learning method for predicting molecular formulas from MS/MS spectra. This repository contains the full research codebase for model training, evaluation, and paper reproduction, plus the [`msfiddle`](https://pypi.org/project/msfiddle/) PyPI package (in [`./msfiddle/`](https://github.com/JosieHong/FIDDLE/tree/main/msfiddle)) for inference-only use.
+
+- **Paper:** [Nature Communications (2025)](https://www.nature.com/articles/s41467-025-66060-9)
+- **CLI and Python API:** `pip install "msfiddle[inference]"` — see the [package docs](https://msfiddle.readthedocs.io)
+- **Try this demo!** [FIDDLE on Hugging Face](https://huggingface.co/spaces/J0siee/FIDDLE)
+
+<br clear="right"/>
+
+> **Breaking change (v2.0.0):** The rescore model has been redesigned (Siamese architecture), see details in [CHANGELOG.md](https://github.com/JosieHong/FIDDLE/blob/main/CHANGELOG.md).
+
+## Set up
+
+### Requirements
+
+1. Install [Anaconda](https://docs.anaconda.com/anaconda/install/), if not already installed.
+
+2. Create the environment with the necessary packages:
+
+```bash
+conda env create -f environment.yml
+```
+
+On Apple Silicon (macOS ARM64), use `environment_mac.yml` instead.
+
+3. (optional) Install [BUDDY](https://github.com/Philipbear/msbuddy) and [SIRIUS](https://v6.docs.sirius-ms.io/) following the respective installation instructions provided in each tool's documentation. 
+
+### Pre-trained Model Weights
+
+To use the pre-trained models, please use the following scripts to download the weights from the [release page](https://github.com/JosieHong/FIDDLE/releases/tag/v2.2.0) and place them in the `./check_point/` directory:
+
+- **Orbitrap models**:
+  - `fiddle_tcn_orbitrap.pt`: formula prediction model on Orbitrap spectra
+  - `fiddle_rescore_orbitrap.pt`: rescore model on Orbitrap spectra
+- **Q-TOF models**:
+  - `fiddle_tcn_qtof.pt`: formula prediction model on Q-TOF spectra
+  - `fiddle_rescore_qtof.pt`: rescore model on Q-TOF spectra
+
+```bash
+bash ./running_scripts/download_models.sh
+```
+
+## Usage
+
+The input format is `mgf`, where `title`, `precursor_mz`, `precursor_type`, `collision_energy` fields are required. Here, we sampled 21 spectra from the EMBL-MCF 2.0 dataset as an example.
+
+```mgf
+BEGIN IONS
+TITLE=EMBL_MCF_2_0_HRMS_Library000531
+PEPMASS=129.01941
+CHARGE=1-
+PRECURSOR_TYPE=[M-H]-
+PRECURSOR_MZ=129.01941
+COLLISION_ENERGY=50.0
+SMILES=[H]OC(=O)C([H])=C(C(=O)O[H])C([H])([H])[H]
+FORMULA=C5H6O4
+THEORETICAL_PRECURSOR_MZ=129.018785
+PPM=4.844255818912111
+SIMULATED_PRECURSOR_MZ=129.02032113281717
+41.2041 0.410228
+55.7698 0.503672
+56.8647 0.461943
+85.0296 100.0
+129.0196 8.036902
+END IONS
+```
+
+**Run FIDDLE!**
+
+```bash
+python scripts/run_fiddle.py --test_data ./msfiddle/demo/input_msms.mgf \
+                    --config_path ./msfiddle/config/fiddle_tcn_orbitrap.yml \
+                    --resume_path ./check_point/fiddle_tcn_orbitrap.pt \
+                    --rescore_resume_path ./check_point/fiddle_rescore_orbitrap.pt \
+                    --result_path ./msfiddle/demo/output_fiddle.csv --device 0
+```
+
+If you'd like to integrate the results from SIRIUS and BUDDY, please organize the results in the format shown in `./msfiddle/demo/buddy_output.csv` and `./msfiddle/demo/sirius_output.csv`, and provide them to run FIDDLE:
+
+```bash
+python scripts/run_fiddle.py --test_data ./msfiddle/demo/input_msms.mgf \
+                    --config_path ./msfiddle/config/fiddle_tcn_orbitrap.yml \
+                    --resume_path ./check_point/fiddle_tcn_orbitrap.pt \
+                    --rescore_resume_path ./check_point/fiddle_rescore_orbitrap.pt \
+                    --buddy_path ./msfiddle/demo/output_buddy.csv \
+                    --sirius_path ./msfiddle/demo/output_sirius.csv \
+                    --result_path ./msfiddle/demo/output_fiddle_all.csv --device 0
+```
+
+See [`scripts/test_caffeine.py`](https://github.com/JosieHong/FIDDLE/blob/main/scripts/test_caffeine.py) for a worked example running FIDDLE on a caffeine Orbitrap spectrum fetched live from GNPS.
+
+## Reproduce paper results
+
+All scripts should be run from the repository root (`FIDDLE/`).
+
+| Script | Description |
+|---|---|
+| `running_scripts/experiments_test_benchmark.sh` | Evaluate on external benchmarks (CASMI 2016, CASMI 2017, EMBL-MCF 2.0) |
+| `running_scripts/experiments_test_nist23.sh` | Evaluate on NIST23 |
+| `running_scripts/experiments_test_chimeric.sh` | Evaluate on chimeric spectra |
+| `running_scripts/experiments_test_noised.sh` | Evaluate under noise conditions |
+| `running_scripts/experiments_ablation_study.sh` | Run ablation study |
+| `running_scripts/experiments_demo.sh` | Run demo experiment |
+| `running_scripts/train_released_models.sh` | Train TCN and rescore models for both Orbitrap and Q-TOF |
+
+For training from scratch, see the train scripts (`scripts/train_tcn_gpus.py`, `scripts/train_tcn_gpus_cl.py`, `scripts/train_rescore.py`) and the corresponding config files in `./msfiddle/config/`.
+
+## Citation
+
+```
+@article{hong2025fiddle,
+  title={FIDDLE: a deep learning method for chemical formulas prediction from tandem mass spectra},
+  author={Hong, Yuhui and Li, Sujun and Ye, Yuzhen and Tang, Haixu},
+  journal={Nature Communications},
+  volume={16},
+  number={1},
+  pages={11102},
+  year={2025},
+  publisher={Nature Publishing Group UK London}
+}
+```

{msfiddle-2.1.0 → msfiddle-2.2.0}/msfiddle/config/fiddle_tcn_orbitrap.yml

@@ -131,12 +131,14 @@ model:
 # ------------------------
 train: 
   epochs: 500
-  batch_size: 512 # 128
+  batch_size: 512 # 128 fits 11GB GPUs (CL training holds 2 forward graphs, memory ~ 2x batch)
+  grad_accum_steps: 1 # optimizer step every N batches; effective batch = batch_size x N (256 x 2 on 2 x 11GB GPUs = 512)
   num_workers: 8
   lr: 0.001
   weight_decay: 0.01
   warmup_ratio: 0.1
-  early_stop_step: 10
+  early_stop_step: 15 # ~3x plateau patience so >=2 LR halvings get a chance
+  early_stop_min_epochs: 5 # patience does not count before this epoch
   patience: 5
   noised_times: 0
   
@@ -150,15 +152,18 @@ train_rescore:
 # ------------------------
 # postprocessing settings
 # ------------------------
-post_processing: 
+post_processing:
   # Orbitrap
-  refine_atom_type: ['C', 'O', 'N', 'H']
-  refine_atom_num: [3, 3, 3, -1]
+  # D0c: halogens in the refinement search so a missed halogen can be introduced and
+  # H rebalanced to the exact precursor mass. (Orbitrap is halogen-rich, so the model
+  # usually predicts halogens already; this mainly helps the cases it misses.)
+  refine_atom_type: ['C', 'O', 'N', 'H', 'Cl', 'Br', 'F', 'I', 'S', 'P']
+  refine_atom_num: [5, 5, 5, -1, 2, 1, 2, 1, 2, 1] # config E caps (08/2026 sweep): recall@20 88.5% -> 94.3% at 10ppm, 95.7% if 5ppm window validates
 
-  mass_tolerance: 20.0 # target space
+  mass_tolerance: 5.0 # config E (08/2026 sweep): matches the data filters' ppm_tolerance 5 (no admitted spectrum can be orphaned) and Orbitrap instrument accuracy; halves the decoy count vs 10 ppm, recall@20 ceiling 95.7%
   ppm_mode: True
   time_out: 300
   maxium_miss_atom_num: 20 # search space
-  top_k: 5
+  top_k: 20 # config E (08/2026 sweep): +2 points of rescorer ceiling over K=10; keep pair-generation and inference matched
   buddy_fdr_thr: 0.6
   sirius_score_thr: -5

{msfiddle-2.1.0 → msfiddle-2.2.0}/msfiddle/config/fiddle_tcn_qtof.yml

@@ -165,17 +165,19 @@ model:
 # ------------------------
 train: 
   epochs: 200
-  batch_size: 512 # 128
+  batch_size: 512 # 128 fits 11GB GPUs (CL training holds 2 forward graphs, memory ~ 2x batch)
+  grad_accum_steps: 1 # optimizer step every N batches; effective batch = batch_size x N (256 x 2 on 2 x 11GB GPUs = 512)
   num_workers: 8
   lr: 0.002 # 0.001
   weight_decay: 0.01
   warmup_ratio: 0.1
-  early_stop_step: 10
+  early_stop_step: 15 # ~3x plateau patience so >=2 LR halvings get a chance
+  early_stop_min_epochs: 5 # patience does not count before this epoch
   patience: 5
   noised_times: 3
   
 train_rescore:
-  epochs: 100
+  epochs: 20
   batch_size: 128
   num_workers: 8
   lr: 0.0001
@@ -184,14 +186,19 @@ train_rescore:
 # ------------------------
 # postprocessing settings
 # ------------------------
-post_processing: 
-  refine_atom_type: ['C', 'O', 'N', 'H']
-  refine_atom_num: [5, 5, 5, -1]
+post_processing:
+  # D0c (tuned): let refinement INTRODUCE a missed halogen and rebalance H to the
+  # exact precursor mass. Base was ['C','O','N','H'] -> halogens unreachable.
+  # S/P deliberately NOT in base (already reachable when the model predicts them;
+  # adding them only created spurious candidates). Halogen caps kept tight to limit
+  # near-isobaric false positives on the CONH majority.
+  refine_atom_type: ['C', 'O', 'N', 'H', 'Cl', 'Br', 'F', 'I', 'S', 'P']
+  refine_atom_num: [5, 5, 5, -1, 2, 1, 2, 1, 2, 1] # config B (08/2026 sweep): S/P reachable; recall@20 92.3% -> 94.0%
 
   mass_tolerance: 10.0 # target space
   ppm_mode: True
   time_out: -1 # do not limit the time
   maxium_miss_atom_num: 50 # search space
-  top_k: 5
+  top_k: 20 # config B (08/2026 sweep): +2 points of rescorer ceiling over K=10; keep pair-generation and inference matched
   buddy_fdr_thr: 0.6
   sirius_score_thr: -5

{msfiddle-2.1.0 → msfiddle-2.2.0}/msfiddle/dataset.py

@@ -55,7 +55,7 @@ class MS2FDataset_CL(Dataset):
         for data_item in self.data:
             spec = data_item["spec"]
             data_item["spec"] = np.concatenate(
-                (spec, np.zeros((padding_dim, spec.shape[1]))), axis=0
+                (spec, np.zeros((padding_dim, spec.shape[1]), dtype=spec.dtype)), axis=0
             )
 
     def augment_with_noise(self, data, pairs, noised_times):
@@ -76,11 +76,16 @@ class MS2FDataset_CL(Dataset):
         """
         original_length = len(data)
         augmented_data = copy.deepcopy(data)
-        for data_item in tqdm(data, desc="Data Augmentation"):
-            spec = data_item["spec"][:, 0]
-            spec_mask = np.where(spec > 0, 1.0, 0.0)
+        # Copies are appended in whole-dataset blocks (round t of every
+        # spectrum, then round t+1, ...) so that augmented copy t of item j
+        # sits at index original_length * (t + 1) + j — the layout the pair
+        # index shift below assumes. Interleaving copies per spectrum instead
+        # silently pairs unrelated molecules once noised_times >= 2.
+        for _ in range(noised_times):
+            for data_item in tqdm(data, desc="Data Augmentation"):
+                spec = data_item["spec"][:, 0]
+                spec_mask = np.where(spec > 0, 1.0, 0.0)
 
-            for _ in range(noised_times):
                 noise = np.random.normal(0, 0.1, len(spec))
                 new_spec = spec + noise * spec_mask
 
@@ -151,7 +156,7 @@ class MS2FDataset(Dataset):
         for data_item in self.data:
             spec = data_item["spec"]
             data_item["spec"] = np.concatenate(
-                (spec, np.zeros((padding_dim, spec.shape[1]))), axis=0
+                (spec, np.zeros((padding_dim, spec.shape[1]), dtype=spec.dtype)), axis=0
             )
 
     def augment_with_noise(self, data, noised_times):
@@ -293,6 +298,18 @@ class MGFDataset(Dataset):
                          'precursor_type', and 'use_simulated_precursor_mz'.
         """
         for spectrum in spectra:
+            # Skip (not crash on) adducts the encoder does not support: user
+            # MGFs routinely contain a few, and one unsupported PRECURSOR_TYPE
+            # must not abort the whole run with a KeyError.
+            precursor_type = spectrum["params"]["precursor_type"]
+            if precursor_type not in encoder["type2charge"]:
+                print(
+                    "Skip {}: unsupported precursor type {}".format(
+                        spectrum["params"].get("title", "?"), precursor_type
+                    )
+                )
+                continue
+
             good_spec, _, _, spec_arr = generate_ms(
                 x=spectrum["m/z array"],
                 y=spectrum["intensity array"],
@@ -309,6 +326,15 @@ class MGFDataset(Dataset):
             adjust_neutral_add_vec, adjust_precursor_type = self.melt_neutral_precursor(
                 spectrum["params"]["precursor_type"]
             )
+            if adjust_precursor_type not in encoder["precursor_type"]:
+                print(
+                    "Skip {}: unsupported precursor type {} (from {})".format(
+                        spectrum["params"].get("title", "?"),
+                        adjust_precursor_type,
+                        precursor_type,
+                    )
+                )
+                continue
 
             ce, nce = parse_collision_energy(
                 ce_str=spectrum["params"]["collision_energy"],

{msfiddle-2.1.0 → msfiddle-2.2.0}/msfiddle/download.py

@@ -16,11 +16,18 @@ from importlib.metadata import (
 
 
 def checkpoint_release_for_package_version(package_version):
-    """Return the major-version FIDDLE release tag that contains checkpoints."""
-    match = re.match(r"^(\d+)", str(package_version))
+    """Return the FIDDLE release tag that holds this version's checkpoints.
+
+    From 2.2 on, each minor version has its own checkpoints (tag vX.Y.0,
+    shared by its patch releases); 2.0 and 2.1 share the v2.0.0 assets.
+    """
+    match = re.match(r"^(\d+)(?:\.(\d+))?", str(package_version))
     if match is None:
         return f"v{package_version}"
-    return f"v{match.group(1)}.0.0"
+    major, minor = int(match.group(1)), int(match.group(2) or 0)
+    if (major, minor) < (2, 2):
+        return f"v{major}.0.0"
+    return f"v{major}.{minor}.0"
 
 
 try:
@@ -56,10 +63,14 @@ def download_url(url, output_path):
 def get_checkpoint_dir():
     """Get the directory where model checkpoints are stored.
 
-    Returns ``~/.msfiddle/check_point``, a user-level directory that is
-    consistent regardless of working directory or install method.
+    Returns ``~/.msfiddle/check_point/<release tag>``, a user-level directory
+    that is consistent regardless of working directory or install method. The
+    tag keeps checkpoints from different releases apart, so an upgrade never
+    silently reuses an older release's weights.
     """
-    return os.path.join(os.path.expanduser("~"), ".msfiddle", "check_point")
+    return os.path.join(
+        os.path.expanduser("~"), ".msfiddle", "check_point", FIDDLE_RELEASE
+    )
 
 
 def download_models(destination=None, models=None):

msfiddle-2.2.0/msfiddle/optim.py

@@ -0,0 +1,160 @@
+"""Muon optimizer (MomentUm Orthogonalized by Newton-Schulz).
+
+Muon applies momentum SGD to matrix-shaped hidden weights and orthogonalizes
+each update with a Newton-Schulz iteration before applying it, which keeps the
+update's singular values near 1 (an approximate steepest descent under the
+spectral norm). Reference: Keller Jordan et al., 2024
+(https://kellerjordan.github.io/posts/muon/).
+
+Only 2D+ hidden weights benefit; embeddings, biases, and per-channel gains
+(weight_norm's weight_g) keep AdamW. The Muon class below therefore carries
+two param groups, one per update rule, so schedulers and checkpointing see a
+single ordinary optimizer.
+"""
+
+import torch
+import torch.nn as nn
+from torch.optim import Optimizer
+
+
+def zeropower_via_newtonschulz5(G, steps=5, eps=1e-7):
+    """Approximate UV^T of the SVD of G via a quintic Newton-Schulz iteration.
+
+    Runs in float32 (not bfloat16 as in the reference) for correctness on
+    pre-Ampere GPUs; the matrices here are small enough that this is cheap.
+    The iteration's singular values land in roughly [0.7, 1.2], which is all
+    Muon needs.
+    """
+    assert G.ndim == 2
+    a, b, c = (3.4445, -4.7750, 2.0315)
+    X = G.float()
+    transposed = G.size(0) > G.size(1)
+    if transposed:
+        X = X.T
+    X = X / (X.norm() + eps)
+    for _ in range(steps):
+        A = X @ X.T
+        B = b * A + c * A @ A
+        X = a * X + B @ X
+    if transposed:
+        X = X.T
+    return X.to(G.dtype)
+
+
+def split_matrix_params(model):
+    """Partition parameters into (muon_params, adamw_params).
+
+    Muon gets matrix-shaped hidden weights: ndim >= 2, not an nn.Embedding
+    weight, not weight_norm's per-channel gain (weight_g). Everything else
+    (biases, gains, embeddings) goes to AdamW.
+    """
+    if isinstance(model, nn.DataParallel):
+        model = model.module
+    emb_ids = set()
+    for m in model.modules():
+        if isinstance(m, nn.Embedding):
+            emb_ids.update(id(p) for p in m.parameters(recurse=False))
+    muon_params, adamw_params = [], []
+    for name, p in model.named_parameters():
+        if not p.requires_grad:
+            continue
+        if p.ndim >= 2 and id(p) not in emb_ids and not name.endswith("weight_g"):
+            muon_params.append(p)
+        else:
+            adamw_params.append(p)
+    return muon_params, adamw_params
+
+
+class Muon(Optimizer):
+    """Muon for matrix params, decoupled AdamW for everything else.
+
+    Conv kernels (out, in, k) are flattened to (out, in*k) for the
+    orthogonalization, following the Moonlight/Kimi treatment. The update is
+    scaled by sqrt(max(1, rows/cols)) so its RMS matches AdamW's, letting one
+    LR schedule drive both groups.
+    """
+
+    def __init__(
+        self,
+        muon_params,
+        adamw_params,
+        lr=0.02,
+        momentum=0.95,
+        nesterov=True,
+        ns_steps=5,
+        adamw_lr=1e-3,
+        adamw_betas=(0.9, 0.999),
+        adamw_eps=1e-8,
+        weight_decay=0.01,
+    ):
+        groups = [
+            dict(
+                params=list(muon_params),
+                lr=lr,
+                momentum=momentum,
+                nesterov=nesterov,
+                ns_steps=ns_steps,
+                weight_decay=weight_decay,
+                use_muon=True,
+            ),
+            dict(
+                params=list(adamw_params),
+                lr=adamw_lr,
+                betas=adamw_betas,
+                eps=adamw_eps,
+                weight_decay=weight_decay,
+                use_muon=False,
+            ),
+        ]
+        super().__init__(groups, {})
+
+    @torch.no_grad()
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            with torch.enable_grad():
+                loss = closure()
+        for group in self.param_groups:
+            if group["use_muon"]:
+                self._muon_group_step(group)
+            else:
+                self._adamw_group_step(group)
+        return loss
+
+    def _muon_group_step(self, group):
+        for p in group["params"]:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            if "momentum_buffer" not in state:
+                state["momentum_buffer"] = torch.zeros_like(g)
+            buf = state["momentum_buffer"]
+            buf.mul_(group["momentum"]).add_(g)
+            update = g.add(buf, alpha=group["momentum"]) if group["nesterov"] else buf
+            u2 = update.reshape(update.size(0), -1)
+            u2 = zeropower_via_newtonschulz5(u2, steps=group["ns_steps"])
+            scale = max(1.0, u2.size(0) / u2.size(1)) ** 0.5
+            p.mul_(1.0 - group["lr"] * group["weight_decay"])
+            p.add_(u2.reshape_as(p), alpha=-group["lr"] * scale)
+
+    def _adamw_group_step(self, group):
+        beta1, beta2 = group["betas"]
+        for p in group["params"]:
+            g = p.grad
+            if g is None:
+                continue
+            state = self.state[p]
+            if "exp_avg" not in state:
+                state["step"] = 0
+                state["exp_avg"] = torch.zeros_like(g)
+                state["exp_avg_sq"] = torch.zeros_like(g)
+            state["step"] += 1
+            exp_avg, exp_avg_sq = state["exp_avg"], state["exp_avg_sq"]
+            exp_avg.mul_(beta1).add_(g, alpha=1 - beta1)
+            exp_avg_sq.mul_(beta2).addcmul_(g, g, value=1 - beta2)
+            bc1 = 1 - beta1 ** state["step"]
+            bc2 = 1 - beta2 ** state["step"]
+            denom = (exp_avg_sq / bc2).sqrt_().add_(group["eps"])
+            p.mul_(1.0 - group["lr"] * group["weight_decay"])
+            p.addcdiv_(exp_avg, denom, value=-group["lr"] / bc1)