From bcc2f987a5b9718e5748646d9204b35696bafe5a Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 27 Jul 2026 10:08:58 +0200
Subject: [PATCH 1/5] feat: model_based_curation + unitrefine can run from
 computed metrics

---
 .../curation/model_based_curation.py          | 193 ++++++++++--------
 .../tests/test_model_based_curation.py        |  76 ++++---
 .../tests/test_unitrefine_curation.py         |  22 ++
 .../curation/unitrefine_curation.py           |  34 ++-
 4 files changed, 210 insertions(+), 115 deletions(-)

diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py
index 992f3ecd12..3fab6fbb94 100644
--- a/src/spikeinterface/curation/model_based_curation.py
+++ b/src/spikeinterface/curation/model_based_curation.py
@@ -28,6 +28,8 @@ class ModelBasedClassification:
     ----------
     sorting_analyzer : SortingAnalyzer
         The sorting analyzer object containing the spike sorting data.
+    metrics : pd.DataFrame
+        A DataFrame containing the metrics for the units.
     pipeline : Pipeline
         The pipeline object representing the trained classification model.
 
@@ -37,18 +39,28 @@ class ModelBasedClassification:
         Predicts the labels for the spike sorting data using the trained model.
     """
 
-    def __init__(self, sorting_analyzer: SortingAnalyzer, pipeline):
+    def __init__(
+        self, sorting_analyzer: SortingAnalyzer | None = None, metrics: "pd.DataFrame | None" = None, pipeline=None
+    ):
         from sklearn.pipeline import Pipeline
 
         if not isinstance(pipeline, Pipeline):
             raise ValueError("The `pipeline` must be an instance of sklearn.pipeline.Pipeline")
 
+        if sorting_analyzer is None and metrics is None:
+            raise ValueError("At least one of `sorting_analyzer` or `metrics` must be provided.")
         self.sorting_analyzer = sorting_analyzer
+        self.metrics = metrics
         self.pipeline = pipeline
         self.required_metrics = pipeline.feature_names_in_
 
     def predict_labels(
-        self, label_conversion=None, input_data=None, export_to_phy=False, model_info=None, enforce_metric_params=False
+        self,
+        label_conversion: dict[int, str] | None = None,
+        export_to_phy: bool = False,
+        phy_folder: Path | None = None,
+        model_info: dict | None = None,
+        enforce_metric_params: bool = False,
     ):
         """
         Predicts the labels for the spike sorting data using the trained model.
