onnxruntime_extensions 0.14.0__cp313-cp313-macosx_11_0_x86_64.whl

onnxruntime_extensions/_ortapi2.py

@@ -0,0 +1,274 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+###############################################################################
+
+"""
+_ortapi2.py: ONNXRuntime-Extensions Python API
+"""
+
+import copy
+import numpy as np
+from ._ocos import default_opset_domain, get_library_path, Opdef
+from ._cuops import onnx, onnx_proto, SingleOpGraph
+
+_ort_check_passed = False
+try:
+    from packaging import version as _ver
+    import onnxruntime as _ort
+
+    if _ver.parse(_ort.__version__) >= _ver.parse("1.10.0"):
+        _ort_check_passed = True
+except ImportError:
+    pass
+
+if not _ort_check_passed:
+    raise RuntimeError("please install ONNXRuntime/ONNXRuntime-GPU >= 1.10.0")
+
+
+def _ensure_opset_domain(model):
+    op_domain_name = default_opset_domain()
+    domain_missing = True
+    for oi_ in model.opset_import:
+        if oi_.domain == op_domain_name:
+            domain_missing = False
+
+    if domain_missing:
+        model.opset_import.extend([onnx.helper.make_operatorsetid(op_domain_name, 1)])
+
+    return model
+
+
+def hook_model_op(model, node_name, hook_func, input_types):
+    """
+    Add a hook function node in the ONNX Model, which could be used for the model diagnosis.
+    :param model: The ONNX model loaded as ModelProto
+    :param node_name: The node name where the hook will be installed
+    :param hook_func: The hook function, callback on the model inference
+    :param input_types: The input types as a list
+    :return: The ONNX model with the hook installed
+    """
+
+    # onnx.shape_inference is very unstable, useless.
+    # hkd_model = shape_inference.infer_shapes(model)
+    hkd_model = model
+
+    n_idx = 0
+    hnode, nnode = (None, None)
+    nodes = list(hkd_model.graph.node)
+    brkpt_name = node_name + "_hkd"
+    optype_name = "op_{}_{}".format(hook_func.__name__, node_name)
+    for n_ in nodes:
+        if n_.name == node_name:
+            input_names = list(n_.input)
+            brk_output_name = [i_ + "_hkd" for i_ in input_names]
+            hnode = onnx.helper.make_node(
+                optype_name, n_.input, brk_output_name, name=brkpt_name, domain=default_opset_domain()
+            )
+            nnode = n_
+            del nnode.input[:]
+            nnode.input.extend(brk_output_name)
+            break
+        n_idx += 1
+
+    if hnode is None:
+        raise ValueError("{} is not an operator node name".format(node_name))
+
+    repacked = nodes[:n_idx] + [hnode, nnode] + nodes[n_idx + 1 :]
+    del hkd_model.graph.node[:]
+    hkd_model.graph.node.extend(repacked)
+
+    Opdef.create(hook_func, op_type=optype_name, inputs=input_types, outputs=input_types)
+    return _ensure_opset_domain(hkd_model)
+
+
+def expand_onnx_inputs(model, target_input, extra_nodes, new_inputs):
+    """
+    Replace the existing inputs of a model with the new inputs, plus some extra nodes
+    :param model: The ONNX model loaded as ModelProto
+    :param target_input: The input name to be replaced
+    :param extra_nodes: The extra nodes to be added
+    :param new_inputs: The new input (type: ValueInfoProto) sequence
+    :return: The ONNX model after modification
+    """
+    graph = model.graph
+    new_inputs = [n for n in graph.input if n.name != target_input] + new_inputs
+    new_nodes = list(model.graph.node) + extra_nodes
+    new_graph = onnx.helper.make_graph(new_nodes, graph.name, new_inputs, list(graph.output), list(graph.initializer))
+
+    new_model = copy.deepcopy(model)
+    new_model.graph.CopyFrom(new_graph)
+
+    return _ensure_opset_domain(new_model)
+
+
+def get_opset_version_from_ort():
+    _ORT_OPSET_SUPPORT_TABLE = {
+        "1.5": 11,
+        "1.6": 12,
+        "1.7": 13,
+        "1.8": 14,
+        "1.9": 15,
+        "1.10": 15,
+        "1.11": 16,
+        "1.12": 17,
+        "1.13": 17,
+        "1.14": 18,
+        "1.15": 18,
+    }
+
+    ort_ver_string = ".".join(_ort.__version__.split(".")[0:2])
+    max_ver = max(_ORT_OPSET_SUPPORT_TABLE, key=_ORT_OPSET_SUPPORT_TABLE.get)
+    if ort_ver_string > max_ver:
+        ort_ver_string = max_ver
+    return _ORT_OPSET_SUPPORT_TABLE.get(ort_ver_string, 11)
+
+
+def make_onnx_model(graph, opset_version=0, extra_domain=default_opset_domain(), extra_opset_version=1):
+    if opset_version == 0:
+        opset_version = get_opset_version_from_ort()
+    fn_mm = (
+        onnx.helper.make_model_gen_version if hasattr(onnx.helper, "make_model_gen_version") else onnx.helper.make_model
+    )
+    model = fn_mm(graph, opset_imports=[onnx.helper.make_operatorsetid("ai.onnx", opset_version)])
+    model.opset_import.extend([onnx.helper.make_operatorsetid(extra_domain, extra_opset_version)])
+    return model
+
+
+class OrtPyFunction:
+    """
+    OrtPyFunction is a convenience class that serves as a wrapper around the ONNXRuntime InferenceSession,
+    equipped with registered onnxruntime-extensions. This allows execution of an ONNX model as if it were a
+    standard Python function. The order of the function arguments correlates directly with
+    the sequence of the input/output in the ONNX graph.
+    """
+
+    def get_ort_session_options(self):
+        so = _ort.SessionOptions()
+        for k, v in self.extra_session_options.items():
+            so.__setattr__(k, v)
+        so.register_custom_ops_library(get_library_path())
+        return so
+
+    def __init__(self, path_or_model=None, cpu_only=None):
+        self._onnx_model = None
+        self.ort_session = None
+        self.default_inputs = {}
+        self.execution_providers = ["CPUExecutionProvider"]
+        if not cpu_only:
+            if _ort.get_device() == "GPU":
+                self.execution_providers = ["CUDAExecutionProvider"]
+        self.extra_session_options = {}
+        mpath = None
+        if isinstance(path_or_model, str):
+            oxml = onnx.load_model(path_or_model)
+            mpath = path_or_model
+        else:
+            oxml = path_or_model
+        if path_or_model is not None:
+            self._bind(oxml, mpath)
+
+    def create_from_customop(self, op_type, *args, **kwargs):
+        graph = SingleOpGraph.build_graph(op_type, *args, **kwargs)
+        self._bind(make_onnx_model(graph))
+        return self
+
+    def add_default_input(self, **kwargs):
+        inputs = {
+            ky_: val_ if isinstance(val_, (np.ndarray, np.generic)) else np.asarray(list(val_), dtype=np.uint8)
+            for ky_, val_ in kwargs.items()
+        }
+
+        self.default_inputs.update(inputs)
+
+    @property
+    def onnx_model(self):
+        assert self._oxml is not None, "No onnx model attached yet."
+        return self._oxml
+
+    @property
+    def input_names(self):
+        return [vi_.name for vi_ in self.onnx_model.graph.input]
+
+    @property
+    def output_names(self):
+        return [vi_.name for vi_ in self.onnx_model.graph.output]
+
+    def _bind(self, oxml, model_path=None):
+        self.inputs = list(oxml.graph.input)
+        self.outputs = list(oxml.graph.output)
+        self._oxml = oxml
+        if model_path is not None:
+            self.ort_session = _ort.InferenceSession(
+                model_path, self.get_ort_session_options(), self.execution_providers
+            )
+        return self
+
+    def _ensure_ort_session(self):
+        if self.ort_session is None:
+            sess = _ort.InferenceSession(
+                self.onnx_model.SerializeToString(), self.get_ort_session_options(), self.execution_providers
+            )
+            self.ort_session = sess
+
+        return self.ort_session
+
+    @staticmethod
+    def _get_kwarg_device(kwargs):
+        cpuonly = kwargs.get("cpu_only", None)
+        if cpuonly is not None:
+            del kwargs["cpu_only"]
+        return cpuonly
+
+    @classmethod
+    def from_customop(cls, op_type, *args, **kwargs):
+        return cls(cpu_only=cls._get_kwarg_device(kwargs)).create_from_customop(op_type, *args, **kwargs)
+
+    @classmethod
+    def from_model(cls, path_or_model, *args, **kwargs):
+        fn = cls(path_or_model, cls._get_kwarg_device(kwargs))
+        return fn
+
+    def _argument_map(self, *args, **kwargs):
+        idx = 0
+        feed = {}
+        for i_ in self.inputs:
+            if i_.name in self.default_inputs:
+                feed[i_.name] = self.default_inputs[i_.name]
+                continue
+
+            x = args[idx]
+            ts_x = np.array(x) if isinstance(x, (int, float, bool)) else x
+            # numpy by default is int32 in some platforms, sometimes it is int64.
+            feed[i_.name] = (
+                ts_x.astype(np.int64) if i_.type.tensor_type.elem_type == onnx_proto.TensorProto.INT64 else ts_x
+            )
+            idx += 1
+
+        feed.update(kwargs)
+        return feed
+
+    def __call__(self, *args, **kwargs):
+        self._ensure_ort_session()
+        outputs = self.ort_session.run(None, self._argument_map(*args, **kwargs))
+        return outputs[0] if len(outputs) == 1 else tuple(outputs)
+
+
+def ort_inference(model, *args, cpu_only=True, **kwargs):
+    """
+    Run an ONNX model with ORT where args are inputs and return values are outputs.
+    """
+    return OrtPyFunction(model, cpu_only=cpu_only)(*args, **kwargs)
+
+
+def optimize_model(model_or_file, output_file):
+    sess_options = OrtPyFunction().get_ort_session_options()
+    sess_options.graph_optimization_level = _ort.GraphOptimizationLevel.ORT_ENABLE_BASIC
+    sess_options.optimized_model_filepath = output_file
+    _ort.InferenceSession(
+        model_or_file if isinstance(model_or_file, str) else model_or_file.SerializeToString(), sess_options
+    )
+
+
+ONNXRuntimeError = _ort.capi.onnxruntime_pybind11_state.Fail
+ONNXRuntimeException = _ort.capi.onnxruntime_pybind11_state.RuntimeException

