diff --git a/examples/tutorials/curation/plot_1_automated_curation.py b/examples/tutorials/curation/plot_1_automated_curation.py index 8f87be9de1..97fc05bc6e 100644 --- a/examples/tutorials/curation/plot_1_automated_curation.py +++ b/examples/tutorials/curation/plot_1_automated_curation.py @@ -43,8 +43,8 @@ # model and some metadata about the model. model, model_info = sc.load_model( - repo_id = "SpikeInterface/toy_tetrode_model", - trusted = ['numpy.dtype'] + repo_id="SpikeInterface/toy_tetrode_model", + trust_model=True ) @@ -87,9 +87,9 @@ # a confidence for each unit contained in the ``sorting_analyzer``. labels = sc.model_based_label_units( - sorting_analyzer = sorting_analyzer, - repo_id = "SpikeInterface/toy_tetrode_model", - trusted = ['numpy.dtype'] + sorting_analyzer=sorting_analyzer, + repo_id="SpikeInterface/toy_tetrode_model", + trust_model=True ) print(labels) diff --git a/src/spikeinterface/curation/model_based_curation.py b/src/spikeinterface/curation/model_based_curation.py index 2c3f43c6dc..69c7f75419 100644 --- a/src/spikeinterface/curation/model_based_curation.py +++ b/src/spikeinterface/curation/model_based_curation.py @@ -473,14 +473,8 @@ def _load_model_from_folder(model_folder=None, model_name=None, trust_model=Fals skops_file = skops_files[0] if trust_model and trusted is None: - try: - model = skio.load(skops_file) - except UntrustedTypesFoundException as e: - exception_msg = str(e) - # the exception message contains the list of untrusted objects. The following - # search assumes it is the only list in the message. - string_list = re.search(r"\[(.*?)\]", exception_msg).group() - trusted = [list_item for list_item in string_list.split("'") if len(list_item) > 2] + untrusted = skio.get_untrusted_types(file=skops_file) + trusted = untrusted model = skio.load(skops_file, trusted=trusted) _patch_sklearn_imputer_compatibility(model) diff --git a/src/spikeinterface/curation/tests/test_model_based_curation.py b/src/spikeinterface/curation/tests/test_model_based_curation.py index a42ba2250d..800817cdda 100644 --- a/src/spikeinterface/curation/tests/test_model_based_curation.py +++ b/src/spikeinterface/curation/tests/test_model_based_curation.py @@ -25,7 +25,10 @@ def model(trained_pipeline_path): It has been trained locally and, when applied to `sorting_analyzer_for_unitrefine_curation` will label its 10 units with the following labels: [1,1,1,1,1,1,1,1,1,1,0,0,0,0,0,0,0,0,0,0].""" - model = load_model(trained_pipeline_path, trusted=["numpy.dtype"]) + model = load_model( + trained_pipeline_path, + trust_model=True, + ) return model @@ -66,7 +69,7 @@ def test_metric_ordering_independence(sorting_analyzer_for_unitrefine_curation, prediction_prob_dataframe_1 = model_based_label_units( sorting_analyzer=sorting_analyzer_for_unitrefine_curation, model_folder=trained_pipeline_path, - trusted=["numpy.dtype"], + trust_model=True, ) sorting_analyzer_for_unitrefine_curation.compute( @@ -76,7 +79,7 @@ def test_metric_ordering_independence(sorting_analyzer_for_unitrefine_curation, prediction_prob_dataframe_2 = model_based_label_units( sorting_analyzer=sorting_analyzer_for_unitrefine_curation, model_folder=trained_pipeline_path, - trusted=["numpy.dtype"], + trust_model=True, ) assert prediction_prob_dataframe_1.equals(prediction_prob_dataframe_2) @@ -194,7 +197,7 @@ def test_exception_raised_when_metric_params_not_equal(sorting_analyzer_for_unit "template_metrics", metric_names=["half_width", "peak_to_trough_duration", "number_of_peaks"] ) - model, model_info = load_model(model_folder=trained_pipeline_path, trusted=["numpy.dtype"]) + model, model_info = load_model(model_folder=trained_pipeline_path, trust_model=True) model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model) # an error should be raised if `enforce_metric_params` is True @@ -215,6 +218,6 @@ def test_exception_raised_when_metric_params_not_equal(sorting_analyzer_for_unit "template_metrics", metric_names=["half_width", "peak_to_trough_duration"] ) - model, model_info = load_model(model_folder=trained_pipeline_path, trusted=["numpy.dtype"]) + model, model_info = load_model(model_folder=trained_pipeline_path, trust_model=True) model_based_classification = ModelBasedClassification(sorting_analyzer_for_unitrefine_curation, model) model_based_classification._check_params_for_classification(enforce_metric_params=True, model_info=model_info) diff --git a/src/spikeinterface/preprocessing/tests/test_decimate.py b/src/spikeinterface/preprocessing/tests/test_decimate.py index 99a9e1128c..05d43a2c3e 100644 --- a/src/spikeinterface/preprocessing/tests/test_decimate.py +++ b/src/spikeinterface/preprocessing/tests/test_decimate.py @@ -126,7 +126,7 @@ def test_decimate_antialias_by_chunks(decimation_factor): for chunk_size in [137, int(decimated_rate * 2)]: rec3 = rec2.save(format="memory", chunk_size=chunk_size, n_jobs=1, progress_bar=False) traces3 = rec3.get_traces() - np.testing.assert_allclose(traces3, traces2, rtol=1e-6, atol=1e-6) + np.testing.assert_allclose(traces3, traces2, rtol=1e-5, atol=1e-5) @pytest.mark.parametrize("decimation_factor", [6, 10])