@@ -61,10 +73,14 @@ def predict_labels(
         label_conversion : dict or None, default: None
             A dictionary for converting the predicted labels (which are integers) to custom labels. If None,
             tries to find in `model_info` file. The dictionary should have the format {old_label: new_label}.
-        input_data : pandas.DataFrame or None, default: None
-            The input data for classification. If not provided, the method will extract metrics stored in the sorting analyzer.
         export_to_phy : bool, default: False.
             Whether to export the classified units to Phy format. Default is False.
+        phy_folder : Path or None, default: None
+            The path to the Phy folder where the classified units will be exported. If None,
+            the Phy folder will be inferred from the sorting object. If the sorting object does not have a Phy folder,
+            the esport will be skipped.
+        model_info : dict or None, default: None
+            Dictionary of model info containing provenance of the model.
         enforce_metric_params : bool, default: False
             If True and the parameters used to compute the metrics in `sorting_analyzer` are different than the parmeters
             used to compute the metrics used to train the model, this function will raise an error. Otherwise, a warning is raised.
@@ -78,16 +94,19 @@ def predict_labels(
         import pandas as pd
 
         # Get metrics DataFrame for classification
-        if input_data is None:
-            input_data = self.sorting_analyzer.get_metrics_extension_data()
+        if self.metrics is None:
+            metrics = self.sorting_analyzer.get_metrics_extension_data()
+            unit_ids = self.sorting_analyzer.unit_ids
         else:
-            if not isinstance(input_data, pd.DataFrame):
+            metrics = self.metrics
+            if not isinstance(metrics, pd.DataFrame):
                 raise ValueError("Input data must be a pandas DataFrame")
+            unit_ids = metrics.index.to_list()
 
-        input_data = self.handle_backwards_compatibility_in_metrics(input_data, model_info=model_info)
-        input_data = self._check_required_metrics_are_present(input_data)
+        metrics = _handle_backwards_compatibility_in_metrics(metrics, model_info=model_info)
+        metrics = _check_required_metrics_are_present(self.required_metrics, metrics)
 
-        if model_info is not None:
+        if model_info is not None and self.sorting_analyzer is not None:
             self._check_params_for_classification(enforce_metric_params, model_info=model_info)
 
         if model_info is not None and label_conversion is None:
@@ -100,11 +119,11 @@ def predict_labels(
             except:
                 warnings.warn("Could not find `label_conversion` key in `model_info.json` file")
 
-        input_data = _format_metric_dataframe(input_data)
+        metrics = _format_metric_dataframe(metrics)
 
         # Apply classifier
-        predictions = self.pipeline.predict(input_data)
-        probabilities = self.pipeline.predict_proba(input_data)
+        predictions = self.pipeline.predict(metrics)
+        probabilities = self.pipeline.predict_proba(metrics)
         probabilities = np.max(probabilities, axis=1)
 
         if isinstance(label_conversion, dict):
@@ -114,68 +133,21 @@ def predict_labels(
             predictions = [label_conversion[label] for label in predictions]
 
         classified_units = pd.DataFrame(
-            zip(predictions, probabilities), columns=["prediction", "probability"], index=self.sorting_analyzer.unit_ids
+            zip(predictions, probabilities), columns=["prediction", "probability"], index=unit_ids
         )
 
         # Set predictions and probability as sorting properties
-        self.sorting_analyzer.set_sorting_property("classifier_label", predictions)
-        self.sorting_analyzer.set_sorting_property("classifier_probability", probabilities)
+        if self.sorting_analyzer is not None:
+            self.sorting_analyzer.set_sorting_property("classifier_label", predictions)
+            self.sorting_analyzer.set_sorting_property("classifier_probability", probabilities)
 
         if export_to_phy:
-            self._export_to_phy(classified_units)
-
-        return classified_units
-
-    def handle_backwards_compatibility_in_metrics(self, calculated_metrics, model_info):
-        """
-        Handles backwards compatibility in metric names for models trained with older versions of SpikeInterface.
-        In recent versions, some metric names have been changed for clarity. In addition, the sign of some metrics
-        has been inverted to maintain consistency.
-
-        Parameters
-        ----------
-        calculated_metrics : pd.DataFrame
-            The DataFrame containing the calculated metrics.
-        model_info : dict or None
-            Dictionary of model info containing provenance of the model.
 
-        Returns
-        -------
-        pd.DataFrame
-            The DataFrame with updated metric names for compatibility.
-        """
-        if model_info is None:
-            return calculated_metrics
-        si_version = model_info["requirements"].get("spikeinterface", None)
-        if si_version is not None and parse(si_version) < parse("0.103.2"):
-            # if the model was trained with SI version < 0.103.2, we need to rename some metrics
-            calculated_metrics = calculated_metrics.copy()
-            # peak_to_trough_duration was named peak_to_valley
-            if "peak_to_trough_duration" in calculated_metrics.columns:
-                calculated_metrics = calculated_metrics.rename(columns={"peak_to_trough_duration": "peak_to_valley"})
-            # peak_after_to_trough_ratio was named peak_trough_ratio and had inverted sign
-            if "peak_after_to_trough_ratio" in calculated_metrics.columns:
-                calculated_metrics = calculated_metrics.rename(
-                    columns={"peak_after_to_trough_ratio": "peak_trough_ratio"}
-                )
-                calculated_metrics["peak_trough_ratio"] = -1 * calculated_metrics["peak_trough_ratio"]
-            # trough_half_width was named half_width
-            if "trough_half_width" in calculated_metrics.columns:
-                calculated_metrics = calculated_metrics.rename(columns={"trough_half_width": "half_width"})
-        return calculated_metrics
+            if phy_folder is None:
+                raise ValueError("Phy folder must be provided using the `phy_folder` parameter.")
+            classified_units.to_csv(f"{phy_folder}/cluster_prediction.tsv", sep="\t", index_label="cluster_id")
 
-    def _check_required_metrics_are_present(self, calculated_metrics):
-        # Check all the required metrics have been calculated
-        required_metrics = set(self.required_metrics)
-        if required_metrics.issubset(set(calculated_metrics)):
-            input_data = calculated_metrics[self.required_metrics]
-        else:
-            raise ValueError(
-                "Input data does not contain all required metrics for classification",
-                f"Missing metrics: {required_metrics.difference(calculated_metrics)}",
-            )
-
-        return input_data
+        return classified_units
 
     def _check_params_for_classification(self, enforce_metric_params=False, model_info=None):
         """
@@ -226,22 +198,10 @@ def _check_params_for_classification(self, enforce_metric_params=False, model_in
                     else:
                         warnings.warn(warning_message)
 
-    def _export_to_phy(self, classified_df):
-        """Export the classified units to Phy as cluster_prediction.tsv file"""
-
-        # Export to Phy format
-        try:
-            sorting_path = self.sorting_analyzer.sorting.get_annotation("phy_folder")
-            assert sorting_path is not None
-            assert Path(sorting_path).is_dir()
-        except AssertionError:
-            raise ValueError("Phy folder not found in sorting annotations, or is not a directory")
-
-        classified_df.to_csv(f"{sorting_path}/cluster_prediction.tsv", sep="\t", index_label="cluster_id")
-
 
 def model_based_label_units(
-    sorting_analyzer: SortingAnalyzer,
+    sorting_analyzer: SortingAnalyzer | None,
+    metrics=None,
     model_folder=None,
     repo_id=None,
     model_name=None,
@@ -261,15 +221,17 @@ def model_based_label_units(
 
     Parameters
     ----------
-    sorting_analyzer : SortingAnalyzer
+    sorting_analyzer : SortingAnalyzer | None
         The sorting analyzer object containing the spike sorting results.
+    metrics : pd.DataFrame | None, default: None
+        A DataFrame with metrics for the units. If None, metrics will be computed from the sorting_analyzer.
     model_folder : str or Path, default: None
         The path to the folder containing the model
     repo_id : str, default: None
         Hugging face repo id which contains the model e.g. 'username/model'
     model_name: str, default: None
         Filename of model e.g. 'my_model.skops'. If None, uses first model found.
-    label_conversion : dic | None, default: None
+    label_conversion : dict | None, default: None
         A dictionary for converting the predicted labels (which are integers) to custom labels. If None,
         tries to extract from `model_info.json` file. The dictionary should have the format {old_label: new_label}.
     export_to_phy : bool, default: False
@@ -305,7 +267,9 @@ def model_based_label_units(
     if not isinstance(model, Pipeline):
         raise ValueError("The model must be an instance of sklearn.pipeline.Pipeline")
 
-    model_based_classification = ModelBasedClassification(sorting_analyzer, model)
+    model_based_classification = ModelBasedClassification(
+        sorting_analyzer=sorting_analyzer, metrics=metrics, pipeline=model
+    )
 
     classified_units = model_based_classification.predict_labels(
         label_conversion=label_conversion,
@@ -452,13 +416,16 @@ def _load_model_from_folder(model_folder=None, model_name=None, trust_model=Fals
     else:
         model_info = json.load(open(model_info_path))
 
-    model_info = handle_backwards_compatibility_metric_params(model_info)
+    model_info = _handle_backwards_compatibility_metric_params(model_info)
 
     return model, model_info
 
 
-def handle_backwards_compatibility_metric_params(model_info):
-
+def _handle_backwards_compatibility_metric_params(model_info):
+    """
+    Handles backwards compatibility in metric parameters for models trained with older versions of SpikeInterface.
+    In recent versions, some metric parameters have been changed for clarity.
+    """
     if (
         model_info.get("metric_params") is not None
         and model_info.get("metric_params").get("quality_metric_params") is not None
@@ -479,3 +446,53 @@ def handle_backwards_compatibility_metric_params(model_info):
             del model_info["metric_params"]["template_metric_params"]["metrics_kwargs"]
 
     return model_info
+
+
+def _handle_backwards_compatibility_in_metrics(calculated_metrics, model_info):
+    """
+    Handles backwards compatibility in metric names for models trained with older versions of SpikeInterface.
+    In recent versions, some metric names have been changed for clarity. In addition, the sign of some metrics
+    has been inverted to maintain consistency.
+
+    Parameters
+    ----------
+    calculated_metrics : pd.DataFrame
+        The DataFrame containing the calculated metrics.
+    model_info : dict or None
+        Dictionary of model info containing provenance of the model.
+
+    Returns
+    -------
+    pd.DataFrame
+        The DataFrame with updated metric names for compatibility.
+    """
+    if model_info is None:
+        return calculated_metrics
+    si_version = model_info["requirements"].get("spikeinterface", None)
+    if si_version is not None and parse(si_version) < parse("0.103.2"):
+        # if the model was trained with SI version < 0.103.2, we need to rename some metrics
+        calculated_metrics = calculated_metrics.copy()
+        # peak_to_trough_duration was named peak_to_valley
+        if "peak_to_trough_duration" in calculated_metrics.columns:
+            calculated_metrics = calculated_metrics.rename(columns={"peak_to_trough_duration": "peak_to_valley"})
+        # peak_after_to_trough_ratio was named peak_trough_ratio and had inverted sign
+        if "peak_after_to_trough_ratio" in calculated_metrics.columns:
+            calculated_metrics = calculated_metrics.rename(columns={"peak_after_to_trough_ratio": "peak_trough_ratio"})
+            calculated_metrics["peak_trough_ratio"] = -1 * calculated_metrics["peak_trough_ratio"]
+        # trough_half_width was named half_width
+        if "trough_half_width" in calculated_metrics.columns:
+            calculated_metrics = calculated_metrics.rename(columns={"trough_half_width": "half_width"})
+    return calculated_metrics
+
+
+def _check_required_metrics_are_present(required_metrics, calculated_metrics):
+    # Check all the required metrics have been calculated, preserving the order expected by the pipeline
+    if set(required_metrics).issubset(set(calculated_metrics.columns)):
+        input_data = calculated_metrics[list(required_metrics)]
+    else:
+        raise ValueError(
+            "Input data does not contain all required metrics for classification",
+            f"Missing metrics: {set(required_metrics).difference(calculated_metrics.columns)}",
+        )
+
+    return input_data
diff --git a/src/spikeinterface/curation/tests/test_model_based_curation.py b/src/spikeinterface/curation/tests/test_model_based_curation.py
index 94b98418bb..1b96c840f2 100644
--- a/src/spikeinterface/curation/tests/test_model_based_curation.py
+++ b/src/spikeinterface/curation/tests/test_model_based_curation.py
@@ -33,7 +33,9 @@ def required_metrics_and_columns():
 def test_model_based_classification_init(sorting_analyzer_for_unitrefine_curation, model):
     """Test that the ModelBasedClassification attributes are correctly initialised"""
 
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0])
+    model_based_classification = ModelBasedClassification(
+        sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0]
+    )
     assert model_based_classification.sorting_analyzer == sorting_analyzer_for_unitrefine_curation
     assert model_based_classification.pipeline == model[0]
     assert np.all(model_based_classification.required_metrics == model_based_classification.pipeline.feature_names_in_)
@@ -74,19 +76,22 @@ def test_model_based_classification_get_metrics_for_classification(
     This test checks that an error occurs when the required metrics have not been computed,
     and that no error is returned when the required metrics have been computed.
     """
+    from spikeinterface.curation.model_based_curation import _check_required_metrics_are_present
 
     sorting_analyzer_for_unitrefine_curation.delete_extension("quality_metrics")
     sorting_analyzer_for_unitrefine_curation.delete_extension("template_metrics")
 
     required_metric_names, required_metric_columns = required_metrics_and_columns
 
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0])
+    model_based_classification = ModelBasedClassification(
+        sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0]
+    )
 
     # Compute some (but not all) of the required metrics in sorting_analyzer, should still error
     sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=[required_metric_names[0]])
     computed_metrics = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data()
     with pytest.raises(ValueError):