onnxruntime_extensions/_torch_cvt.py

@@ -0,0 +1,231 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+###############################################################################
+
+"""
+_torch_cvt.py: Data processing graph converted from PyTorch
+"""
+
+import io
+import onnx
+import torch
+import numpy as np
+
+from onnx import numpy_helper
+
+from ._ortapi2 import make_onnx_model
+from ._cuops import SingleOpGraph
+from ._hf_cvt import HFTokenizerConverter
+from .util import remove_unused_initializers, mel_filterbank
+
+
+class _WhisperHParams:
+    SAMPLE_RATE = 16000
+    N_FFT = 400
+    N_MELS = 80
+    HOP_LENGTH = 160
+    CHUNK_LENGTH = 30
+    N_SAMPLES = CHUNK_LENGTH * SAMPLE_RATE  # 480000 samples in a 30-second chunk
+    N_FRAMES = N_SAMPLES // HOP_LENGTH
+
+
+class CustomOpStftNorm(torch.autograd.Function):
+    @staticmethod
+    def symbolic(g, self, n_fft, hop_length, window):
+        t_n_fft = g.op('Constant', value_t=torch.tensor(
+            n_fft, dtype=torch.int64))
+        t_hop_length = g.op('Constant', value_t=torch.tensor(
+            hop_length, dtype=torch.int64))
+        t_frame_size = g.op(
+            'Constant', value_t=torch.tensor(n_fft, dtype=torch.int64))
+        return g.op("ai.onnx.contrib::StftNorm", self, t_n_fft, t_hop_length, window, t_frame_size)
+
+    @staticmethod
+    def forward(ctx, audio, n_fft, hop_length, window):
+        win_length = window.shape[0]
+        stft = torch.stft(audio, n_fft, hop_length, win_length, window,
+                          center=True, pad_mode="reflect", normalized=False, onesided=True, return_complex=True)
+        return stft.abs() ** 2
+
+
+class WhisperPrePipeline(torch.nn.Module):
+    def __init__(self, sr=_WhisperHParams.SAMPLE_RATE, n_fft=_WhisperHParams.N_FFT,
+                 hop_length=_WhisperHParams.HOP_LENGTH, n_mels=_WhisperHParams.N_MELS,
+                 n_samples=_WhisperHParams.N_SAMPLES):
+        super().__init__()
+        self.n_samples = n_samples
+        self.hop_length = hop_length
+        self.n_fft = n_fft
+        self.window = torch.hann_window(n_fft)
+        self.mel_filters = torch.from_numpy(
+            mel_filterbank(sr=sr, n_fft=n_fft, n_mels=n_mels))
+
+    def forward(self, audio_pcm: torch.Tensor):
+        stft_norm = CustomOpStftNorm.apply(audio_pcm,
+                                           self.n_fft,
+                                           self.hop_length,
+                                           self.window)
+        magnitudes = stft_norm[:, :, :-1]
+        mel_spec = self.mel_filters @ magnitudes
+        log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+        spec_min = log_spec.max() - 8.0
+        log_spec = torch.maximum(log_spec, spec_min)
+        spec_shape = log_spec.shape
+        padding_spec = torch.ones(spec_shape[0],
+                                  spec_shape[1],
+                                  self.n_samples // self.hop_length -
+                                  spec_shape[2],
+                                  dtype=torch.float)
+        padding_spec *= spec_min
+        log_spec = torch.cat((log_spec, padding_spec), dim=2)
+        log_spec = (log_spec + 4.0) / 4.0
+        return log_spec
+
+
+def _to_onnx_stft(onnx_model, n_fft):
+    """Convert custom-op STFT-Norm to ONNX STFT"""
+    node_idx = 0
+    new_stft_nodes = []
+    stft_norm_node = None
+    for node in onnx_model.graph.node:
+        if node.op_type == "StftNorm":
+            stft_norm_node = node
+            break
+        node_idx += 1
+
+    if stft_norm_node is None:
+        raise RuntimeError("Cannot find STFTNorm node in the graph")
+
+    make_node = onnx.helper.make_node
+    replaced_nodes = [
+        make_node('Constant', inputs=[], outputs=['const_minus_1_output_0'], name='const_minus_1',
+                  value=numpy_helper.from_array(np.array([-1], dtype='int64'))),
+        make_node('Constant', inputs=[], outputs=['const_14_output_0'], name='const_14',
+                  value=numpy_helper.from_array(np.array([0,
+                                                          n_fft // 2, 0,
+                                                          n_fft // 2], dtype='int64'),
+                                                name='const_14')),
+        make_node('Pad',
+                  inputs=[stft_norm_node.input[0], 'const_14_output_0'],
+                  outputs=['pad_1_output_0'], mode='reflect'),
+        make_node('Unsqueeze',
+                  inputs=['pad_1_output_0', 'const_minus_1_output_0'],
+                  outputs=['unsqueeze_1_output_0'],
+                  name='unsqueeze_1'),
+        make_node('STFT',
+                  inputs=['unsqueeze_1_output_0', stft_norm_node.input[2],
+                          stft_norm_node.input[3], stft_norm_node.input[4]],
+                  outputs=['stft_output_0'], name='stft', onesided=1),
+        make_node('Transpose', inputs=['stft_output_0'], outputs=['transpose_1_output_0'], name='transpose_1',
+                  perm=[0, 2, 1, 3]),
+        make_node('Constant', inputs=[], outputs=['const_17_output_0'], name='const_17',
+                  value=numpy_helper.from_array(np.array([2], dtype='int64'), name='')),
+        make_node('Constant', inputs=[], outputs=['const_18_output_0'], name='const_18',
+                  value=numpy_helper.from_array(np.array([0], dtype='int64'), name='')),
+        make_node('Constant', inputs=[], outputs=['const_20_output_0'], name='const_20',
+                  value=numpy_helper.from_array(np.array([1], dtype='int64'), name='')),
+        make_node('Slice', inputs=['transpose_1_output_0', 'const_18_output_0', 'const_minus_1_output_0',
+                                   'const_17_output_0', 'const_20_output_0'], outputs=['slice_1_output_0'],
+                  name='slice_1'),
+        make_node('Constant', inputs=[], outputs=[
+                  'const0_output_0'], name='const0', value_int=0),
+        make_node('Constant', inputs=[], outputs=[
+                  'const1_output_0'], name='const1', value_int=1),
+        make_node('Gather', inputs=['slice_1_output_0', 'const0_output_0'], outputs=['gather_4_output_0'],
+                  name='gather_4', axis=3),
+        make_node('Gather', inputs=['slice_1_output_0', 'const1_output_0'], outputs=['gather_5_output_0'],
+                  name='gather_5', axis=3),
+        make_node('Mul', inputs=['gather_4_output_0', 'gather_4_output_0'], outputs=[
+                  'mul_output_0'], name='mul0'),
+        make_node('Mul', inputs=['gather_5_output_0', 'gather_5_output_0'], outputs=[
+                  'mul_1_output_0'], name='mul1'),
+        make_node('Add', inputs=['mul_output_0', 'mul_1_output_0'], outputs=[
+                  stft_norm_node.output[0]], name='add0'),
+    ]
+    new_stft_nodes.extend(onnx_model.graph.node[:node_idx])
+    new_stft_nodes.extend(replaced_nodes)
+    new_stft_nodes.extend(onnx_model.graph.node[node_idx + 1:])
+    del onnx_model.graph.node[:]
+    onnx_model.graph.node.extend(new_stft_nodes)
+    onnx.checker.check_model(onnx_model)
+    return onnx_model
+
+
+def _torch_export(*arg, **kwargs):
+    with io.BytesIO() as f:
+        torch.onnx.export(*arg, f, **kwargs)
+        return onnx.load_from_string(f.getvalue())
+
+
+class WhisperDataProcGraph:
+    def __init__(self, processor, **kwargs):
+        self.hf_processor = processor
+        _opset = kwargs.pop('opset', 17)
+        self.opset_version = _opset if _opset else 17
+
+    def pre_processing(self, **kwargs):
+        use_audio_decoder = kwargs.pop('USE_AUDIO_DECODER', True)
+        use_onnx_stft = kwargs.pop('USE_ONNX_STFT', True)
+        feature_extractor = self.hf_processor.feature_extractor
+        whisper_processing = WhisperPrePipeline(
+            feature_extractor.sampling_rate,
+            feature_extractor.n_fft,
+            feature_extractor.hop_length,
+            feature_extractor.feature_size,
+            feature_extractor.n_samples)
+
+        audio_pcm = torch.rand((1, 32000), dtype=torch.float32)
+        model_args = (audio_pcm,)
+        pre_model = _torch_export(
+            whisper_processing,
+            model_args,
+            input_names=["audio_pcm"],
+            output_names=["log_mel"],
+            do_constant_folding=True,
+            export_params=True,
+            opset_version=self.opset_version,
+            dynamic_axes={
+                "audio_pcm": {1: "sample_len"},
+            }
+        )
+        if use_onnx_stft:
+            pre_model = _to_onnx_stft(pre_model, feature_extractor.n_fft)
+            remove_unused_initializers(pre_model.graph)
+
+        pre_full = pre_model
+        if use_audio_decoder:
+            audecoder_g = SingleOpGraph.build_graph(
+                "AudioDecoder",
+                downsampling_rate=feature_extractor.sampling_rate,
+                stereo_to_mono=1)
+            audecoder_m = make_onnx_model(audecoder_g)
+            pre_full = onnx.compose.merge_models(
+                audecoder_m,
+                pre_model,
+                io_map=[("floatPCM", "audio_pcm")])
+
+        return pre_full
+
+    def post_processing(self, **kwargs):
+        skip_special_tokens = kwargs.get('skip_special_tokens', True)
+        g = SingleOpGraph.build_graph(
+            "BpeDecoder",
+            cvt=HFTokenizerConverter(self.hf_processor.tokenizer).bpe_decoder,
+            skip_special_tokens=skip_special_tokens)
+
+        bpenode = g.node[0]
+        bpenode.input[0] = "generated_ids"
+        nodes = [onnx.helper.make_node('Cast', ['sequences'], ["generated_ids"], to=onnx.TensorProto.INT64),
+                 bpenode]
+        del g.node[:]
+        g.node.extend(nodes)
+
+        inputs = [onnx.helper.make_tensor_value_info(
+            "sequences", onnx.TensorProto.INT32, ['N', 'seq_len', 'ids'])]
+        del g.input[:]
+        g.input.extend(inputs)
+        g.output[0].type.CopyFrom(onnx.helper.make_tensor_type_proto(
+            onnx.TensorProto.STRING, ['N', 'text']))
+
+        return make_onnx_model(g, opset_version=self.opset_version)

onnxruntime_extensions/_version.py

@@ -0,0 +1,2 @@
+# Generated by setup.py, DON'T MANUALLY UPDATE IT!
+__version__ = "0.14.0"

onnxruntime_extensions/cmd.py

@@ -0,0 +1,66 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License. See License.txt in the project root for
+# license information.
+###############################################################################
+
+"""
+cmd.py: cli commands for onnxruntime_extensions
+"""
+
+import os
+import argparse
+import onnx
+import numpy
+
+from onnx import onnx_pb, save_tensor, numpy_helper
+from ._ortapi2 import OrtPyFunction
+
+
+class ORTExtCommands:
+    def __init__(self, model='model.onnx', testdata_dir=None) -> None:
+        self._model = model
+        self._testdata_dir = testdata_dir
+
+    def run(self, *args):
+        """
+        Run an onnx model with the arguments as its inputs
+        """
+        op_func = OrtPyFunction.from_model(self._model)
+        np_args = [numpy.asarray(_x) for _x in args]
+        for _idx, _sch in enumerate(op_func.inputs):
+            if _sch.type.tensor_type.elem_type == onnx_pb.TensorProto.FLOAT:
+                np_args[_idx] = np_args[_idx].astype(numpy.float32)
+
+        print(op_func(*np_args))
+        if self._testdata_dir:
+            testdir = os.path.expanduser(self._testdata_dir)
+            target_dir = os.path.join(testdir, 'test_data_set_0')
+            os.makedirs(target_dir, exist_ok=True)
+            for _idx, _x in enumerate(np_args):
+                fn = os.path.join(target_dir, "input_{}.pb".format(_idx))
+                save_tensor(numpy_helper.from_array(_x, op_func.inputs[_idx].name), fn)
+            onnx.save_model(op_func.onnx_model, os.path.join(testdir, 'model.onnx'))
+
+    def selfcheck(self, *args):
+        print("The extensions loaded, status: OK.")
+
+
+def main():
+    parser = argparse.ArgumentParser(description="ORT Extension commands")
+    parser.add_argument("command", choices=["run", "selfcheck"])
+    parser.add_argument("--model", default="model.onnx", help="Path to the ONNX model file")
+    parser.add_argument("--testdata-dir", help="Path to the test data directory")
+    parser.add_argument("args", nargs=argparse.REMAINDER, help="Additional arguments")
+
+    args = parser.parse_args()
+
+    ort_commands = ORTExtCommands(model=args.model, testdata_dir=args.testdata_dir)
+
+    if args.command == "run":
+        ort_commands.run(*args.args)
+    elif args.command == "selfcheck":
+        ort_commands.selfcheck(*args.args)
+
+
+if __name__ == '__main__':
+    main()