From 44136f6bdc1d04ae3380cf3bb0c713ac6b64a828 Mon Sep 17 00:00:00 2001
From: dbogunowicz <damian@neuralmagic.com>
Date: Tue, 6 Feb 2024 07:06:05 +0000
Subject: [PATCH] initial commit

---
 hehe.py                             | 0
 src/sparsezoo/model/model.py        | 3 +++
 src/sparsezoo/model/utils.py        | 2 ++
 tests/sparsezoo/model/test_model.py | 1 +
 4 files changed, 6 insertions(+)
 create mode 100644 hehe.py

diff --git a/hehe.py b/hehe.py
new file mode 100644
index 00000000..e69de29b
diff --git a/src/sparsezoo/model/model.py b/src/sparsezoo/model/model.py
index b14ff503..21576337 100644
--- a/src/sparsezoo/model/model.py
+++ b/src/sparsezoo/model/model.py
@@ -188,6 +188,8 @@ def __init__(self, source: str, download_path: Optional[str] = None):
         )
         self.eval_results: File = self._file_from_files(files, display_name="eval.yaml")
 
+        self.metrics: File = self._file_from_files(files, display_name="metrics.yaml")
+
         # plaintext validation metrics optionally parsed from a zoo stub
         self.validation_results: Optional[
             Dict[str, List[ModelResult]]
@@ -248,6 +250,7 @@ def __init__(self, source: str, download_path: Optional[str] = None):
             "benchmarks": self.benchmarks,
             "benchmark": self.benchmark,
             "eval_results": self.eval_results,
+            "metrics": self.metrics,
         }
 
         self.inference_runner = InferenceRunner(
diff --git a/src/sparsezoo/model/utils.py b/src/sparsezoo/model/utils.py
index 3e061055..c1feb47b 100644
--- a/src/sparsezoo/model/utils.py
+++ b/src/sparsezoo/model/utils.py
@@ -59,6 +59,7 @@
     "outputs",
     "onnx_gz",
     "benchmark",
+    "metrics",
 }
 
 _LOGGER = logging.getLogger(__name__)
@@ -345,6 +346,7 @@ def restructure_request_json(
         that will not be filtered out during restructuring
     :return: restructured files
     """
+
     # create `training` folder
     training_dicts_list = fetch_from_request_json(
         request_json, "file_type", "framework"
diff --git a/tests/sparsezoo/model/test_model.py b/tests/sparsezoo/model/test_model.py
index 9a1bdeb0..dc2a7fa0 100644
--- a/tests/sparsezoo/model/test_model.py
+++ b/tests/sparsezoo/model/test_model.py
@@ -46,6 +46,7 @@
     "eval.yaml",
     "analysis.yaml",
     "model.md",
+    "metrics.yaml",
 }
 
 files_nlp = copy.copy(files_ic)