onnxruntime_extensions 0.14.0__cp313-cp313-macosx_11_0_arm64.whl

onnxruntime_extensions/pnp/_base.py

@@ -0,0 +1,124 @@
+import io
+import onnx
+import torch
+from typing import Any
+from onnx.onnx_pb import TensorProto
+from torch.onnx import TrainingMode, export as _export
+
+from ._onnx_ops import OPSET_TO_IR_VERSION
+
+
+def _export_f(model, *args,
+              opset_version=None,
+              output_path=None,
+              output_seq=0,
+              export_params=True,
+              verbose=False,
+              input_names=None,
+              output_names=None,
+              operator_export_type=None,
+              do_constant_folding=True,
+              dynamic_axes=None,
+              keep_initializers_as_inputs=None,
+              custom_opsets=None):
+
+    with io.BytesIO() as f:
+        _export(model, args, f,
+                export_params=export_params, verbose=verbose,
+                training=TrainingMode.EVAL, input_names=input_names,
+                output_names=output_names,
+                operator_export_type=operator_export_type, opset_version=opset_version,
+                do_constant_folding=do_constant_folding,
+                dynamic_axes=dynamic_axes,
+                keep_initializers_as_inputs=keep_initializers_as_inputs,
+                custom_opsets=custom_opsets)
+
+        mdl = onnx.load_model(io.BytesIO(f.getvalue()))
+        for ops in mdl.opset_import:
+            if ops.domain in ('', 'ai.onnx'):
+                mdl.ir_version = OPSET_TO_IR_VERSION[ops.version]
+        if output_path is not None:
+            if output_seq > 0:
+                output_path.replace('.onnx', '.{}.onnx'.format(output_seq))
+            onnx.save_model(mdl, output_path)
+        return mdl
+
+
+class _ProcessingModule:
+
+    def __init__(self):
+        super(_ProcessingModule, self).__init__()
+        _ProcessingModule.register_customops()
+
+    @staticmethod
+    @torch.jit.unused
+    def _argsort(g, x, dim, descending):
+        return g.op('ai.onnx.contrib::ArgSort', x, dim)
+
+    @classmethod
+    @torch.jit.unused
+    def register_customops(cls):
+        if hasattr(cls, 'loaded'):
+            return True
+
+        torch.onnx.register_custom_op_symbolic('::argsort', cls._argsort, 1)
+        # ... more
+
+        cls.loaded = True
+        return True
+
+    @torch.jit.unused
+    def export(self, *args, opset_version=None, script_mode=False, output_path=None, output_seq=0, **kwargs):
+        if opset_version is None:
+            raise RuntimeError('No opset_version found in the kwargs.')
+        mod = self
+        if script_mode and not isinstance(mod, torch.jit.ScriptModule):
+            mod = torch.jit.script(mod)
+
+        return _export_f(mod,
+                         *args,
+                         opset_version=opset_version,
+                         output_path=output_path,
+                         output_seq=output_seq, **kwargs)
+
+
+class ProcessingTracedModule(torch.nn.Module, _ProcessingModule):
+    def __init__(self, func_obj=None):
+        super().__init__()
+        self.func_obj = func_obj
+
+    def forward(self, *args):
+        assert self.func_obj is not None, "No forward method found."
+        return self.func_obj(*args)
+
+
+class ProcessingScriptModule(torch.nn.Module, _ProcessingModule):
+
+    @torch.jit.unused
+    def export(self, *args, **kwargs):
+        return super().export(*args, script_mode=True, **kwargs)
+
+
+class CustomFunction(torch.autograd.Function):
+    @staticmethod
+    def jvp(ctx: Any, *grad_inputs: Any) -> Any:
+        pass
+
+    @staticmethod
+    def backward(ctx: Any, *grad_outputs: Any) -> Any:
+        return grad_outputs
+
+    @classmethod
+    def forward(cls, ctx: Any, *args: Any, **kwargs: Any) -> Any:
+        pass
+
+    @classmethod
+    def symbolic(cls, g, *args):
+        return g.op('ai.onnx.contrib::' + cls.__name__, *args)
+
+
+tensor_data_type = TensorProto
+
+
+def is_processing_module(m):
+    return isinstance(m, _ProcessingModule)

onnxruntime_extensions/pnp/_imagenet.py