-        model_based_classification._check_required_metrics_are_present(computed_metrics)
+        _check_required_metrics_are_present(required_metric_names, computed_metrics)
 
     # Compute all of the required metrics in sorting_analyzer, no more error
     sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=required_metric_names[0:2])
@@ -97,25 +102,6 @@ def test_model_based_classification_get_metrics_for_classification(
     assert set(metrics_data.columns.to_list()) == set(required_metric_columns)
 
 
-def test_model_based_classification_export_to_phy(sorting_analyzer_for_unitrefine_curation, model):
-    import pandas as pd
-
-    # Test the _export_to_phy() method of ModelBasedClassification
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0])
-
-    classified_units = pd.DataFrame.from_dict({0: (1, 0.5), 1: (0, 0.5), 2: (1, 0.5), 3: (0, 0.5), 4: (1, 0.5)})
-    # Function should fail here
-    with pytest.raises(ValueError):
-        model_based_classification._export_to_phy(classified_units)
-    # Make temp output folder and set as phy_folder
-    phy_folder = cache_folder / "phy_folder"
-    phy_folder.mkdir(parents=True, exist_ok=True)
-
-    model_based_classification.sorting_analyzer.sorting.annotate(phy_folder=phy_folder)
-    model_based_classification._export_to_phy(classified_units)
-    assert (phy_folder / "cluster_prediction.tsv").exists()
-
-
 def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefine_curation, model):
     """The model `model` has been trained on the `sorting_analyzer` used in this test with
     the labels `[1, 0, 1, 0, 1]`. Hence if we apply the model to this `sorting_analyzer`
@@ -128,7 +114,9 @@ def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefi
     sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"])
 
     # Test the predict_labels() method of ModelBasedClassification
-    model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model[0])
+    model_based_classification = ModelBasedClassification(
+        sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0]
+    )
     classified_units = model_based_classification.predict_labels()
     predictions = classified_units["prediction"].values
 
@@ -142,6 +130,48 @@ def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefi
     assert np.all(predictions_labelled == expected_result_converted)
 
 
+def test_predict_labels_with_phy_export(sorting_analyzer_for_unitrefine_curation, model):
+    """Test that the predict_labels() method of ModelBasedClassification correctly exports to Phy format when requested."""
+
+    sorting_analyzer_for_unitrefine_curation.compute(
+        "template_metrics", metric_names=["half_width", "peak_to_trough_duration", "number_of_peaks"]
+    )
+    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"])
+
+    phy_folder = cache_folder / "phy_export"
+    phy_folder.mkdir(parents=True, exist_ok=True)
+
+    model_based_classification = ModelBasedClassification(
+        sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0]
+    )
+    classified_units = model_based_classification.predict_labels(export_to_phy=True, phy_folder=phy_folder)
+
+    # Check that the cluster_prediction.tsv file was created in the specified phy_folder
+    assert (phy_folder / "cluster_prediction.tsv").exists()
+
+    # Using export_to_phy=True without providing a phy_folder should raise a ValueError
+    with pytest.raises(ValueError):
+        model_based_classification.predict_labels(export_to_phy=True, phy_folder=None)
+
+
+def test_model_based_classification_from_dataframe(sorting_analyzer_for_unitrefine_curation, model):
+    """Test that the ModelBasedClassification can be initialised from a DataFrame of metrics."""
+
+    sorting_analyzer_for_unitrefine_curation.compute(
+        "template_metrics", metric_names=["half_width", "peak_to_trough_duration", "number_of_peaks"]
+    )
+    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=["num_spikes", "snr"])
+
+    metrics_dataframe = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data()
+
+    model_based_classification = ModelBasedClassification(metrics=metrics_dataframe, pipeline=model[0])
+    classified_units = model_based_classification.predict_labels()
+    predictions = classified_units["prediction"].values
+
+    expected_result = np.array([1] * 6 + [0] * 6)
+    assert np.all(predictions == expected_result)
+
+
 @pytest.mark.skip(reason="We need to retrain the model to reflect any changes in metric computation")
 def test_exception_raised_when_metric_params_not_equal(sorting_analyzer_for_unitrefine_curation, trained_pipeline_path):
     """We track whether the metric parameters used to compute the metrics used to train
