onnx 1.13.0__cp311-cp311-macosx_10_12_universal2.whl → 1.14.0__cp311-cp311-macosx_10_12_universal2.whl

onnx/hub.py

@@ -1,3 +1,5 @@
+# Copyright (c) ONNX Project Contributors
+#
 # SPDX-License-Identifier: Apache-2.0
 """ONNX Model Hub
 
@@ -7,6 +9,7 @@ import hashlib
 import json
 import os
 import sys
+import tarfile
 from io import BytesIO
 from os.path import join
 from typing import IO, Any, Dict, List, Optional, Set, Tuple, cast
@@ -58,9 +61,7 @@ class ModelInfo:
         self.raw_model_info: Dict[str, Any] = raw_model_info
 
     def __str__(self) -> str:
-        return "ModelInfo(model={}, opset={}, path={}, metadata={})".format(
-            self.model, self.opset, self.model_path, self.metadata
-        )
+        return f"ModelInfo(model={self.model}, opset={self.opset}, path={self.model_path}, metadata={self.metadata})"
 
     def __repr__(self) -> str:
         return self.__str__()
@@ -72,7 +73,7 @@ def set_dir(new_dir: str) -> None:
 
     :param new_dir: location of new model hub cache
     """
-    global _ONNX_HUB_DIR
+    global _ONNX_HUB_DIR  # pylint: disable=global-statement
     _ONNX_HUB_DIR = new_dir
 
 
@@ -120,8 +121,7 @@ def _get_base_url(repo: str, lfs: bool = False) -> str:
 
     if lfs:
         return f"https://media.githubusercontent.com/media/{repo_owner}/{repo_name}/{repo_ref}/"
-    else:
-        return f"https://raw.githubusercontent.com/{repo_owner}/{repo_name}/{repo_ref}/"
+    return f"https://raw.githubusercontent.com/{repo_owner}/{repo_name}/{repo_ref}/"
 
 
 def _download_file(url: str, file_name: str) -> None:
@@ -175,14 +175,14 @@ def list_models(
     # Filter by tags
     if tags is None:
         return matching_models
-    else:
-        canonical_tags = {t.lower() for t in tags}
-        matching_info_list: List[ModelInfo] = []
-        for m in matching_models:
-            model_tags = {t.lower() for t in m.tags}
-            if len(canonical_tags.intersection(model_tags)) > 0:
-                matching_info_list.append(m)
-        return matching_info_list
+
+    canonical_tags = {t.lower() for t in tags}
+    matching_info_list: List[ModelInfo] = []
+    for m in matching_models:
+        model_tags = {t.lower() for t in m.tags}
+        if len(canonical_tags.intersection(model_tags)) > 0:
+            matching_info_list.append(m)
+    return matching_info_list
 
 
 def get_model_info(
@@ -241,10 +241,7 @@ def load(
 
     if force_reload or not os.path.exists(local_model_path):
         if not _verify_repo_ref(repo) and not silent:
-            msg = (
-                'The model repo specification "{}" is not trusted and may'
-                " contain security vulnerabilities. Only continue if you trust this repo."
-            ).format(repo)
+            msg = f"The model repo specification {repo} is not trusted and may contain security vulnerabilities. Only continue if you trust this repo."
 
             print(msg, file=sys.stderr)
             print("Continue?[y/n]")
@@ -266,9 +263,181 @@ def load(
         if not downloaded_sha == selected_model.model_sha:
             raise AssertionError(
                 (
-                    "The cached model {} has SHA256 {} while checksum should be {}. "
+                    f"The cached model {selected_model.model} has SHA256 {downloaded_sha} while checksum should be {selected_model.model_sha}."
                     + "The model in the hub may have been updated. Use force_reload to download the model from the model hub."
-                ).format(selected_model.model, downloaded_sha, selected_model.model_sha)
+                )
             )
 
     return onnx.load(cast(IO[bytes], BytesIO(model_bytes)))
+
+
+def download_model_with_test_data(
+    model: str,
+    repo: str = "onnx/models:main",
+    opset: Optional[int] = None,
+    force_reload: bool = False,
+    silent: bool = False,
+) -> Optional[str]:
+    """
+    Downloads a model along with test data by name from the onnx model hub and returns the directory to which the files have been extracted.
+
+    :param model: The name of the onnx model in the manifest. This field is case-sensitive
+    :param repo: The location of the model repo in format "user/repo[:branch]".
+        If no branch is found will default to "main"
+    :param opset: The opset of the model to download. The default of `None` automatically chooses the largest opset
+    :param force_reload: Whether to force the model to re-download even if its already found in the cache
+    :param silent: Whether to suppress the warning message if the repo is not trusted.
+    :return: str or None
+    """
+    selected_model = get_model_info(model, repo, opset)
+
+    local_model_with_data_path_arr = selected_model.metadata[
+        "model_with_data_path"
+    ].split("/")
+
+    model_with_data_sha = selected_model.metadata["model_with_data_sha"]
+
+    if model_with_data_sha is not None:
+        local_model_with_data_path_arr[
+            -1
+        ] = f"{model_with_data_sha}_{local_model_with_data_path_arr[-1]}"
+    local_model_with_data_path = join(
+        _ONNX_HUB_DIR, os.sep.join(local_model_with_data_path_arr)
+    )
+
+    if force_reload or not os.path.exists(local_model_with_data_path):
+        if not _verify_repo_ref(repo) and not silent:
+            msg = f"The model repo specification {repo} is not trusted and may contain security vulnerabilities. Only continue if you trust this repo."
+
+            print(msg, file=sys.stderr)
+            print("Continue?[y/n]")
+            if input().lower() != "y":
+                return None
+
+        os.makedirs(os.path.dirname(local_model_with_data_path), exist_ok=True)
+        lfs_url = _get_base_url(repo, True)
+        print(f"Downloading {model} to local path {local_model_with_data_path}")
+        _download_file(
+            lfs_url + selected_model.metadata["model_with_data_path"],
+            local_model_with_data_path,
+        )
+    else:
+        print(f"Using cached {model} model from {local_model_with_data_path}")
+
+    with open(local_model_with_data_path, "rb") as f:
+        model_with_data_bytes = f.read()
+
+    if model_with_data_sha is not None:
+        downloaded_sha = hashlib.sha256(model_with_data_bytes).hexdigest()
+        if not downloaded_sha == model_with_data_sha:
+            raise AssertionError(
+                (
+                    f"The cached model {selected_model.model} has SHA256 {downloaded_sha} while checksum should be {model_with_data_sha}."
+                    + "The model in the hub may have been updated. Use force_reload to download the model from the model hub."
+                )
+            )
+
+    with tarfile.open(local_model_with_data_path) as model_with_data_zipped:
+        # FIXME: Avoid index manipulation with magic numbers
+        local_model_with_data_dir_path = local_model_with_data_path[
+            0 : len(local_model_with_data_path) - 7
+        ]
+        model_with_data_zipped.extractall(local_model_with_data_dir_path)
+    model_with_data_path = (
+        local_model_with_data_dir_path
+        + "/"
+        + os.listdir(local_model_with_data_dir_path)[0]
+    )
+
+    return model_with_data_path
+
+
+def load_composite_model(
+    network_model: str,
+    preprocessing_model: str,
+    network_repo: str = "onnx/models:main",
+    preprocessing_repo: str = "onnx/models:main",
+    opset: Optional[int] = None,
+    force_reload: bool = False,
+    silent: bool = False,
+) -> Optional[onnx.ModelProto]:
+    """
+    Builds a composite model including data preprocessing by downloading a network and a preprocessing model
+    and combine it into a single model
+
+    :param model: The name of the onnx model in the manifest. This field is case-sensitive
+    :param repo: The location of the model repo in format "user/repo[:branch]".
+        If no branch is found will default to "main"
+    :param opset: The opset of the model to download. The default of `None` automatically chooses the largest opset
+    :param force_reload: Whether to force the model to re-download even if its already found in the cache
+    :param silent: Whether to suppress the warning message if the repo is not trusted.
+    :return: ModelProto or None
+    """
+    preprocessing = load(
+        preprocessing_model, preprocessing_repo, opset, force_reload, silent
+    )
+    if preprocessing is None:
+        raise RuntimeError(
+            f"Could not load the preprocessing model: {preprocessing_model}"
+        )
+    network = load(network_model, network_repo, opset, force_reload, silent)
+    if network is None:
+        raise RuntimeError(f"Could not load the network model: {network_model}")
+
+    all_domains: Set[str] = set()
+    domains_to_version_network: Dict[str, int] = {}
+    domains_to_version_preprocessing: Dict[str, int] = {}
+
+    for opset_import_entry in network.opset_import:
+        domain = (
+            "ai.onnx" if opset_import_entry.domain == "" else opset_import_entry.domain
+        )
+        all_domains.add(domain)
+        domains_to_version_network[domain] = opset_import_entry.version
+
+    for opset_import_entry in preprocessing.opset_import:
+        domain = (
+            "ai.onnx" if opset_import_entry.domain == "" else opset_import_entry.domain
+        )
+        all_domains.add(domain)
+        domains_to_version_preprocessing[domain] = opset_import_entry.version
+
+    preprocessing_opset_version = -1
+    network_opset_version = -1
+    for domain in all_domains:
+        if domain == "ai.onnx":
+            preprocessing_opset_version = domains_to_version_preprocessing[domain]
+            network_opset_version = domains_to_version_network[domain]
+        elif (
+            domain in domains_to_version_preprocessing
+            and domain in domains_to_version_network
+            and domains_to_version_preprocessing[domain]
+            != domains_to_version_preprocessing[domain]
+        ):
+            raise ValueError(
+                f"Can not merge {preprocessing_model} and {network_model} because they contain "
+                f"different opset versions for domain {domain} ({domains_to_version_preprocessing[domain]}) "
+                f"and {domains_to_version_network[domain]}). Only the default domain can be "
+                "automatically converted to the highest version of the two."
+            )
+    if preprocessing_opset_version > network_opset_version:
+        network = onnx.version_converter.convert_version(
+            network, preprocessing_opset_version
+        )
+        network.ir_version = preprocessing.ir_version
+        onnx.checker.check_model(network)
+    elif network_opset_version > preprocessing_opset_version:
+        preprocessing = onnx.version_converter.convert_version(
+            preprocessing, network_opset_version
+        )
+        preprocessing.ir_version = network.ir_version
+        onnx.checker.check_model(preprocessing)
+
+    io_map = []
+    for out_entry, in_entry in zip(preprocessing.graph.output, network.graph.input):
+        io_map.append((out_entry.name, in_entry.name))
+
+    model_with_preprocessing = onnx.compose.merge_models(
+        preprocessing, network, io_map=io_map
+    )
+    return model_with_preprocessing

onnx/mapping.py

@@ -1,3 +1,7 @@
+# Copyright (c) ONNX Project Contributors
+#
+# SPDX-License-Identifier: Apache-2.0
+
 import warnings
 from typing import Any, Dict, NamedTuple, Union, cast
 
@@ -38,8 +42,7 @@ TENSOR_TYPE_MAP = {
     int(TensorProto.FLOAT16): TensorDtypeMap(
         np.dtype("float16"), int(TensorProto.UINT16), "TensorProto.FLOAT16"
     ),
-    # Native numpy does not support bfloat16 so now use float32 for bf16 values
-    # TODO ONNX should dirtectly use bfloat16 for bf16 values after numpy has supported bfloat16 type
+    # Native numpy does not support bfloat16 so now use float32.
     int(TensorProto.BFLOAT16): TensorDtypeMap(
         np.dtype("float32"), int(TensorProto.UINT16), "TensorProto.BFLOAT16"
     ),
@@ -61,6 +64,19 @@ TENSOR_TYPE_MAP = {
     int(TensorProto.STRING): TensorDtypeMap(
         np.dtype("object"), int(TensorProto.STRING), "TensorProto.STRING"
     ),
+    # Native numpy does not support float8 types, so now use float32 for these types.
+    int(TensorProto.FLOAT8E4M3FN): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E4M3FN"
+    ),
+    int(TensorProto.FLOAT8E4M3FNUZ): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E4M3FNUZ"
+    ),
+    int(TensorProto.FLOAT8E5M2): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2"
+    ),
+    int(TensorProto.FLOAT8E5M2FNUZ): TensorDtypeMap(
+        np.dtype("float32"), int(TensorProto.UINT8), "TensorProto.FLOAT8E5M2FNUZ"
+    ),
 }
 
 
@@ -87,8 +103,8 @@ class DeprecatedWarningDict(dict):  # type: ignore
         if not self._future_function:
             warnings.warn(
                 str(
-                    f"`mapping.{self._origin_function}` is now deprecated and will be removed in the next release or so."
-                    + "To silence this warning, please simply use if-else statement to get the corresponding value."
+                    f"`mapping.{self._origin_function}` is now deprecated and will be removed in a future release."
+                    "To silence this warning, please simply use if-else statement to get the corresponding value."
                 ),
                 DeprecationWarning,
                 stacklevel=2,
@@ -96,8 +112,8 @@ class DeprecatedWarningDict(dict):  # type: ignore
         else:
             warnings.warn(
                 str(
-                    f"`mapping.{self._origin_function}` is now deprecated and will be removed in the next release or so."
-                    + "To silence this warning, please use `helper.{self._future_function}` instead."
+                    f"`mapping.{self._origin_function}` is now deprecated and will be removed in a future release."
+                    f"To silence this warning, please use `helper.{self._future_function}` instead."
                 ),
                 DeprecationWarning,
                 stacklevel=2,
@@ -126,7 +142,16 @@ TENSOR_TYPE_TO_STORAGE_TENSOR_TYPE = DeprecatedWarningDict(
 # NP_TYPE_TO_TENSOR_TYPE will be eventually removed in the future
 # and _NP_TYPE_TO_TENSOR_TYPE will only be used internally
 _NP_TYPE_TO_TENSOR_TYPE = {
-    v: k for k, v in TENSOR_TYPE_TO_NP_TYPE.items() if k != TensorProto.BFLOAT16
+    v: k
+    for k, v in TENSOR_TYPE_TO_NP_TYPE.items()
+    if k
+    not in (
+        TensorProto.BFLOAT16,
+        TensorProto.FLOAT8E4M3FN,
+        TensorProto.FLOAT8E4M3FNUZ,
+        TensorProto.FLOAT8E5M2,
+        TensorProto.FLOAT8E5M2FNUZ,
+    )
 }
 
 # Currently native numpy does not support bfloat16 so TensorProto.BFLOAT16 is ignored for now
@@ -143,6 +168,7 @@ _STORAGE_TENSOR_TYPE_TO_FIELD = {
     int(TensorProto.FLOAT): "float_data",
     int(TensorProto.INT32): "int32_data",
     int(TensorProto.INT64): "int64_data",
+    int(TensorProto.UINT8): "int32_data",
     int(TensorProto.UINT16): "int32_data",
     int(TensorProto.DOUBLE): "double_data",
     int(TensorProto.COMPLEX64): "float_data",