# Copyright (c) 2021 - present / Neuralmagic, Inc. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""
Code related to base functionality for making requests
"""
import os
import warnings
from typing import Any, Dict, List, Optional, Tuple, Union
from sparsezoo.utils import convert_to_bool
__all__ = [
"BASE_API_URL",
"ModelArgs",
"MODELS_API_URL",
"SPARSEZOO_TEST_MODE",
"RecipeArgs",
"RECIPES_API_URL",
"parse_zoo_stub",
"LATEST_PACKAGE_VERSION_URL",
]
SPARSEZOO_TEST_MODE = convert_to_bool(os.getenv("SPARSEZOO_TEST_MODE"))
BASE_API_URL = (
os.getenv("SPARSEZOO_API_URL")
if os.getenv("SPARSEZOO_API_URL")
else "https://api.neuralmagic.com"
)
MODELS_API_URL = f"{BASE_API_URL}/models"
RECIPES_API_URL = f"{BASE_API_URL}/recipes"
LATEST_PACKAGE_VERSION_URL = f"{BASE_API_URL}/packages/check-latest"
# optional prefix for stubs
ZOO_STUB_PREFIX = "zoo:"
[docs]def parse_zoo_stub(
stub: str, valid_params: Optional[List[str]] = None
) -> Tuple[str, Dict[str, str]]:
"""
:param stub: A SparseZoo model stub. i.e. 'model/stub/path',
'zoo:model/stub/path', 'zoo:model/stub/path?param1=value1¶m2=value2'
:param valid_params: list of expected parameter names to be encoded in the
stub. Will raise a warning if any unexpected param names are given. Leave
as None to not raise any warnings. Default is None
:return: the parsed base stub and a dictionary of parameter names and their values
"""
# strip optional zoo stub prefix
if stub.startswith(ZOO_STUB_PREFIX):
stub = stub[len(ZOO_STUB_PREFIX) :]
if "?" not in stub:
return stub, {}
stub_parts = stub.split("?")
if len(stub_parts) > 2:
raise ValueError(
"Invalid SparseZoo stub, query string must be preceded by only one '?'"
f"given {stub}"
)
stub, params = stub_parts
params = dict(param.split("=") for param in params.split("&"))
if valid_params is not None and any(param not in valid_params for param in params):
warnings.warn(
f"Invalid query string for stub {stub} valid params include {valid_params},"
f" given {list(params.keys())}"
)
return stub, params
[docs]class ModelArgs:
"""
Arguments for making requests into the sparsezoo
:param domain: The domain of the model the object belongs to;
e.g. cv, nlp
:param sub_domain: The sub domain of the model the object belongs to;
e.g. classification, segmentation
:param architecture: The architecture of the model the object belongs to;
e.g. resnet_v1, mobilenet_v1
:param sub_architecture: The sub architecture (scaling factor) of the model
the object belongs to; e.g. 50, 101, 152
:param framework: The framework the model the object belongs to was trained on;
e.g. pytorch, tensorflow
:param repo: the source repo for the model the object belongs to;
e.g. sparseml, torchvision
:param dataset: The dataset the model the object belongs to was trained on;
e.g. imagenet, cifar10
:param training_scheme: The training scheme used on the model the object belongs
to if any; e.g. augmented
:param sparse_name: The name describing the sparsification of the model
the object belongs to, e.g. base, pruned, pruned_quant,
:param sparse_category: The degree of sparsification of the model the object
belongs to; e.g. none, conservative (~100% baseline metric),
moderate (>=99% baseline metric), aggressive (<99% baseline metric)
:param sparse_target: The deployment target of sparsification of the model
the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu
:param release_version: The sparsezoo release version for the model
"""
def __init__(
self,
domain: Union[str, None] = None,
sub_domain: Union[str, None] = None,
architecture: Union[str, List[str], None] = None,
sub_architecture: Union[str, List[str], None] = None,
framework: Union[str, List[str], None] = None,
repo: Union[str, List[str], None] = None,
dataset: Union[str, List[str], None] = None,
training_scheme: Union[str, List[str], None] = None,
sparse_name: Union[str, List[str], None] = None,
sparse_category: Union[str, List[str], None] = None,
sparse_target: Union[str, List[str], None] = None,
release_version: Union[str, Any, None] = None,
**kwargs,
):
self._domain = domain
self._sub_domain = sub_domain
self._architecture = architecture
self._sub_architecture = sub_architecture
self._framework = framework
self._repo = repo
self._dataset = dataset
self._training_scheme = training_scheme
self._sparse_name = sparse_name
self._sparse_category = sparse_category
self._sparse_target = sparse_target
self._release_version = release_version
@property
def domain(self) -> Union[str, None]:
"""
:return: The domain of the model the object belongs to;
e.g. cv, nlp
"""
return self._domain
@property
def sub_domain(self) -> Union[str, None]:
"""
:return: The sub domain of the model the object belongs to;
e.g. classification, segmentation
"""
return self._sub_domain
@property
def architecture(self) -> Union[str, List[str], None]:
"""
:return: The architecture of the model the object belongs to;
e.g. resnet_v1, mobilenet_v1
"""
return self._architecture
@property
def sub_architecture(self) -> Union[str, List[str], None]:
"""
:return: The sub architecture (scaling factor) of the model
the object belongs to; e.g. 50, 101, 152
"""
return self._sub_architecture
@property
def framework(self) -> Union[str, List[str], None]:
"""
:return: The framework the model the object belongs to was trained on;
e.g. pytorch, tensorflow
"""
return self._framework
@property
def repo(self) -> Union[str, List[str], None]:
"""
:return: the source repo for the model the object belongs to;
e.g. sparseml, torchvision
"""
return self._repo
@property
def dataset(self) -> Union[str, List[str], None]:
"""
:return: The dataset the model the object belongs to was trained on;
e.g. imagenet, cifar10
"""
return self._dataset
@property
def training_scheme(self) -> Union[str, List[str], None]:
"""
:return: The training scheme used on the model the object belongs to if any;
e.g. augmented
"""
return self._training_scheme
@property
def sparse_name(self) -> Union[str, List[str], None]:
"""
:return: The name describing the sparsification of the model
the object belongs to, e.g. base, pruned, pruned_quant
"""
return self._sparse_name
@property
def sparse_category(self) -> Union[str, List[str], None]:
"""
:return: The degree of sparsification of the model the object belongs to;
e.g. none, conservative (~100% baseline metric),
moderate (>=99% baseline metric), aggressive (<99% baseline metric)
"""
return self._sparse_category
@property
def sparse_target(self) -> Union[str, List[str], None]:
"""
:return: The deployment target of sparsification of the model
the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu
"""
return self._sparse_target
@property
def release_version(self) -> Union[str, None]:
"""
:return: The sparsezoo release version for the model
"""
return self._release_version
@property
def architecture_id(self) -> str:
"""
:return: Unique id for the model architecture containing both the
architecture and sub_architecture
"""
if not self.architecture:
return ""
if not self.sub_architecture:
return f"{self.architecture}"
return f"{self.architecture}-{self.sub_architecture}"
@property
def training_id(self) -> str:
"""
:return: Unique id for how the model was trained containing both the
dataset and training_scheme
"""
if not self.dataset:
return ""
if not self.training_scheme:
return f"{self.dataset}"
return f"{self.dataset}-{self.training_scheme}"
@property
def sparse_id(self) -> str:
"""
:return: Unique id for how the model was sparse containing the
sparse_name, sparse_category, sparse_target
"""
if not self.sparse_name:
return ""
if not self.sparse_category:
return f"{self.sparse_name}"
if not self.sparse_target:
return f"{self.sparse_name}-{self.sparse_category}"
return f"{self.sparse_name}-{self.sparse_category}-{self.sparse_target}"
@property
def model_url_root(self) -> str:
"""
:return: root path for where the model is located in the sparsezoo
"""
if not self.domain:
return ""
if not self.sub_domain:
return f"{self.domain}"
return f"{self.domain}/{self.sub_domain}"
@property
def stub(self) -> str:
"""
:return: full path for where the model is located in the sparsezoo
"""
return "/".join(
[
self.model_url_root,
self.architecture_id,
f"{self.framework}" if self.framework else "",
f"{self.repo}" if self.repo else "",
self.training_id,
self.sparse_id,
]
)
@property
def model_url_args(self) -> List[str]:
"""
:return: arguments for searching in the sparsezoo
"""
args = []
for key in [
"architecture",
"sub_architecture",
"framework",
"repo",
"dataset",
"training_scheme",
"sparse_name",
"sparse_category",
"sparse_target",
]:
value = getattr(self, key)
if value and isinstance(value, List):
args.extend([f"{key}={item}" for item in value])
elif value:
args.append(f"{key}={value}")
return args
[docs]class RecipeArgs(ModelArgs):
"""
Arguments for making recipe requests into the sparsezoo
:param domain: The domain of the model the object belongs to;
e.g. cv, nlp
:param sub_domain: The sub domain of the model the object belongs to;
e.g. classification, segmentation
:param architecture: The architecture of the model the object belongs to;
e.g. resnet_v1, mobilenet_v1
:param sub_architecture: The sub architecture (scaling factor) of the model
the object belongs to; e.g. 50, 101, 152
:param framework: The framework the model the object belongs to was trained on;
e.g. pytorch, tensorflow
:param repo: the source repo for the model the object belongs to;
e.g. sparseml, torchvision
:param dataset: The dataset the model the object belongs to was trained on;
e.g. imagenet, cifar10
:param training_scheme: The training scheme used on the model the object belongs
to if any; e.g. augmented
:param sparse_name: The name describing the sparsification of the model
the object belongs to, e.g. base, pruned, pruned_quant,
:param sparse_category: The degree of sparsification of the model the object
belongs to; e.g. none, conservative (~100% baseline metric),
moderate (>=99% baseline metric), aggressive (<99% baseline metric)
:param sparse_target: The deployment target of sparsification of the model
the object belongs to; e.g. edge, deepsparse, deepsparse_throughput, gpu
:param release_version: The sparsezoo release version for the model
:param recipe_type: The recipe type; e.g. original, transfer_learn
"""
def __init__(
self,
domain: Union[str, None] = None,
sub_domain: Union[str, None] = None,
architecture: Union[str, List[str], None] = None,
sub_architecture: Union[str, List[str], None] = None,
framework: Union[str, List[str], None] = None,
repo: Union[str, List[str], None] = None,
dataset: Union[str, List[str], None] = None,
training_scheme: Union[str, List[str], None] = None,
sparse_name: Union[str, List[str], None] = None,
sparse_category: Union[str, List[str], None] = None,
sparse_target: Union[str, List[str], None] = None,
release_version: Union[str, Any, None] = None,
recipe_type: Union[str, None] = None,
**kwargs,
):
super(RecipeArgs, self).__init__(
domain=domain,
sub_domain=sub_domain,
architecture=architecture,
sub_architecture=sub_architecture,
framework=framework,
repo=repo,
dataset=dataset,
training_scheme=training_scheme,
sparse_name=sparse_name,
sparse_category=sparse_category,
sparse_target=sparse_target,
release_version=release_version,
**kwargs,
)
self._recipe_type = recipe_type
@property
def recipe_type(self) -> Union[str, None]:
"""
:return: The recipe type; e.g. original, transfer_learn
"""
return self._recipe_type
@property
def model_url_args(self) -> List[str]:
"""
:return: arguments for searching in the sparsezoo
"""
args = []
for key in [
"architecture",
"sub_architecture",
"framework",
"repo",
"dataset",
"training_scheme",
"sparse_name",
"sparse_category",
"sparse_target",
"recipe_type",
]:
value = getattr(self, key)
if value and isinstance(value, List):
args.extend([f"{key}={item}" for item in value])
elif value:
args.append(f"{key}={value}")
return args