--- a/nn_classifier.xml	Thu Aug 11 09:54:23 2022 +0000
+++ b/nn_classifier.xml	Wed Aug 09 12:57:05 2023 +0000
@@ -1,4 +1,4 @@
-<tool id="sklearn_nn_classifier" name="Nearest Neighbors Classification" version="@VERSION@" profile="20.05">
+<tool id="sklearn_nn_classifier" name="Nearest Neighbors Classification" version="@VERSION@" profile="@PROFILE@">
     <description></description>
     <macros>
         <import>main_macros.xml</import>
@@ -19,9 +19,9 @@
 import numpy as np
 import sklearn.neighbors
 import pandas
-import pickle
 
-from galaxy_ml.utils import load_model, get_X_y
+from galaxy_ml.model_persist import dump_model_to_h5, load_model_from_h5
+from galaxy_ml.utils import clean_params, get_X_y
 
 
 input_json_path = sys.argv[1]
@@ -30,8 +30,8 @@
 
 #if $selected_tasks.selected_task == "load":
 
-with open("$infile_model", 'rb') as model_handler:
-    classifier_object = load_model(model_handler)
+classifier_object = load_model_from_h5('$infile_model')
+classifier_object = clean_params(classifier_object)
 
 header = 'infer' if params["selected_tasks"]["header"] else None
 data = pandas.read_csv("$selected_tasks.infile_data", sep='\t', header=header, index_col=None, parse_dates=True, encoding=None)
@@ -59,8 +59,7 @@
 classifier_object = my_class(**options)
 classifier_object.fit(X, y)
 
-with open("$outfile_fit", 'wb') as out_handler:
-    pickle.dump(classifier_object, out_handler)
+dump_model_to_h5(classifier_object, '$outfile_fit')
 
 #end if
 
@@ -68,7 +67,7 @@
         </configfile>
     </configfiles>
     <inputs>
-        <expand macro="sl_Conditional" model="zip">            <!--Todo: add sparse to targets-->
+        <expand macro="sl_Conditional" model="h5mlm">            <!--Todo: add sparse to targets-->
             <param name="selected_algorithm" type="select" label="Classifier type">
                 <option value="nneighbors">Nearest Neighbors</option>
                 <option value="ncentroid">Nearest Centroid</option>
@@ -128,6 +127,7 @@
             <param name="selected_task" value="train" />
             <param name="selected_algorithm" value="nneighbors" />
             <param name="sampling_method" value="RadiusNeighborsClassifier" />
+            <param name="algorithm" value="brute" />
             <output name="outfile_fit" file="nn_model02" compare="sim_size" delta="5" />
         </test>
         <test>
@@ -142,21 +142,21 @@
             <output name="outfile_fit" file="nn_model03" compare="sim_size" delta="5" />
         </test>
         <test>
-            <param name="infile_model" value="nn_model01" ftype="zip" />
+            <param name="infile_model" value="nn_model01" ftype="h5mlm" />
             <param name="infile_data" value="test_set.tabular" ftype="tabular" />
             <param name="header" value="True" />
             <param name="selected_task" value="load" />
             <output name="outfile_predict" file="nn_prediction_result01.tabular" />
         </test>
         <test>
-            <param name="infile_model" value="nn_model02" ftype="zip" />
+            <param name="infile_model" value="nn_model02" ftype="h5mlm" />
             <param name="infile_data" value="test_set.tabular" ftype="tabular" />
             <param name="header" value="True" />
             <param name="selected_task" value="load" />
             <output name="outfile_predict" file="nn_prediction_result02.tabular" />
         </test>
         <test>
-            <param name="infile_model" value="nn_model03" ftype="zip" />
+            <param name="infile_model" value="nn_model03" ftype="h5mlm" />
             <param name="infile_data" value="test_set.tabular" ftype="tabular" />
             <param name="header" value="True" />
             <param name="selected_task" value="load" />