@@ -0,0 +1,65 @@
+import torch
+from typing import Tuple
+from torch.nn.functional import interpolate
+from ._base import ProcessingTracedModule
+from ._torchext import onnx_where, onnx_greater
+
+
+def _resize_param(img, size):
+    y, x = tuple(img.shape[-2:])
+    scale_y = size / y
+    scale_x = size / x
+    return onnx_where(onnx_greater(scale_x, scale_y), scale_x, scale_y)
+
+
+class ImageNetPreProcessing(ProcessingTracedModule):
+    def __init__(self, size, resize_image=True):
+        super(ImageNetPreProcessing, self).__init__()
+        self.target_size = size
+        self.resize_image = resize_image
+
+    def forward(self, img):
+        if not isinstance(img, torch.Tensor):
+            img = torch.tensor(img)
+        assert img.shape[-1] == 3, 'the input image should be in RGB channels'
+        img = torch.permute(img, (2, 0, 1))
+        img = img.to(torch.float32).unsqueeze(0)
+        # T.Resize(256),
+        if self.resize_image:
+            scale = _resize_param(img, torch.tensor(256))
+            img = interpolate(img, scale_factor=scale,
+                              recompute_scale_factor=True,
+                              mode="bilinear", align_corners=False)
+        # T.CenterCrop(224),
+        width, height = self.target_size, self.target_size
+        img_h, img_w = img.shape[-2:]
+        s_h = torch.div((img_h - height), 2, rounding_mode='trunc')
+        s_w = torch.div((img_w - width), 2, rounding_mode='trunc')
+        x = img[:, :, s_h:s_h + height, s_w:s_w + width]
+        # T.ToTensor(),
+        x /= 255.  # ToTensor
+        # T.Normalize(
+        #     mean=[0.485, 0.456, 0.406],
+        #     std=[0.229, 0.224, 0.225]
+        # )
+        mean = [0.485, 0.456, 0.406]
+        std = [0.229, 0.224, 0.225]
+        x -= torch.reshape(torch.tensor(mean), (3, 1, 1))
+        x /= torch.reshape(torch.tensor(std), (3, 1, 1))
+        return x
+
+
+class ImageNetPostProcessing(ProcessingTracedModule):
+    def forward(self, scores: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        probabilities = torch.softmax(scores, dim=1)
+        top10_prob, top10_ids = probabilities.topk(k=10, dim=1, largest=True, sorted=True)
+        return top10_ids, top10_prob
+
+
+class PreMobileNet(ImageNetPreProcessing):
+    def __init__(self, size=None):
+        super(PreMobileNet, self).__init__(224 if size is None else size)
+
+
+class PostMobileNet(ImageNetPostProcessing):
+    pass

onnxruntime_extensions/pnp/_nlp.py

@@ -0,0 +1,148 @@
+import json
+from collections import OrderedDict
+
+from ._base import ProcessingTracedModule, tensor_data_type as _dt
+from ._torchext import create_op_function
+from ._onnx_ops import schema
+from .._ocos import default_opset_domain
+
+
+def make_custom_op(ctx, op_type, input_names, output_names, container, operator_name=None, **kwargs):
+    op_name = container.get_unique_operator_name(op_type) if operator_name is None else operator_name
+    container.add_node(op_type, input_names, output_names,
+                       op_version=1, name=op_name, op_domain=default_opset_domain(), **kwargs)
+
+
+def create_bert_tokenizer(ctx, name, input_names, output_names, container, operator_name=None, **kwargs):
+    if 'hf_tok' in kwargs:
+        hf_bert_tokenizer = kwargs['hf_tok']
+        ordered_vocab = OrderedDict(sorted(hf_bert_tokenizer.vocab.items(), key=lambda item: int(item[1])))
+        vocab = '\n'.join(ordered_vocab.keys())
+        attrs = dict(vocab_file=vocab)
+        # Unfortunately, there's no specific accessor function on
+        # transformers.BertTokenizer to query for strip_accents.
+        attrs['strip_accents'] = 1 if 'strip_accents' in hf_bert_tokenizer.init_kwargs and hf_bert_tokenizer.init_kwargs.get('strip_accents') else 0
+        attrs['do_lower_case'] = 1 if hasattr(hf_bert_tokenizer, 'do_lower_case') and hf_bert_tokenizer.do_lower_case else 0
+    elif 'vocab_file' in kwargs:
+        vocab = None
+        vocab_file = kwargs['vocab_file']
+        with open(vocab_file, "r", encoding='utf-8') as vf:
+            lines = vf.readlines()
+            vocab = '\n'.join(lines)
+        if vocab is None:
+            raise RuntimeError("Cannot load vocabulary file {}!".format(vocab_file))
+        attrs = dict(vocab_file=vocab)
+        if 'strip_accents' in kwargs:
+            attrs['strip_accents'] = kwargs['strip_accents']
+        if 'do_lower_case' in kwargs:
+            attrs['do_lower_case'] = kwargs['do_lower_case']
+    else:
+        raise RuntimeError("Need hf_tok/vocab_file parameter to build the tokenizer")
+
+    return make_custom_op(ctx, name, input_names,
+                          output_names, container, operator_name=operator_name, **attrs)
+
+
+@schema(inputs=((_dt.STRING, []),),
+        outputs=((_dt.INT64, []), (_dt.INT64, []), (_dt.INT64, [])))
+def bert_tokenizer(ctx, input_names, output_names, container, operator_name=None, **kwargs):
+    return create_bert_tokenizer(ctx, 'BertTokenizer', input_names, output_names,
+                                 container, operator_name=operator_name, **kwargs)
+
+
+@schema(inputs=((_dt.STRING, []),),
+        outputs=((_dt.INT64, []), (_dt.INT64, []), (_dt.INT64, [])))
+def hf_bert_tokenizer(ctx, input_names, output_names, container, operator_name=None, **kwargs):
+    return create_bert_tokenizer(ctx, 'HfBertTokenizer', input_names, output_names,
+                                 container, operator_name=operator_name, **kwargs)
+
+
+@schema(inputs=((_dt.STRING, []),),
+        outputs=((_dt.INT64, []), (_dt.INT64, [])))
+def gpt2_tokenize(ctx, input_names, output_names, container, operator_name=None, **kwargs):
+    if 'hf_tok' in kwargs:
+        hf_gpt2_tokenizer = kwargs['hf_tok']
+        attrs = {'vocab': json.dumps(hf_gpt2_tokenizer.encoder, separators=(',', ':'))}
+        sorted_merges = {v_: k_ for k_, v_ in hf_gpt2_tokenizer.bpe_ranks.items()}
+        attrs['merges'] = '\n'.join("{} {}".format(*sorted_merges[n_]) for n_ in range(len(sorted_merges)))
+    elif 'vocab' in kwargs:
+        attrs = dict(
+            vocab=kwargs['vocab'],
+            merges=kwargs['merges'])
+    else:
+        raise RuntimeError("Need hf_tok/vocab parameter to build the tokenizer")
+    padding_len = -1
+    if 'padding_length' in kwargs:
+        padding_len = kwargs['padding_length']
+    attrs['padding_length'] = padding_len
+
+    return make_custom_op(ctx, 'GPT2Tokenizer', input_names,
+                          output_names, container, operator_name=operator_name, **attrs)
+
+
+def _get_file_content(path):
+    with open(path, "rb") as file:
+        return file.read()
+
+
+def _get_bound_object(func):
+    return func.__self__
+
+# v1. Order of outputs - input_ids, token_type_ids, attention_mask
+#    (this is NOT consistent with the HuggingFace implementation of the tokenizer)
+class PreHuggingFaceBert(ProcessingTracedModule):
+    def __init__(self, hf_tok=None, vocab_file=None, do_lower_case=0, strip_accents=1):
+        super(PreHuggingFaceBert, self).__init__()
+        if hf_tok is None:
+            self.onnx_bert_tokenizer = create_op_function('BertTokenizer', bert_tokenizer,
+                                                          vocab_file=vocab_file,
+                                                          do_lower_case=do_lower_case,
+                                                          strip_accents=strip_accents)
+        else:
+            self.onnx_bert_tokenizer = create_op_function('BertTokenizer', bert_tokenizer,
+                                                          hf_tok=hf_tok)
+
+    def forward(self, text):
+        return self.onnx_bert_tokenizer(text)
+
+    def export(self, *args, **kwargs):
+        return _get_bound_object(self.onnx_bert_tokenizer).build_model(kwargs.get('opset_version', 0), *args)
+
+
+# v2. Order of outputs - input_ids, attention_mask, token_type_ids
+#    (this is consistent with the HuggingFace implementation of the tokenizer)
+class HfBertTokenizer(ProcessingTracedModule):
+    def __init__(self, hf_tok=None, vocab_file=None, do_lower_case=0, strip_accents=1):
+        super(HfBertTokenizer, self).__init__()
+        if hf_tok is None:
+            self.onnx_bert_tokenizer = create_op_function('HfBertTokenizer', hf_bert_tokenizer,
+                                                          vocab_file=vocab_file,
+                                                          do_lower_case=do_lower_case,
+                                                          strip_accents=strip_accents)
+        else:
+            self.onnx_bert_tokenizer = create_op_function('HfBertTokenizer', hf_bert_tokenizer,
+                                                          hf_tok=hf_tok)
+
+    def forward(self, text):
+        return self.onnx_bert_tokenizer(text)
+
+    def export(self, *args, **kwargs):
+        return _get_bound_object(self.onnx_bert_tokenizer).build_model(kwargs.get('opset_version', 0), *args)
+
+
+class PreHuggingFaceGPT2(ProcessingTracedModule):
+    def __init__(self, hf_tok=None, vocab_file=None, merges_file=None, padding_length=-1):
+        super(PreHuggingFaceGPT2, self).__init__()
+        if hf_tok is None:
+            self.onnx_gpt2_tokenize = create_op_function('GPT2Tokenizer', gpt2_tokenize,
+                                                         vocab=_get_file_content(vocab_file),
+                                                         merges=_get_file_content(merges_file),
+                                                         padding_length=padding_length)
+        else:
+            self.onnx_gpt2_tokenize = create_op_function('GPT2Tokenizer', gpt2_tokenize, hf_tok=hf_tok)
+
+    def forward(self, text):
+        return self.onnx_gpt2_tokenize(text)
+
+    def export(self, *args, **kwargs):
+        return _get_bound_object(self.onnx_gpt2_tokenize).build_model(kwargs.get('opset_version', 0), *args)