diff --git a/src/spikeinterface/curation/tests/test_unitrefine_curation.py b/src/spikeinterface/curation/tests/test_unitrefine_curation.py
index 314fcb5878..e30b016b53 100644
--- a/src/spikeinterface/curation/tests/test_unitrefine_curation.py
+++ b/src/spikeinterface/curation/tests/test_unitrefine_curation.py
@@ -89,3 +89,25 @@ def test_unitrefine_label_units_with_local_models(sorting_analyzer_for_unitrefin
         sorting_analyzer_for_unitrefine_curation,
         noise_neural_classifier=trained_pipeline_path / "best_model.skops",
     )
+
+
+def test_unitrefine_label_units_with_metrics(sorting_analyzer_for_unitrefine_curation):
+    # test passing metrics instead of sorting_analyzer
+    sorting_analyzer_for_unitrefine_curation.compute(
+        {
+            "spike_amplitudes": {},
+            "template_metrics": {"include_multi_channel_metrics": True},
+            "quality_metrics": {},
+        }
+    )
+    metrics = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data()
+
+    labels = unitrefine_label_units(
+        metrics=metrics,
+        noise_neural_classifier="SpikeInterface/UnitRefine_noise_neural_classifier_lightweight",
+        sua_mua_classifier="SpikeInterface/UnitRefine_sua_mua_classifier_lightweight",
+    )
+
+    assert "unitrefine_label" in labels.columns
+    assert "unitrefine_probability" in labels.columns
+    assert labels.shape[0] == len(metrics)
diff --git a/src/spikeinterface/curation/unitrefine_curation.py b/src/spikeinterface/curation/unitrefine_curation.py
index 309e924bb8..2a25b68fcb 100644
--- a/src/spikeinterface/curation/unitrefine_curation.py
+++ b/src/spikeinterface/curation/unitrefine_curation.py
@@ -6,7 +6,8 @@
 
 
 def unitrefine_label_units(
-    sorting_analyzer: SortingAnalyzer,
+    sorting_analyzer: SortingAnalyzer | None = None,
+    metrics: "pd.DataFrame | None" = None,
     noise_neural_classifier: str | Path | None = None,
     sua_mua_classifier: str | Path | None = None,
 ):
