[SparseZoo v2 Bridge] Save IC training artifacts to ModelDirectory directory

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

@@ -51,6 +51,7 @@
 )
 from sparseml.utils import create_dirs
 from sparsezoo import Zoo
+from sparsezoo.v2.helpers import setup_model_directory
 
 
 _LOGGER = logging.getLogger(__name__)
@@ -67,9 +68,52 @@
     "get_loss_wrapper",
     "ddp_aware_model_move",
     "extract_metadata",
+    "get_model_directory",
 ]
 
 
+def get_model_directory(
+    output_dir: str, training_outputs_dir: str, logs_path: Optional[str] = None
+) -> None:
+    """
+    Takes the `training_outputs_dir`
+    (the directory where the pipeline saves its training artifacts),
+    and saves the training artifacts to `output_dir` in the `ModelDirectory` structure.
+
+    :param output_dir: The output path where the artifacts are saved
+        (adhering to the in `ModelDirectory` structure)
+    :param training_outputs_dir: The path to the existing directory
+        with the saved training artifacts
+    :param logs_path: Optional directory where the training logs reside
+    """
+    for root_file in ["model.onnx", "sample_inputs", "sample_outputs", "sample_labels"]:
+        root_file_path = os.path.join(training_outputs_dir, root_file)
+        if not os.path.exists(root_file_path):
+            raise ValueError(
+                f"File {root_file_path} missing. To create this file, "
+                "make sure that the `export` script (for exporting image "
+                "classification models) has been evoked."
+            )
+
+    setup_model_directory(
+        output_dir=output_dir,
+        training=os.path.join(training_outputs_dir, "training"),
+        deployment=os.path.join(training_outputs_dir, "model.onnx"),
+        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"),
+        model_card=os.path.join(training_outputs_dir, "model.md"),
+        logs=logs_path,
+        sample_originals=None,
+        analysis=None,
+        benchmarks=None,
+        eval_results=None,
+        recipes=None,
+    )
+    _LOGGER.info(f"Created `ModelDirectory` folder locally in {output_dir}")
+
+
 @unique
 class Tasks(Enum):
     """

src/sparseml/pytorch/utils/exporter.py

@@ -268,7 +268,7 @@ def export_pytorch(
         :param arch_key: if provided, the `arch_key` will be saved in the
             checkpoint
         """
-        pytorch_path = os.path.join(self._output_dir, "framework")
+        pytorch_path = os.path.join(self._output_dir, "training")
         pth_path = os.path.join(pytorch_path, name)
         create_parent_dirs(pth_path)
 
@@ -305,10 +305,10 @@ 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")
-        originals_dir = os.path.join(self._output_dir, "sample-originals")
+        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():
             for batch, lab, orig in zip(