comments

Author: George Ohashi

src/sparsezoo/api/graphql.py

@@ -82,7 +82,7 @@ def make_request(
         query = self.parse_query(
             operation_body=operation_body, arguments=arguments, fields=fields
         )
-
+        print(url or f"{BASE_API_URL}/v2/graphql")
         response = requests.post(
             url=url or f"{BASE_API_URL}/v2/graphql",
             json={

src/sparsezoo/api/utils.py

@@ -12,7 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Callable, Dict, List
+from typing import Any, Callable, Dict, List
 
 
 def to_camel_case(string: str):
@@ -32,7 +32,7 @@ def to_snake_case(string: str):
 
 
 def map_keys(
-    dictionary: Dict[str, str], mapper: Callable[[str], str]
+    dictionary: Dict[str, Any], mapper: Callable[[str], str]
 ) -> Dict[str, str]:
     """
     Given a dictionary, update its keys to a given mapper callable.
@@ -44,8 +44,7 @@ def map_keys(
         if isinstance(value, List) or isinstance(value, Dict):
             value_type = type(value)
             mapped_dict[mapper(key)] = value_type(
-                map_keys(dictionary=sub_dict, mapper=to_snake_case)
-                for sub_dict in value
+                map_keys(dictionary=sub_dict, mapper=mapper) for sub_dict in value
             )
         else:
             mapped_dict[mapper(key)] = value

src/sparsezoo/model/utils.py

@@ -75,11 +75,12 @@
     r"/(?P<architecture>[\.A-z0-9_]+)(-(?P<sub_architecture>[\.A-z0-9_]+))?"
     r"/(?P<framework>[\.A-z0-9_]+)"
     r"/(?P<repo>[\.A-z0-9_]+)"
-    r"/(?P<dataset>[\.A-z0-9_]+)"
+    r"/(?P<dataset>[\.A-z0-9_]+)(-(?P<training_scheme>[\.A-z0-9_]+))?"
     r"/(?P<sparse_tag>[\.A-z0-9_-]+)"
 )
 
 STUB_V2_REGEX_EXPR = (
+    r"^(zoo:)?"
     r"(?P<architecture>[\.A-z0-9_]+)"
     r"(-(?P<sub_architecture>[\.A-z0-9_]+))?"
     r"-(?P<source_dataset>[\.A-z0-9_]+)"
@@ -148,22 +149,35 @@ def load_files_from_stub(
         ],
     )
 
-    if len(models):
+    matching_models = len(models)
+    if matching_models == 0:
+        raise ValueError(
+            f"No matching models found with stub: {stub}." "Please try another stub"
+        )
+    if matching_models > 1:
+        logging.warning(
+            f"{len(models)} found from the stub: {stub}"
+            "Using the first model to obtain metadata."
+            "Proceed with caution"
+        )
+
+    if matching_models:
+        model = models[0]
 
-        model_id = models[0]["model_id"]
+        model_id = model["model_id"]
 
-        files = models[0].get("files")
+        files = model.get("files")
         include_file_download_url(files)
         files = restructure_request_json(request_json=files)
 
         if params is not None:
             files = filter_files(files=files, params=params)
 
-        training_results = models[0].get("training_results")
+        training_results = model.get("training_results")
 
-        benchmark_results = models[0].get("benchmark_results")
+        benchmark_results = model.get("benchmark_results")
 
-        model_onnx_size_compressed_bytes = models[0]["model_onnx_size_compressed_bytes"]
+        model_onnx_size_compressed_bytes = model["model_onnx_size_compressed_bytes"]
 
         throughput_results = _parse_results_metrics(
             results=benchmark_results, parser=ThroughputResults
@@ -601,9 +615,8 @@ def get_model_metadata_from_stub(stub: str) -> Dict[str, str]:
     return {}
 
 
-def is_stub(candidate: str):
-    if candidate.startswith(ZOO_STUB_PREFIX):
-        return True
-    if bool(get_model_metadata_from_stub(candidate)):
-        return True
-    return False
+def is_stub(candidate: str) -> bool:
+    return bool(
+        re.match(STUB_V1_REGEX_EXPR, candidate)
+        or re.match(STUB_V2_REGEX_EXPR, candidate)
+    )

tests/sparsezoo/model/test_model.py

@@ -137,32 +137,44 @@ def _assert_validation_results_exist(model):
     "stub, clone_sample_outputs, expected_files",
     [
         (
-            "zoo:cv/classification/mobilenet_v1-1.0/pytorch/sparseml/imagenet/pruned-moderate",  # noqa E501
+            (
+                "zoo:"
+                "cv/classification/mobilenet_v1-1.0/"
+                "pytorch/sparseml/imagenet/pruned-moderate"
+            ),
             True,
             files_ic,
         ),
         (
-            "zoo:nlp/question_answering/distilbert-none/pytorch/huggingface/squad/pruned80_quant-none-vnni",  # noqa E501
+            (
+                "zoo:"
+                "nlp/question_answering/distilbert-none/"
+                "pytorch/huggingface/squad/pruned80_quant-none-vnni"
+            ),
             False,
             files_nlp,
         ),
         (
-            "zoo:cv/detection/yolov5-s/pytorch/ultralytics/coco/pruned_quant-aggressive_94",  # noqa E501
+            (
+                "zoo:"
+                "cv/detection/yolov5-s/"
+                "pytorch/ultralytics/coco/pruned_quant-aggressive_94"
+            ),
             True,
             files_yolo,
         ),
         (
-            "mobilebert-squad_wikipedia_bookcorpus-14layer_pruned50.4block",  # noqa E501
+            "mobilebert-squad_wikipedia_bookcorpus-14layer_pruned50.4block",
             False,
             files_yolo,
         ),
         (
-            "yolov5-s-coco-pruned85_quantized",  # noqa E501
+            "yolov5-s-coco-pruned85_quantized",
             False,
             files_yolo,
         ),
         (
-            "resnet_v1-50-imagenet-channel30_pruned91",  # noqa E501
+            "resnet_v1-50-imagenet-channel30_pruned91",
             False,
             files_yolo,
         ),