Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+architecture.png filter=lfs diff=lfs merge=lfs -text

DecisionTree.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:a2ec37a3f117ef6cd281d12ffbf2a91779513a4c80b230ea7def4552628e57df
+size 4984593

README.md

@@ -14,7 +14,7 @@ datasets:
 # DOM Formula Assignment using Machine Learning
 
 
-![Model Type](https://img.shields.io/badge/Model-KNN-blue)
+![Model Type](https://img.shields.io/badge/Model-ML-blue)
 ![Data](https://img.shields.io/badge/Data-FT--ICR_MS-green)
 ![License](https://img.shields.io/badge/License-CC_BY_NC_ND_4-yellow)
 [![GitHub](https://img.shields.io/badge/GitHub-pcdslab/dom--formula--assignment--using--ml-blue?logo=github)](https://github.com/pcdslab/dom-formula-assignment-using-ml)
@@ -30,47 +30,153 @@ datasets:
 ## Abstract
 A machine learning approach to molecular formula assignment is crucial for unlocking the full potential of ultra-high resolution mass spectrometry (UHRMS) when analyzing complex mixtures. By combining data-driven models with rigorous benchmarking, the accuracy, consistency, and speed in identifying plausible molecular formulas from vast spectral datasets can be improved. Compared with traditional de novo methods that rely heavily on rule-based heuristics, and manual parameter tuning, machine learning approaches can capture complex patterns in data and adapt more readily to diverse sample types. In this paper, we describe the application of a machine learning methods using the k-nearest neighbors (KNN) algorithm trained on curated chemical formula datasets of UHRMS analysis of dissolved organic matter (DOM) covering the saline river continuum and tropical wet/dry season variability. The influence of the mass accuracy (training set with 0.15-1ppm) was evaluated on a blind test set of DOMs of different geographical origins. A Decision Tree Regressor (DTR) and Random Forest Regressor (RFR) based on mass accuracy (<1ppm) was used. Results from our ML models exhibit 43% more formulas annotated than traditional methods (5796 vs 4047), Model-Synthetic achieved 99.9% assignment rate and annotated/assigned 2x more formulas (8,268 vs 4047). DTR and RFR achieved formula-level accuracies (FA) of 86.5% and 60.4%, respectively. Overall, results show an increase in formula assignment when compared with traditional methods. This ultimately enables more reliable characterization of complex natural and engineered systems, supporting advances in fields such as environmental science, metabolomics, and petroleomics. Furthermore, the novel data set produced for this study is made publicly available, establishing an initial benchmark for molecular formula assignment in UHRMS using machine learning. The dataset and code are publicly available at: https://github.com/pcdslab/dom-formula-assignment-using-ml
 
-## Available Models
+## Architecture
+
+![DOM formula assignment architecture](architecture.png)
+(a) KNN Pipeline for Formula Assignment/Annotation. For each train set, the first KNN model is fitted, and then formula assignment/annotation is performed using the closest peaks in the training set. (b) Decision Tree Regressor (DTR) and Random Forest Regressor (RFR) are trained to predict element counts for CHONS using squared-error criterion.
+## Usage
+
+Install the dependencies:
+
+```bash
+pip install "numpy==2.4.1" "scikit-learn==1.8.0" joblib huggingface_hub transformers torch pandas
+```
+
+
+### KNN Formula Assignment
+
+Load the default KNN model:
+
+```python
+from transformers import AutoModel
+import torch
+
+model = AutoModel.from_pretrained(
+    "SaeedLab/dom-formula-assignment-using-ml",
+    trust_remote_code=True,
+).eval()
+
+# Mass-only input. Each value is one mass to assign.
+masses = torch.tensor([350.123456, 421.987654], dtype=torch.float64)
+
+with torch.no_grad():
+    outputs = model(features=masses, return_neighbors=True, n_neighbors=3)
+
+print("Predicted formulas:", outputs.predictions)
+print("Neighbor indices:", outputs.indices)
+```
+
+Load a specific KNN model by passing `model_name`:
+
+```python
+from transformers import AutoModel
+
+model = AutoModel.from_pretrained(
+    "SaeedLab/dom-formula-assignment-using-ml",
+    trust_remote_code=True,
+    model_name="L1-L3_K3_Manhattan_Ensemble",
+).eval()
+```
+
+Use the `Model Name` value from the table below. For example, `Synthetic_K3_Euclidean_Ensemble` selects the default synthetic KNN ensemble.
+
+There is no tokenizer step for these KNN models. Inputs are 1D numeric vectors of mass values.
+
+### Get nearest neighbors
+
+KNN models can also return the closest training examples for each input sample:
+
+```python
+neighbor_indices = model.neighbor_indices(masses, n_neighbors=3)
+
+print("Neighbor indices:")
+print(neighbor_indices)
+```
+
+Use the model variant that matches your experiment setup:
+
+- `L1`, `L3`, `L1-L3`, or `Synthetic` for the training dataset.
+- `K1` or `K3` for the number of neighbors.
+- `Euclidean` or `Manhattan` for the distance metric.
+- `Ensemble`, `7T`, `21T`, or `SYN` for combined, field-strength, or synthetic backing models.
+
+### Decision Tree and Random Forest Formula Regressors
+
+The Decision Tree and Random Forest models predict CHONS element counts from mass and ion mobility features. Inputs must be numeric rows in this order:
+
+```text
+[mz, inv_k0, ccs]
+```
+
+Here, `inv_k0` is `1/k0`, and both `inv_k0` and `ccs` are mobility-derived values.
+
+```python
+from transformers import AutoModel
+import torch
+
+model = AutoModel.from_pretrained(
+    "SaeedLab/dom-formula-assignment-using-ml",
+    trust_remote_code=True,
+    model_name="RandomForest",
+).eval()
+
+features = torch.tensor(
+    [
+        [191.03498, 0.6808667247437136, 146.2048497276655],
+        [191.071388, 0.7087887538442335, 152.19879035313514],
+        [191.10775, 0.6935974811483341, 148.93494377261592],
+    ],
+    dtype=torch.float64,
+)
+
+with torch.no_grad():
+    outputs = model(features=features)
+
+print("Predicted formulas:", outputs.formulas)
+print("Predicted CHONS counts:", outputs.formula_counts)
+```
+
+The CHONS output columns are ordered as:
+
+```text
+[C, H, O, N, S]
+```
+
+## KNN Models
 
 This repository contains pre-trained KNN models used for molecular formula assignment. The models vary by:
-- **Dataset**: L1, L3, L1-L3, Synthetic
+- **Dataset**: L1 (7T), L3 (21T), L1-L3, Synthetic
 - **Neighbors (K)**: 1, 3
 - **Distance Metric**: Euclidean, Manhattan
-- **Field Strength/Type**: 7T, 21T, SYN (Synthetic)
+- **Field Strength/Type**: Ensemble, 7T, 21T, SYN (Synthetic)
 
 ### Model List
 
-| Model Filename | Description |
+| Model Name | Description |
 |---|---|
-| `knn_model_Model-L1_K1_Euclidean.joblib` | L1 Dataset, K=1, Euclidean Distance |
-| `knn_model_Model-L1_K1_Manhattan.joblib` | L1 Dataset, K=1, Manhattan Distance |
-| `knn_model_Model-L1_K3_Euclidean.joblib` | L1 Dataset, K=3, Euclidean Distance |
-| `knn_model_Model-L1_K3_Manhattan.joblib` | L1 Dataset, K=3, Manhattan Distance |
-| `knn_model_Model-L3_K1_Euclidean.joblib` | L3 Dataset, K=1, Euclidean Distance |
-| `knn_model_Model-L3_K1_Manhattan.joblib` | L3 Dataset, K=1, Manhattan Distance |
-| `knn_model_Model-L3_K3_Euclidean.joblib` | L3 Dataset, K=3, Euclidean Distance |
-| `knn_model_Model-L3_K3_Manhattan.joblib` | L3 Dataset, K=3, Manhattan Distance |
-| `knn_model_Model-L1-L3_K1_Euclidean_21T.joblib` | L1-L3 Combined, K=1, Euclidean, 21T |
-| `knn_model_Model-L1-L3_K1_Euclidean_7T.joblib` | L1-L3 Combined, K=1, Euclidean, 7T |
-| `knn_model_Model-L1-L3_K1_Manhattan_21T.joblib` | L1-L3 Combined, K=1, Manhattan, 21T |
-| `knn_model_Model-L1-L3_K1_Manhattan_7T.joblib` | L1-L3 Combined, K=1, Manhattan, 7T |
-| `knn_model_Model-L1-L3_K3_Euclidean_21T.joblib` | L1-L3 Combined, K=3, Euclidean, 21T |
-| `knn_model_Model-L1-L3_K3_Euclidean_7T.joblib` | L1-L3 Combined, K=3, Euclidean, 7T |
-| `knn_model_Model-L1-L3_K3_Manhattan_21T.joblib` | L1-L3 Combined, K=3, Manhattan, 21T |
-| `knn_model_Model-L1-L3_K3_Manhattan_7T.joblib` | L1-L3 Combined, K=3, Manhattan, 7T |
-| `knn_model_Model-Synthetic_K1_Euclidean_21T.joblib` | Synthetic, K=1, Euclidean, 21T |
-| `knn_model_Model-Synthetic_K1_Euclidean_7T.joblib` | Synthetic, K=1, Euclidean, 7T |
-| `knn_model_Model-Synthetic_K1_Euclidean_SYN.joblib` | Synthetic, K=1, Euclidean, SYN |
-| `knn_model_Model-Synthetic_K1_Manhattan_21T.joblib` | Synthetic, K=1, Manhattan, 21T |
-| `knn_model_Model-Synthetic_K1_Manhattan_7T.joblib` | Synthetic, K=1, Manhattan, 7T |
-| `knn_model_Model-Synthetic_K1_Manhattan_SYN.joblib` | Synthetic, K=1, Manhattan, SYN |
-| `knn_model_Model-Synthetic_K3_Euclidean_21T.joblib` | Synthetic, K=3, Euclidean, 21T |
-| `knn_model_Model-Synthetic_K3_Euclidean_7T.joblib` | Synthetic, K=3, Euclidean, 7T |
-| `knn_model_Model-Synthetic_K3_Euclidean_SYN.joblib` | Synthetic, K=3, Euclidean, SYN |
-| `knn_model_Model-Synthetic_K3_Manhattan_21T.joblib` | Synthetic, K=3, Manhattan, 21T |
-| `knn_model_Model-Synthetic_K3_Manhattan_7T.joblib` | Synthetic, K=3, Manhattan, 7T |
-| `knn_model_Model-Synthetic_K3_Manhattan_SYN.joblib` | Synthetic, K=3, Manhattan, SYN |
-
+| `L1_K1_Euclidean` | L1 (7T) KNN model with 1 neighbor and Euclidean distance |
+| `L1_K1_Manhattan` | L1 (7T) KNN model with 1 neighbor and Manhattan distance |
+| `L1_K3_Euclidean` | L1 (7T) KNN model with 3 neighbors and Euclidean distance |
+| `L1_K3_Manhattan` | L1 (7T) KNN model with 3 neighbors and Manhattan distance |
+| `L3_K1_Euclidean` | L3 (21T) KNN model with 1 neighbor and Euclidean distance |
+| `L3_K1_Manhattan` | L3 (21T) KNN model with 1 neighbor and Manhattan distance |
+| `L3_K3_Euclidean` | L3 (21T) KNN model with 3 neighbors and Euclidean distance |
+| `L3_K3_Manhattan` | L3 (21T) KNN model with 3 neighbors and Manhattan distance |
+| `L1-L3_K1_Euclidean_Ensemble` | Ensemble of L1 (7T) and L3 (21T) KNN models with 1 neighbor and Euclidean distance |
+| `L1-L3_K1_Manhattan_Ensemble` | Ensemble of L1 (7T) and L3 (21T) KNN models with 1 neighbor and Manhattan distance |
+| `L1-L3_K3_Euclidean_Ensemble` | Ensemble of L1 (7T) and L3 (21T) KNN models with 3 neighbors and Euclidean distance |
+| `L1-L3_K3_Manhattan_Ensemble` | Ensemble of L1 (7T) and L3 (21T) KNN models with 3 neighbors and Manhattan distance |
+| `Synthetic_K1_Euclidean_Ensemble` | Ensemble of Synthetic 7T, 21T, and SYN KNN models with 1 neighbor and Euclidean distance |
+| `Synthetic_K1_Manhattan_Ensemble` | Ensemble of Synthetic 7T, 21T, and SYN KNN models with 1 neighbor and Manhattan distance |
+| `Synthetic_K3_Euclidean_Ensemble` | Ensemble of Synthetic 7T, 21T, and SYN KNN models with 3 neighbors and Euclidean distance |
+| `Synthetic_K3_Manhattan_Ensemble` | Ensemble of Synthetic 7T, 21T, and SYN KNN models with 3 neighbors and Manhattan distance |
+
+## Decision Tree and Random Forest Models
+
+| Model Name | Description |
+|---|---|
+| `DecisionTree` | Predicts CHONS counts from `mz`, `inv_k0`, and `ccs` |
+| `RandomForest` | Predicts CHONS counts from `mz`, `inv_k0`, and `ccs` |
 
 ## License
 

RandomForest.joblib

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b6b93e850427cebb857a42e7714b25ae672425ee644ffae03a7d3d52b97adcfb
+size 233828641

config.json

@@ -0,0 +1,15 @@
+{
+  "architectures": [
+    "DomMLModel"
+  ],
+  "auto_map": {
+    "AutoConfig": "configuration_dom_ml.DomMLConfig",
+    "AutoModel": "modeling_dom_ml.DomMLModel"
+  },
+  "feature_names": [
+    "mass"
+  ],
+  "library_name": "transformers",
+  "model_name": "Synthetic_K3_Euclidean_Ensemble",
+  "model_type": "dom_ml"
+}

configuration_dom_ml.py

@@ -0,0 +1,19 @@
+from transformers import PretrainedConfig
+
+
+class DomMLConfig(PretrainedConfig):
+    model_type = "dom_ml"
+
+    def __init__(
+        self,
+        model_name="Synthetic_K3_Euclidean_Ensemble",
+        model_file=None,
+        feature_names=None,
+        model_kind=None,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.model_name = model_name
+        self.model_file = model_file
+        self.model_kind = model_kind
+        self.feature_names = feature_names or ["mass"]

modeling_dom_ml.py

@@ -0,0 +1,354 @@
+import numpy as np
+import torch
+from huggingface_hub import hf_hub_download
+from joblib import load
+from transformers import PreTrainedModel
+from transformers.utils import ModelOutput
+import os
+from dataclasses import dataclass
+from typing import Any, Optional
+
+from .configuration_dom_ml import DomMLConfig
+
+
+ELEMENTS = ("C", "H", "O", "N", "S")
+FORMULA_REGRESSOR_NAMES = {"DecisionTree", "RandomForest"}
+
+
+@dataclass
+class DomMLOutput(ModelOutput):
+    predictions: Any = None
+    formula_counts: Any = None
+    formulas: Any = None
+    distances: Optional[Any] = None
+    indices: Optional[Any] = None
+
+
+class DomMLModel(PreTrainedModel):
+    config_class = DomMLConfig
+    base_model_prefix = "dom_ml"
+    main_input_name = "features"
+
+    def __init__(self, config, estimator=None, estimators=None):
+        super().__init__(config)
+        self.estimators = estimators or ([estimator] if estimator is not None else [])
+        self.estimator = self.estimators[0] if self.estimators else None
+
+    @staticmethod
+    def _model_name_to_file(model_name):
+        if not model_name:
+            return None
+        if model_name in FORMULA_REGRESSOR_NAMES:
+            return f"{model_name}.joblib"
+        if model_name.endswith(".joblib"):
+            return model_name
+        if model_name.startswith("knn_model_"):
+            return f"{model_name}.joblib"
+        if model_name.startswith("Model-"):
+            return f"knn_model_{model_name}.joblib"
+        return f"knn_model_Model-{model_name}.joblib"
+
+    @classmethod
+    def _model_name_to_files(cls, model_name):
+        if not model_name:
+            return None
+        if model_name.startswith("L1-L3_") and model_name.endswith("_Ensemble"):
+            base_name = model_name[: -len("_Ensemble")]
+            return [
+                cls._model_name_to_file(f"{base_name}_7T"),
+                cls._model_name_to_file(f"{base_name}_21T"),
+            ]
+        if model_name.startswith("Synthetic_") and model_name.endswith("_Ensemble"):
+            base_name = model_name[: -len("_Ensemble")]
+            return [
+                cls._model_name_to_file(f"{base_name}_7T"),
+                cls._model_name_to_file(f"{base_name}_21T"),
+                cls._model_name_to_file(f"{base_name}_SYN"),
+            ]
+        return [cls._model_name_to_file(model_name)]
+
+    @staticmethod
+    def _infer_model_kind(model_name, model_files):
+        if model_name in FORMULA_REGRESSOR_NAMES:
+            return "formula_regressor"
+        if model_files and all(
+            os.path.basename(model_file) in {"DecisionTree.joblib", "RandomForest.joblib"}
+            for model_file in model_files
+        ):
+            return "formula_regressor"
+        return "knn"
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        config = kwargs.pop("config", None)
+        model_name = kwargs.pop("model_name", None)
+        model_file = kwargs.pop("model_file", None)
+        cache_dir = kwargs.pop("cache_dir", None)
+        force_download = kwargs.pop("force_download", False)
+        local_files_only = kwargs.pop("local_files_only", False)
+        token = kwargs.pop("token", None)
+        revision = kwargs.pop("revision", None)
+        subfolder = kwargs.pop("subfolder", "")
+
+        kwargs.pop("trust_remote_code", None)
+        kwargs.pop("code_revision", None)
+        kwargs.pop("_commit_hash", None)
+
+        if config is None:
+            config = DomMLConfig.from_pretrained(
+                pretrained_model_name_or_path,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                local_files_only=local_files_only,
+                token=token,
+                revision=revision,
+                subfolder=subfolder,
+            )
+
+        model_files = None
+        if model_file is not None:
+            model_files = [model_file]
+        else:
+            model_name = model_name or getattr(config, "model_name", None)
+            model_files = cls._model_name_to_files(model_name)
+        if model_files is None:
+            configured_model_file = getattr(config, "model_file", None)
+            if configured_model_file is not None:
+                model_files = [configured_model_file]
+        if model_files is None:
+            raise ValueError("Pass model_name=... to select one of the available models.")
+
+        config.model_name = model_name
+        config.model_file = model_files[0]
+        config.model_files = model_files
+        config.model_kind = cls._infer_model_kind(model_name, model_files)
+        if config.model_kind == "formula_regressor":
+            config.feature_names = ["mz", "inv_k0", "ccs"]
+
+        model_paths = []
+        for current_model_file in model_files:
+            if os.path.isdir(pretrained_model_name_or_path):
+                model_path = os.path.join(
+                    pretrained_model_name_or_path,
+                    subfolder,
+                    current_model_file,
+                )
+            else:
+                model_path = hf_hub_download(
+                    repo_id=pretrained_model_name_or_path,
+                    filename=current_model_file,
+                    cache_dir=cache_dir,
+                    force_download=force_download,
+                    local_files_only=local_files_only,
+                    token=token,
+                    revision=revision,
+                    subfolder=subfolder,
+                )
+            model_paths.append(model_path)
+
+        estimators = [load(model_path) for model_path in model_paths]
+        model = cls(config, estimators=estimators)
+        model.eval()
+        return model
+
+    def _to_numpy(self, features):
+        if isinstance(features, torch.Tensor):
+            values = features.detach().cpu().numpy()
+        elif hasattr(features, "loc") and hasattr(features, "columns"):
+            feature_names = getattr(self.config, "feature_names", None)
+            if feature_names and all(name in features.columns for name in feature_names):
+                values = features.loc[:, feature_names].to_numpy(dtype=np.float64)
+            else:
+                values = features.to_numpy(dtype=np.float64)
+        elif hasattr(features, "to_numpy"):
+            values = features.to_numpy(dtype=np.float64)
+        else:
+            values = np.asarray(features, dtype=np.float64)
+
+        if getattr(self.config, "model_kind", None) == "formula_regressor":
+            if values.ndim == 1:
+                if values.size != 3:
+                    raise ValueError(
+                        "DecisionTree and RandomForest inputs must use [mz, inv_k0, ccs]."
+                    )
+                return values.reshape(1, 3)
+            return values
+
+        if values.ndim == 0:
+            return values.reshape(1, 1)
+        if values.ndim == 1:
+            return values.reshape(-1, 1)
+        return values
+
+    @staticmethod
+    def _maybe_tensor(array, as_numpy, dtype=None):
+        if as_numpy or array is None:
+            return array
+        return torch.as_tensor(array, dtype=dtype)
+
+    @staticmethod
+    def _counts_to_formulas(counts):
+        formulas = []
+        for row in np.asarray(counts, dtype=int):
+            formula = ""
+            for element, count in zip(ELEMENTS, row):
+                if count > 0:
+                    formula += element
+                    if count != 1:
+                        formula += str(int(count))
+            formulas.append(formula)
+        return np.asarray(formulas)
+
+    def predict_counts(self, features, as_numpy=True):
+        if not self.estimators:
+            raise ValueError("No model is loaded.")
+        if getattr(self.config, "model_kind", None) != "formula_regressor":
+            raise ValueError("predict_counts is only available for DecisionTree and RandomForest.")
+
+        values = self._to_numpy(features)
+        raw_counts = self.estimators[0].predict(values)
+        counts = np.rint(raw_counts).astype(int)
+        counts = np.clip(counts, 0, None)
+        if as_numpy:
+            return counts
+        return torch.as_tensor(counts, dtype=torch.long)
+
+    def predict(self, features):
+        if not self.estimators:
+            raise ValueError("No model is loaded.")
+        if getattr(self.config, "model_kind", None) == "formula_regressor":
+            return self._counts_to_formulas(self.predict_counts(features))
+
+        values = self._to_numpy(features)
+        predictions = [model.predict(values) for model in self.estimators]
+        if len(predictions) == 1:
+            return predictions[0]
+
+        stacked = np.vstack(predictions).T
+        voted_predictions = []
+        for row in stacked:
+            counts = {}
+            for prediction in row:
+                counts[prediction] = counts.get(prediction, 0) + 1
+            voted_predictions.append(
+                max(row, key=lambda prediction: counts[prediction])
+            )
+        return np.asarray(voted_predictions)
+
+    def joblib_summary(self, max_items=5):
+        if not self.estimators:
+            raise ValueError("No model is loaded.")
+
+        summaries = []
+        model_files = getattr(self.config, "model_files", None) or [self.config.model_file]
+        for model_file, estimator in zip(model_files, self.estimators):
+            summary = {
+                "type": type(estimator).__name__,
+                "model_file": model_file,
+                "model_kind": getattr(self.config, "model_kind", None),
+                "feature_names": getattr(self.config, "feature_names", None),
+                "n_features_in": getattr(estimator, "n_features_in_", None),
+                "n_samples_fit": getattr(estimator, "n_samples_fit_", None),
+            }
+            if getattr(self.config, "model_kind", None) == "formula_regressor":
+                summary["output_elements"] = list(ELEMENTS)
+
+            classes = getattr(estimator, "classes_", None)
+            if classes is not None:
+                summary["classes_preview"] = classes[:max_items].tolist()
+
+            fit_x = getattr(estimator, "_fit_X", None)
+            if fit_x is not None:
+                summary["fit_masses_preview"] = np.asarray(fit_x[:max_items]).reshape(-1).tolist()
+
+            summaries.append(summary)
+
+        if len(summaries) == 1:
+            return summaries[0]
+        return {
+            "type": "Ensemble",
+            "model_name": self.config.model_name,
+            "models": summaries,
+        }
+
+    def kneighbors(self, features, n_neighbors=None, as_numpy=False):
+        if getattr(self.config, "model_kind", None) == "formula_regressor":
+            raise ValueError("Nearest-neighbor lookup is only available for KNN models.")
+        if not self.estimators:
+            raise ValueError("No KNN model is loaded.")
+        values = self._to_numpy(features)
+        if len(self.estimators) > 1:
+            model_files = getattr(self.config, "model_files", None) or []
+            distances = {}
+            indices = {}
+            for model_file, estimator in zip(model_files, self.estimators):
+                model_distances, model_indices = estimator.kneighbors(
+                    values,
+                    n_neighbors=n_neighbors,
+                )
+                distances[model_file] = self._maybe_tensor(
+                    model_distances,
+                    as_numpy,
+                    dtype=torch.float32,
+                )
+                indices[model_file] = self._maybe_tensor(
+                    model_indices,
+                    as_numpy,
+                    dtype=torch.long,
+                )
+            return distances, indices
+
+        distances, indices = self.estimators[0].kneighbors(
+            values,
+            n_neighbors=n_neighbors,
+        )
+        return (
+            self._maybe_tensor(distances, as_numpy, dtype=torch.float32),
+            self._maybe_tensor(indices, as_numpy, dtype=torch.long),
+        )
+
+    def neighbor_indices(self, features, n_neighbors=None, as_numpy=False):
+        _, indices = self.kneighbors(
+            features,
+            n_neighbors=n_neighbors,
+            as_numpy=as_numpy,
+        )
+        return indices
+
+    def forward(
+        self,
+        features=None,
+        input_features=None,
+        return_neighbors=False,
+        n_neighbors=None,
+        as_numpy=False,
+        **kwargs,
+    ):
+        if features is None:
+            features = input_features
+        if features is None:
+            raise ValueError("Pass model inputs with features=... or input_features=....")
+
+        predictions = self.predict(features)
+        formula_counts = None
+        formulas = None
+        distances = None
+        indices = None
+
+        if getattr(self.config, "model_kind", None) == "formula_regressor":
+            formula_counts = self.predict_counts(features, as_numpy=as_numpy)
+            formulas = predictions
+        elif return_neighbors:
+            distances, indices = self.kneighbors(
+                features,
+                n_neighbors=n_neighbors,
+                as_numpy=as_numpy,
+            )
+
+        return DomMLOutput(
+            predictions=predictions,
+            formula_counts=formula_counts,
+            formulas=formulas,
+            distances=distances,
+            indices=indices,
+        )

requirements.txt

@@ -0,0 +1,7 @@
+numpy==2.4.1
+scikit-learn==1.8.0
+joblib
+huggingface_hub
+transformers
+torch
+pandas