@@ -18,8 +19,10 @@ def unitrefine_label_units(
 
     Parameters
     ----------
-    sorting_analyzer : SortingAnalyzer
+    sorting_analyzer : SortingAnalyzer or None, default: None
         The sorting analyzer object containing the spike sorting results.
+    metrics : pd.DataFrame or None, default: None
+        A DataFrame with metrics for the units. If None, metrics will be computed from the sorting_analyzer.
     noise_neural_classifier : str or Path or None, default: None
         The path to the folder containing the model, a full path to a model (".skops")
         or a string to a repo on HuggingFace.
@@ -49,6 +52,18 @@ def unitrefine_label_units(
             "https://huggingface.co/AnoushkaJain3/models. You can also train models on your own data: "
             "see https://github.com/anoushkajain/UnitRefine for more details."
         )
+    if sorting_analyzer is None and metrics is None:
+        raise ValueError(
+            "At least one of sorting_analyzer or metrics must be provided. "
+            "If you have a Sorting object, you can create a SortingAnalyzer object using "
+            "`sorting_analyzer = SortingAnalyzer(sorting, recording)`."
+        )
+    if sorting_analyzer is not None and metrics is not None:
+        raise ValueError(
+            "Only one of sorting_analyzer or metrics should be provided. "
+            "If you have a Sorting object, you can create a SortingAnalyzer object using "
+            "`sorting_analyzer = SortingAnalyzer(sorting, recording)`."
+        )
 
     if noise_neural_classifier is not None:
         # 1. apply the noise/neural classification and remove noise
@@ -56,6 +71,7 @@ def unitrefine_label_units(
             warnings.filterwarnings("ignore", category=InconsistentVersionWarning)
             noise_neuron_labels = model_based_label_units(
                 sorting_analyzer=sorting_analyzer,
+                metrics=metrics,
                 trust_model=True,
                 **get_model_based_classification_kwargs(noise_neural_classifier),
             )
@@ -65,18 +81,28 @@ def unitrefine_label_units(
                 "Please check the model used for classification."
             )
         noise_units = noise_neuron_labels[noise_neuron_labels["prediction"] == "noise"]
-        sorting_analyzer_neural = sorting_analyzer.remove_units(noise_units.index)
+        if sorting_analyzer is not None:
+            sorting_analyzer_neural = sorting_analyzer.remove_units(noise_units.index)
+            metrics_neural = None
+            unit_ids_neural = sorting_analyzer_neural.unit_ids
+        else:
+            metrics_neural = metrics.drop(index=noise_units.index)
+            sorting_analyzer_neural = None
+            unit_ids_neural = metrics_neural.index
     else:
         sorting_analyzer_neural = sorting_analyzer
+        metrics_neural = metrics
         noise_units = pd.DataFrame(columns=["prediction", "probability"])
+        unit_ids_neural = sorting_analyzer.unit_ids if sorting_analyzer is not None else metrics.index
 
     if sua_mua_classifier is not None:
         # 2. apply the sua/mua classification and aggregate results
-        if len(sorting_analyzer.unit_ids) > len(noise_units):
+        if len(unit_ids_neural) > 0:
             with warnings.catch_warnings():
                 warnings.filterwarnings("ignore", category=InconsistentVersionWarning)
                 sua_mua_labels = model_based_label_units(
                     sorting_analyzer=sorting_analyzer_neural,
+                    metrics=metrics_neural,
                     trust_model=True,
                     **get_model_based_classification_kwargs(sua_mua_classifier),
                 )

From 8120fab7b965775d419ae347e59a4153617c8fc9 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 27 Jul 2026 11:57:46 +0200
Subject: [PATCH 2/5] Add get_required_metrics_from_model helper

---
 doc/api.rst                                   |  2 +
 src/spikeinterface/curation/__init__.py       |  7 ++-
 .../curation/model_based_curation.py          | 49 +++++++++++++----
 .../tests/test_model_based_curation.py        | 52 ++++++++++++++-----
 4 files changed, 85 insertions(+), 25 deletions(-)

diff --git a/doc/api.rst b/doc/api.rst
index ce850d1291..ecc209a9cc 100755
--- a/doc/api.rst
+++ b/doc/api.rst
@@ -411,6 +411,8 @@ spikeinterface.curation
     .. autofunction:: bombcell_label_units
     .. autofunction:: bombcell_get_default_thresholds
     .. autofunction:: model_based_label_units
+    .. autofunction:: get_required_metrics_from_model
+    .. autofunction:: check_required_metrics_are_present
     .. autofunction:: load_model
     .. autofunction:: train_model
     .. autofunction:: unitrefine_label_units
diff --git a/src/spikeinterface/curation/__init__.py b/src/spikeinterface/curation/__init__.py
index 16bdddd870..9d3405ef62 100644
--- a/src/spikeinterface/curation/__init__.py
+++ b/src/spikeinterface/curation/__init__.py
@@ -22,7 +22,12 @@
 # automated curation
 from .curation_tools import get_labeling_summary
 from .threshold_metrics_curation import threshold_metrics_label_units
-from .model_based_curation import model_based_label_units, load_model, auto_label_units
+from .model_based_curation import (
+    model_based_label_units,
+    load_model,
+    get_required_metrics_from_model,
+    check_required_metrics_are_present,
+)
 from .train_manual_curation import train_model, get_default_classifier_search_spaces
 from .unitrefine_curation import unitrefine_label_units
 from .bombcell_curation import (
diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py
index 3fab6fbb94..04e2fe9a1d 100644
--- a/src/spikeinterface/curation/model_based_curation.py
+++ b/src/spikeinterface/curation/model_based_curation.py
@@ -104,7 +104,7 @@ def predict_labels(
             unit_ids = metrics.index.to_list()
 
         metrics = _handle_backwards_compatibility_in_metrics(metrics, model_info=model_info)
-        metrics = _check_required_metrics_are_present(self.required_metrics, metrics)
+        metrics = check_required_metrics_are_present(self.required_metrics, metrics)
 
         if model_info is not None and self.sorting_analyzer is not None:
             self._check_params_for_classification(enforce_metric_params, model_info=model_info)
@@ -281,17 +281,44 @@ def model_based_label_units(
     return classified_units
 
 
-def auto_label_units(*args, **kwargs):
+def get_required_metrics_from_model(
+    model_folder=None, repo_id=None, model_name=None, model=None, trust_model=False, trusted=None
+):
     """
-    Deprecated function. Please use `model_based_label_units` instead.
+    Returns the required metrics for a model, either from a model hosted on HuggingFaceHub or one available in a local folder.
+
+    Parameters
+    ----------
+    model_folder : str or Path, default: None
+        The path to the folder containing the model
+    repo_id : str, default: None
+        Hugging face repo id which contains the model e.g. 'username/model'
+    model_name: str, default: None
+        Filename of model e.g. 'my_model.skops'. If None, uses first model found.
+    model : sklearn.pipeline.Pipeline, default: None
+        A trained sklearn pipeline model. If provided, the required metrics will be extracted from this model
+    trust_model : bool, default: False
+        Whether to trust the model. If True, the `trusted` parameter that is passed to `skops.load` to load the model will be
+        automatically inferred. If False, the `trusted` parameter must be provided to indicate the trusted objects.
+    trusted : list of str, default: None
+        Passed to skops.load. The object will be loaded only if there are only trusted objects and objects of types listed in trusted in the dumped file.
+
+    Returns
+    -------
+    required_metrics : list of str
+        A list of required metrics for the model.
     """
-    warnings.warn(
-        "`auto_label_units` is deprecated and will be removed in v0.105.0. "
-        "Please use `model_based_label_units` instead.",
-        FutureWarning,
-        stacklevel=2,
-    )
-    return model_based_label_units(*args, **kwargs)
+    from sklearn.pipeline import Pipeline
+
+    if model is None:
+        model, _ = load_model(
+            model_folder=model_folder, repo_id=repo_id, model_name=model_name, trust_model=trust_model, trusted=trusted
+        )
+
+    if not isinstance(model, Pipeline):
+        raise ValueError("The model must be an instance of sklearn.pipeline.Pipeline")
+
+    return list(model.feature_names_in_)
 
 
 def load_model(model_folder=None, repo_id=None, model_name=None, trust_model=False, trusted=None):
@@ -485,7 +512,7 @@ def _handle_backwards_compatibility_in_metrics(calculated_metrics, model_info):
     return calculated_metrics
 
 
-def _check_required_metrics_are_present(required_metrics, calculated_metrics):
+def check_required_metrics_are_present(required_metrics, calculated_metrics):
     # Check all the required metrics have been calculated, preserving the order expected by the pipeline
     if set(required_metrics).issubset(set(calculated_metrics.columns)):
         input_data = calculated_metrics[list(required_metrics)]
diff --git a/src/spikeinterface/curation/tests/test_model_based_curation.py b/src/spikeinterface/curation/tests/test_model_based_curation.py
index 1b96c840f2..0d6db7acc6 100644
--- a/src/spikeinterface/curation/tests/test_model_based_curation.py
+++ b/src/spikeinterface/curation/tests/test_model_based_curation.py
@@ -3,7 +3,12 @@
 
 from spikeinterface.curation.tests.common import sorting_analyzer_for_unitrefine_curation, trained_pipeline_path
 from spikeinterface.curation.model_based_curation import ModelBasedClassification
-from spikeinterface.curation import model_based_label_units, load_model
+from spikeinterface.curation import (
+    model_based_label_units,
+    load_model,
+    get_required_metrics_from_model,
+    check_required_metrics_are_present,
+)
 
 
 import numpy as np
@@ -25,9 +30,17 @@ def model(trained_pipeline_path):
 
 
 @pytest.fixture
-def required_metrics_and_columns():
+def required_metrics():
     """These are the metrics which `model` are trained on."""
-    return ["num_spikes", "snr", "half_width"], ["num_spikes", "snr", "trough_half_width", "peak_half_width"]
+    from spikeinterface.metrics import ComputeQualityMetrics, ComputeTemplateMetrics
+
+    all_metric_names = ["snr", "half_width", "peak_to_trough_duration", "number_of_peaks"]
+    quality_metric_names = ["snr"]
+    template_metric_names = ["half_width", "peak_to_trough_duration", "number_of_peaks"]
+    all_metric_columns = ComputeQualityMetrics.get_metric_columns(
+        quality_metric_names
+    ) + ComputeTemplateMetrics.get_metric_columns(template_metric_names)
+    return all_metric_names, all_metric_columns, quality_metric_names, template_metric_names
 
 
 def test_model_based_classification_init(sorting_analyzer_for_unitrefine_curation, model):
@@ -70,36 +83,34 @@ def test_metric_ordering_independence(sorting_analyzer_for_unitrefine_curation,
 
 
 def test_model_based_classification_get_metrics_for_classification(
-    sorting_analyzer_for_unitrefine_curation, model, required_metrics_and_columns
+    sorting_analyzer_for_unitrefine_curation, model, required_metrics
 ):
     """If the user has not computed the required metrics, an error should be returned.
     This test checks that an error occurs when the required metrics have not been computed,
     and that no error is returned when the required metrics have been computed.
     """
-    from spikeinterface.curation.model_based_curation import _check_required_metrics_are_present
-
     sorting_analyzer_for_unitrefine_curation.delete_extension("quality_metrics")
     sorting_analyzer_for_unitrefine_curation.delete_extension("template_metrics")
 
-    required_metric_names, required_metric_columns = required_metrics_and_columns
+    all_metric_names, all_metric_columns, qm_names, tm_names = required_metrics
 
     model_based_classification = ModelBasedClassification(
         sorting_analyzer=sorting_analyzer_for_unitrefine_curation, pipeline=model[0]
     )
 
     # Compute some (but not all) of the required metrics in sorting_analyzer, should still error
-    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=[required_metric_names[0]])
+    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=[all_metric_names[0]])
     computed_metrics = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data()
     with pytest.raises(ValueError):
-        _check_required_metrics_are_present(required_metric_names, computed_metrics)
+        check_required_metrics_are_present(all_metric_columns, computed_metrics)
 
     # Compute all of the required metrics in sorting_analyzer, no more error
-    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=required_metric_names[0:2])
-    sorting_analyzer_for_unitrefine_curation.compute("template_metrics", metric_names=[required_metric_names[2]])
+    sorting_analyzer_for_unitrefine_curation.compute("quality_metrics", metric_names=qm_names)
+    sorting_analyzer_for_unitrefine_curation.compute("template_metrics", metric_names=tm_names)
 
     metrics_data = sorting_analyzer_for_unitrefine_curation.get_metrics_extension_data()
-    assert metrics_data.shape[0] == len(sorting_analyzer_for_unitrefine_curation.sorting.get_unit_ids())
-    assert set(metrics_data.columns.to_list()) == set(required_metric_columns)
+    assert len(metrics_data) == len(sorting_analyzer_for_unitrefine_curation.unit_ids)
+    assert set(metrics_data.columns.to_list()) == set(all_metric_columns)
 
 
 def test_model_based_classification_predict_labels(sorting_analyzer_for_unitrefine_curation, model):
@@ -172,6 +183,21 @@ def test_model_based_classification_from_dataframe(sorting_analyzer_for_unitrefi
     assert np.all(predictions == expected_result)
 
 
+def test_get_required_metrics_from_model(model, required_metrics):
+    """Test that the get_required_metrics_from_model function returns the correct required metrics and columns."""
+
+    required_from_model = get_required_metrics_from_model(model=model[0])
+
+    _, all_metric_columns, _, _ = required_metrics
+    assert set(all_metric_columns) == set(required_from_model)
+
+    # from HF
+    required_metrics_from_model_hf = get_required_metrics_from_model(
+        repo_id="SpikeInterface/UnitRefine_sua_mua_classifier", trust_model=True
+    )
+    assert set(all_metric_columns) != set(required_metrics_from_model_hf[0])
+
+
 @pytest.mark.skip(reason="We need to retrain the model to reflect any changes in metric computation")
 def test_exception_raised_when_metric_params_not_equal(sorting_analyzer_for_unitrefine_curation, trained_pipeline_path):
     """We track whether the metric parameters used to compute the metrics used to train

From cde564887fdd72c4669ddef60af6f18097fad43a Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 27 Jul 2026 12:53:36 +0200
Subject: [PATCH 3/5] fix: handle backward compatibility in metric_names

---
 .../curation/model_based_curation.py          | 52 ++++++++++++++-----
 1 file changed, 40 insertions(+), 12 deletions(-)

diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py
index 04e2fe9a1d..3c05c6ca55 100644
--- a/src/spikeinterface/curation/model_based_curation.py
+++ b/src/spikeinterface/curation/model_based_curation.py
@@ -12,6 +12,13 @@
     _format_metric_dataframe,
 )
 
+# Map old metric column names to new metric column names for backwards compatibility
+_BACKWARD_COMPATIBILITY_MAP = {
+    "peak_to_valley": {"name": "peak_to_trough_duration"},
+    "peak_trough_ratio": {"name": "peak_after_to_trough_ratio", "flip_sign": True},
+    "half_width": {"name": "trough_half_width"},
+}
+
 
 class ModelBasedClassification:
     """
@@ -318,7 +325,7 @@ def get_required_metrics_from_model(
     if not isinstance(model, Pipeline):
         raise ValueError("The model must be an instance of sklearn.pipeline.Pipeline")
 
-    return list(model.feature_names_in_)
+    return _handle_backwards_compatibility_in_metric_names(list(model.feature_names_in_))
 
 
 def load_model(model_folder=None, repo_id=None, model_name=None, trust_model=False, trusted=None):
@@ -497,21 +504,42 @@ def _handle_backwards_compatibility_in_metrics(calculated_metrics, model_info):
         return calculated_metrics
     si_version = model_info["requirements"].get("spikeinterface", None)
     if si_version is not None and parse(si_version) < parse("0.103.2"):
-        # if the model was trained with SI version < 0.103.2, we need to rename some metrics
+        # If the model was trained with SI version < 0.103.2, we need to rename some metrics
         calculated_metrics = calculated_metrics.copy()
-        # peak_to_trough_duration was named peak_to_valley
-        if "peak_to_trough_duration" in calculated_metrics.columns:
-            calculated_metrics = calculated_metrics.rename(columns={"peak_to_trough_duration": "peak_to_valley"})
-        # peak_after_to_trough_ratio was named peak_trough_ratio and had inverted sign
-        if "peak_after_to_trough_ratio" in calculated_metrics.columns:
-            calculated_metrics = calculated_metrics.rename(columns={"peak_after_to_trough_ratio": "peak_trough_ratio"})
-            calculated_metrics["peak_trough_ratio"] = -1 * calculated_metrics["peak_trough_ratio"]
-        # trough_half_width was named half_width
-        if "trough_half_width" in calculated_metrics.columns:
-            calculated_metrics = calculated_metrics.rename(columns={"trough_half_width": "half_width"})
+        # We need to rename the metrics and flip sign when needed
+        for old_name, updated_dict in _BACKWARD_COMPATIBILITY_MAP.items():
+            updated_name = updated_dict["name"]
+            if updated_name in calculated_metrics.columns:
+                calculated_metrics = calculated_metrics.rename(columns={updated_name: old_name})
+                if updated_dict.get("flip_sign", False):
+                    calculated_metrics[old_name] = -1 * calculated_metrics[old_name]
     return calculated_metrics
 
 
+def _handle_backwards_compatibility_in_metric_names(model_metric_names):
+    """
+    Handles backwards compatibility in metric names for models trained with older versions of SpikeInterface.
+    In recent versions, some metric names have been changed for clarity.
+
+    Parameters
+    ----------
+    model_metric_names : list of str
+        The list of metric names used in the model.
+
+    Returns
+    -------
+    list of str
+        The list of updated metric names for compatibility.
+    """
+    updated_metric_names = []
+    for metric_name in model_metric_names:
+        if metric_name in _BACKWARD_COMPATIBILITY_MAP:
+            updated_metric_names.append(_BACKWARD_COMPATIBILITY_MAP[metric_name])
+        else:
+            updated_metric_names.append(metric_name)
+    return updated_metric_names
+
+
 def check_required_metrics_are_present(required_metrics, calculated_metrics):
     # Check all the required metrics have been calculated, preserving the order expected by the pipeline
     if set(required_metrics).issubset(set(calculated_metrics.columns)):

From 3ab7c2f9f9f67ac2b3e5ddb70368f68454c52986 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 27 Jul 2026 12:58:41 +0200
Subject: [PATCH 4/5] oups

---
 src/spikeinterface/curation/model_based_curation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py
index 3c05c6ca55..547deb1c6b 100644
--- a/src/spikeinterface/curation/model_based_curation.py
+++ b/src/spikeinterface/curation/model_based_curation.py
@@ -534,7 +534,7 @@ def _handle_backwards_compatibility_in_metric_names(model_metric_names):
     updated_metric_names = []
     for metric_name in model_metric_names:
         if metric_name in _BACKWARD_COMPATIBILITY_MAP:
-            updated_metric_names.append(_BACKWARD_COMPATIBILITY_MAP[metric_name])
+            updated_metric_names.append(_BACKWARD_COMPATIBILITY_MAP[metric_name["name"]])
         else:
             updated_metric_names.append(metric_name)
     return updated_metric_names

From 5b1c73d0ced29b85527b421279e9d151216b9192 Mon Sep 17 00:00:00 2001
From: Alessio Buccino <alejoe9187@gmail.com>
Date: Mon, 27 Jul 2026 13:01:01 +0200
Subject: [PATCH 5/5] oups 2

---
 src/spikeinterface/curation/model_based_curation.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py
index 547deb1c6b..b17656ddbe 100644
--- a/src/spikeinterface/curation/model_based_curation.py
+++ b/src/spikeinterface/curation/model_based_curation.py
@@ -534,7 +534,7 @@ def _handle_backwards_compatibility_in_metric_names(model_metric_names):
     updated_metric_names = []
     for metric_name in model_metric_names:
         if metric_name in _BACKWARD_COMPATIBILITY_MAP:
-            updated_metric_names.append(_BACKWARD_COMPATIBILITY_MAP[metric_name["name"]])
+            updated_metric_names.append(_BACKWARD_COMPATIBILITY_MAP[metric_name]["name"])
         else:
             updated_metric_names.append(metric_name)
     return updated_metric_names