SparseZoo V2 Bug Fixes

src/sparseml/optim/helpers.py

@@ -21,7 +21,7 @@
 import re
 from contextlib import suppress
 from copy import deepcopy
-from typing import Any, Dict, Optional, Tuple, Union
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import yaml
 
@@ -401,7 +401,10 @@ def _load_yaml_str_from_file(file_path: Union[str, File]) -> str:
     if file_path.startswith("zoo:"):
         # download from zoo stub
         model = Model(file_path)
-        file_path = model.recipes.default.path
+        recipe = model.recipes
+        if isinstance(recipe, List):
+            recipe = recipe[0]
+        file_path = recipe.path
 
     # load the yaml string
     if "\n" in file_path or "\r" in file_path:

src/sparseml/pytorch/image_classification/utils/helpers.py

@@ -91,9 +91,9 @@ def save_zoo_directory(
     """
     for root_file in [
         "model.onnx",
-        "sample_inputs",
-        "sample_outputs",
-        "sample_labels",
+        "sample-inputs",
+        "sample-outputs",
+        "sample-labels",
         "deployment",
     ]:
         root_file_path = os.path.join(training_outputs_dir, root_file)
@@ -109,9 +109,9 @@ def save_zoo_directory(
         training=os.path.join(training_outputs_dir, "training"),
         deployment=os.path.join(training_outputs_dir, "deployment"),
         onnx_model=os.path.join(training_outputs_dir, "model.onnx"),
-        sample_inputs=os.path.join(training_outputs_dir, "sample_inputs"),
-        sample_outputs=os.path.join(training_outputs_dir, "sample_outputs"),
-        sample_labels=os.path.join(training_outputs_dir, "sample_labels"),
+        sample_inputs=os.path.join(training_outputs_dir, "sample-inputs"),
+        sample_outputs=os.path.join(training_outputs_dir, "sample-outputs"),
+        sample_labels=os.path.join(training_outputs_dir, "sample-labels"),
         model_card=os.path.join(training_outputs_dir, "model.md"),
         logs=logs_path,
         sample_originals=None,

src/sparseml/pytorch/utils/exporter.py

@@ -362,9 +362,9 @@ def export_samples(
         :param exp_counter: the counter to start exporting the tensor files at
         """
         sample_batches = [tensors_to_device(batch, "cpu") for batch in sample_batches]
-        inputs_dir = os.path.join(self._output_dir, "sample_inputs")
-        outputs_dir = os.path.join(self._output_dir, "sample_outputs")
-        labels_dir = os.path.join(self._output_dir, "sample_labels")
+        inputs_dir = os.path.join(self._output_dir, "sample-inputs")
+        outputs_dir = os.path.join(self._output_dir, "sample-outputs")
+        labels_dir = os.path.join(self._output_dir, "sample-labels")
         originals_dir = os.path.join(self._output_dir, "sample_originals")
 
         with torch.no_grad():

src/sparseml/pytorch/utils/helpers.py

@@ -1150,9 +1150,10 @@ def download_framework_model_by_recipe_type(
     """
 
     # default to model query params if available
-    recipe_name = recipe_name or (
-        zoo_model.stub_params.get("recipe_type") or zoo_model.stub_params.get("recipe")
-    )
+    model_recipe_name = zoo_model.recipes
+    if model_recipe_name is not None:
+        model_recipe_name = model_recipe_name[0].path
+    recipe_name = recipe_name or (model_recipe_name)
 
     framework_model = None
     if recipe_name and "transfer" in recipe_name.lower():

src/sparseml/transformers/utils/helpers.py

@@ -45,7 +45,7 @@ def save_zoo_directory(
         with the saved training artifacts
     :param logs_path: Optional directory where the training logs reside
     """
-    for root_file in ["sample_inputs", "sample_outputs"]:
+    for root_file in ["sample-inputs", "sample-outputs"]:
         root_file_path = os.path.join(training_outputs_dir, root_file)
         if not os.path.exists(root_file_path):
             logging.warning(
@@ -67,8 +67,8 @@ def save_zoo_directory(
         training=os.path.join(training_outputs_dir, "training"),
         deployment=os.path.join(training_outputs_dir, "deployment"),
         onnx_model=os.path.join(training_outputs_dir, "model.onnx"),
-        sample_inputs=os.path.join(training_outputs_dir, "sample_inputs"),
-        sample_outputs=os.path.join(training_outputs_dir, "sample_outputs"),
+        sample_inputs=os.path.join(training_outputs_dir, "sample-inputs"),
+        sample_outputs=os.path.join(training_outputs_dir, "sample-outputs"),
         model_card=os.path.join(training_outputs_dir, "model.md"),
         logs=logs_path,
         sample_labels=None,

src/sparseml/yolov8/utils/export_samples.py

@@ -78,7 +78,7 @@ def export_sample_inputs_outputs(
     exported_samples = 0
 
     # Sample export directories
-    sample_in_dir = os.path.join(save_dir, "sample_inputs")
+    sample_in_dir = os.path.join(save_dir, "sample-inputs")
     sample_out_dir_torch = os.path.join(save_dir, "sample_outputs_torch")
     sample_out_dir_ort = os.path.join(save_dir, "sample_outputs_onnxruntime")
 

tests/sparseml/optim/test_helpers.py

@@ -626,11 +626,7 @@ def test_evaluate_recipe_yaml_str_equations_invalid(recipe):
         (
             "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/"
             "pruned-conservative"
-        ),
-        (
-            "zoo:cv/classification/resnet_v1-50/pytorch/sparseml/imagenette/"
-            "pruned-conservative?recipe=original"
-        ),
+        )
     ],
 )
 def test_load_recipe_yaml_str_zoo(zoo_path):