GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

microsoft/onnxruntime-extensions

Public

mirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
v0.3.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

onnxruntime-extensions/onnxruntime_extensions

onnxruntime_extensions/_cuops.py

124lines · modecode

RawDownload
Latest commit unavailable.
1# Copyright (c) Microsoft Corporation. All rights reserved.
2# Licensed under the MIT License. See License.txt in the project root for
3# license information.
4###############################################################################
5
6import onnx
7from onnx import onnx_pb as onnx_proto
8from ._ocos import default_opset_domain, get_library_path # noqa
9
10
11class CustomOp:
12 @classmethod
13 def op_type(cls):
14 rcls = cls
15 while CustomOp != rcls.__base__:
16 rcls = rcls.__base__
17 return rcls.__name__
18
19 @classmethod
20 def get_inputs(cls): return None
21
22 @classmethod
23 def get_output(cls): return None
24
25 @classmethod
26 def serialize_attr(cls, attrs):
27 """
28 Only support serialize the basic python type like list or dict,
29 All other types needs to be serialized by the users
30 :param attrs: the dict attributes
31 :return: the dict of serialized data
32 """
33 return attrs
34
35 io_def = onnx.helper.make_tensor_value_info
36
37
38class GPT2Tokenizer(CustomOp):
39 @classmethod
40 def get_inputs(cls):
41 return [cls.io_def('input_text', onnx_proto.TensorProto.STRING, [None])]
42
43 @classmethod
44 def get_outputs(cls):
45 return [cls.io_def("input_ids", onnx.TensorProto.INT64, [None, None]),
46 cls.io_def('attention_mask', onnx.TensorProto.INT64, [None, None])]
47
48
49class VectorToString(CustomOp):
50 @classmethod
51 def get_inputs(cls):
52 return [cls.io_def("token_ids", onnx.TensorProto.INT64, [])]
53
54 @classmethod
55 def get_outputs(cls):
56 return [cls.io_def('text', onnx_proto.TensorProto.STRING, [])]
57
58 @classmethod
59 def serialize_attr(cls, attrs):
60 attr_data = {}
61 for k_, v_ in attrs.items():
62 if k_ == 'map' and isinstance(v_, dict):
63 attr_data[k_] = '\n'.join(k + "\t" + " ".join([str(i) for i in v]) for k, v in v_.items())
64 else:
65 attr_data[k_] = v_
66 return attr_data
67
68
69class StringToVector(CustomOp):
70 @classmethod
71 def get_inputs(cls):
72 return [cls.io_def("text", onnx.TensorProto.STRING, [None])]
73
74 @classmethod
75 def get_outputs(cls):
76 return [cls.io_def('token_ids', onnx_proto.TensorProto.INT64, [])]
77
78 @classmethod
79 def serialize_attr(cls, attrs):
80 attr_data = {}
81 for k_, v_ in attrs.items():
82 if k_ == 'map' and isinstance(v_, dict):
83 attr_data[k_] = '\n'.join(k + "\t" + " ".join([str(i) for i in v]) for k, v in v_.items())
84 elif k_ == 'unk' and isinstance(v_, list):
85 attr_data[k_] = ' '.join(str(i) for i in v_)
86 else:
87 attr_data[k_] = v_
88 return attr_data
89
90
91# TODO: list all custom operators schema here:
92# ...
93# ...
94
95
96class SingleOpGraph:
97 @classmethod
98 def get_next_id(cls):
99 if not hasattr(cls, '_id_counter'):
100 cls._id_counter = 0
101 cls._id_counter += 1
102 return cls._id_counter
103
104 @classmethod
105 def build_my_graph(cls, op_class, *args, **kwargs):
106 op_type = op_class.op_type()
107 inputs = op_class.get_inputs()
108 outputs = op_class.get_outputs()
109 attrs = op_class.serialize_attr(kwargs)
110 cuop = onnx.helper.make_node(op_type,
111 [i_.name for i_ in inputs],
112 [o_.name for o_ in outputs],
113 "{}_{}".format(op_type, cls.get_next_id()),
114 **attrs,
115 domain=default_opset_domain())
116 graph = onnx.helper.make_graph([cuop],
117 "og_{}_{}".format(op_type, cls.get_next_id()),
118 inputs,
119 outputs)
120 return graph
121
122 @staticmethod
123 def get_op_class(op_type):
124 return globals()[op_type]
125