microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
test/test_string_ops.py
351lines · modecode
| 1 | # coding: utf-8 |
| 2 | import unittest |
| 3 | import re |
| 4 | import numpy as np |
| 5 | from onnx import helper, onnx_pb as onnx_proto |
| 6 | import onnxruntime as _ort |
| 7 | from ortcustomops import ( |
| 8 | onnx_op, PyCustomOpDef, |
| 9 | get_library_path as _get_library_path) |
| 10 | |
| 11 | |
| 12 | def _create_test_model_string_upper(prefix, domain='ai.onnx.contrib'): |
| 13 | nodes = [] |
| 14 | nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])] |
| 15 | nodes[1:] = [helper.make_node('%sStringUpper' % prefix, |
| 16 | ['identity1'], ['customout'], |
| 17 | domain=domain)] |
| 18 | |
| 19 | input0 = helper.make_tensor_value_info( |
| 20 | 'input_1', onnx_proto.TensorProto.STRING, [None, None]) |
| 21 | output0 = helper.make_tensor_value_info( |
| 22 | 'customout', onnx_proto.TensorProto.STRING, [None, None]) |
| 23 | |
| 24 | graph = helper.make_graph(nodes, 'test0', [input0], [output0]) |
| 25 | model = helper.make_model( |
| 26 | graph, opset_imports=[helper.make_operatorsetid(domain, 1)]) |
| 27 | return model |
| 28 | |
| 29 | |
| 30 | def _create_test_model_string_join(prefix, domain='ai.onnx.contrib'): |
| 31 | nodes = [] |
| 32 | nodes.append( |
| 33 | helper.make_node('Identity', ['text'], ['identity1'])) |
| 34 | nodes.append( |
| 35 | helper.make_node('Identity', ['sep'], ['identity2'])) |
| 36 | nodes.append( |
| 37 | helper.make_node('Identity', ['axis'], ['identity3'])) |
| 38 | nodes.append( |
| 39 | helper.make_node( |
| 40 | '%sStringJoin' % prefix, ['identity1', 'identity2', 'identity3'], |
| 41 | ['customout'], domain=domain)) |
| 42 | |
| 43 | input0 = helper.make_tensor_value_info( |
| 44 | 'text', onnx_proto.TensorProto.STRING, None) |
| 45 | input1 = helper.make_tensor_value_info( |
| 46 | 'sep', onnx_proto.TensorProto.STRING, [1]) |
| 47 | input2 = helper.make_tensor_value_info( |
| 48 | 'axis', onnx_proto.TensorProto.INT64, [1]) |
| 49 | output0 = helper.make_tensor_value_info( |
| 50 | 'customout', onnx_proto.TensorProto.STRING, None) |
| 51 | |
| 52 | graph = helper.make_graph( |
| 53 | nodes, 'test0', [input0, input1, input2], [output0]) |
| 54 | model = helper.make_model( |
| 55 | graph, opset_imports=[helper.make_operatorsetid(domain, 1)]) |
| 56 | return model |
| 57 | |
| 58 | |
| 59 | def _create_test_model_string_replace(prefix, domain='ai.onnx.contrib'): |
| 60 | nodes = [] |
| 61 | nodes.append( |
| 62 | helper.make_node('Identity', ['text'], ['id1'])) |
| 63 | nodes.append( |
| 64 | helper.make_node('Identity', ['pattern'], ['id2'])) |
| 65 | nodes.append( |
| 66 | helper.make_node('Identity', ['rewrite'], ['id3'])) |
| 67 | nodes.append( |
| 68 | helper.make_node( |
| 69 | '%sStringRegexReplace' % prefix, ['id1', 'id2', 'id3'], |
| 70 | ['customout'], domain=domain)) |
| 71 | |
| 72 | input0 = helper.make_tensor_value_info( |
| 73 | 'text', onnx_proto.TensorProto.STRING, [None, 1]) |
| 74 | input1 = helper.make_tensor_value_info( |
| 75 | 'pattern', onnx_proto.TensorProto.STRING, [1]) |
| 76 | input2 = helper.make_tensor_value_info( |
| 77 | 'rewrite', onnx_proto.TensorProto.STRING, [1]) |
| 78 | output0 = helper.make_tensor_value_info( |
| 79 | 'customout', onnx_proto.TensorProto.STRING, [None, 1]) |
| 80 | |
| 81 | graph = helper.make_graph( |
| 82 | nodes, 'test0', [input0, input1, input2], [output0]) |
| 83 | model = helper.make_model( |
| 84 | graph, opset_imports=[helper.make_operatorsetid(domain, 1)]) |
| 85 | return model |
| 86 | |
| 87 | |
| 88 | class TestPythonOpString(unittest.TestCase): |
| 89 | |
| 90 | _string_join = None |
| 91 | |
| 92 | @classmethod |
| 93 | def setUpClass(cls): |
| 94 | |
| 95 | @onnx_op(op_type="PyStringUpper", |
| 96 | inputs=[PyCustomOpDef.dt_string], |
| 97 | outputs=[PyCustomOpDef.dt_string]) |
| 98 | def string_upper(x): |
| 99 | # The user custom op implementation here. |
| 100 | return np.array([s.upper() for s in x.ravel()]).reshape(x.shape) |
| 101 | |
| 102 | @onnx_op(op_type="PyStringJoin", |
| 103 | inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_string, |
| 104 | PyCustomOpDef.dt_int64], |
| 105 | outputs=[PyCustomOpDef.dt_string]) |
| 106 | def string_join(x, sep, axis): |
| 107 | # The user custom op implementation here. |
| 108 | if sep.shape != (1, ): |
| 109 | raise RuntimeError( |
| 110 | "Unexpected shape {} for 'sep'.".format(sep.shape)) |
| 111 | if axis.shape != (1, ): |
| 112 | raise RuntimeError( |
| 113 | "Unexpected shape {} for 'axis'.".format(axis.shape)) |
| 114 | sp = sep[0] |
| 115 | ax = axis[0] |
| 116 | if ax < 0 or ax >= len(x.shape): |
| 117 | raise RuntimeError("axis must be in [%r,%r] but is" % ( |
| 118 | 0, len(x.shape), ax)) |
| 119 | if len(x.shape) == 1: |
| 120 | return np.array([sp.join(x)]) |
| 121 | dims = np.arange(len(x.shape)) |
| 122 | dims[ax], dims[-1] = dims[-1], dims[ax] |
| 123 | x2 = np.transpose(x, dims) |
| 124 | res_shape = x2.shape[:-1] |
| 125 | x2 = x2.reshape((-1, x2.shape[-1])) |
| 126 | res = np.empty(x2.shape[0], dtype=x.dtype) |
| 127 | for i in range(x2.shape[0]): |
| 128 | res[i] = sp.join(x2[i, :]) |
| 129 | return res.reshape(res_shape) |
| 130 | |
| 131 | @onnx_op(op_type="PyStringRegexReplace", |
| 132 | inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_string, |
| 133 | PyCustomOpDef.dt_string], |
| 134 | outputs=[PyCustomOpDef.dt_string]) |
| 135 | def string_replace(x, pattern, rewrite): |
| 136 | # The user custom op implementation here. |
| 137 | if pattern.shape != (1, ): |
| 138 | raise RuntimeError( |
| 139 | "Unexpected shape {} for 'pattern'.".format(pattern.shape)) |
| 140 | if rewrite.shape != (1, ): |
| 141 | raise RuntimeError( |
| 142 | "Unexpected shape {} for 'rewrite'.".format(rewrite.shape)) |
| 143 | reg = re.compile(pattern[0]) |
| 144 | res = np.array( |
| 145 | list(map(lambda t: reg.sub(rewrite[0], t), x.ravel()))) |
| 146 | return res.reshape(x.shape) |
| 147 | |
| 148 | cls._string_join = string_join |
| 149 | |
| 150 | def test_check_types(self): |
| 151 | def_list = set(dir(PyCustomOpDef)) |
| 152 | type_list = [ |
| 153 | # 'dt_bfloat16', |
| 154 | 'dt_bool', |
| 155 | 'dt_complex128', |
| 156 | 'dt_complex64', |
| 157 | 'dt_double', |
| 158 | 'dt_float', |
| 159 | 'dt_float16', |
| 160 | 'dt_int16', |
| 161 | 'dt_int32', |
| 162 | 'dt_int64', |
| 163 | 'dt_int8', |
| 164 | 'dt_string', |
| 165 | 'dt_uint16', |
| 166 | 'dt_uint32', |
| 167 | 'dt_uint64', |
| 168 | 'dt_uint8'] |
| 169 | for t in type_list: |
| 170 | self.assertIn(t, def_list) |
| 171 | |
| 172 | def test_string_upper_cc(self): |
| 173 | so = _ort.SessionOptions() |
| 174 | so.register_custom_ops_library(_get_library_path()) |
| 175 | onnx_model = _create_test_model_string_upper('') |
| 176 | self.assertIn('op_type: "StringUpper"', str(onnx_model)) |
| 177 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 178 | input_1 = np.array([["Abc"]]) |
| 179 | txout = sess.run(None, {'input_1': input_1}) |
| 180 | self.assertEqual(txout[0].tolist(), np.array([["ABC"]]).tolist()) |
| 181 | |
| 182 | def test_string_upper_cc_accent(self): |
| 183 | so = _ort.SessionOptions() |
| 184 | so.register_custom_ops_library(_get_library_path()) |
| 185 | onnx_model = _create_test_model_string_upper('') |
| 186 | self.assertIn('op_type: "StringUpper"', str(onnx_model)) |
| 187 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 188 | input_1 = np.array([["Abcé"]]) |
| 189 | txout = sess.run(None, {'input_1': input_1}) |
| 190 | self.assertEqual(txout[0].tolist(), np.array([["ABCé"]]).tolist()) |
| 191 | |
| 192 | def test_string_upper_python(self): |
| 193 | so = _ort.SessionOptions() |
| 194 | so.register_custom_ops_library(_get_library_path()) |
| 195 | onnx_model = _create_test_model_string_upper('Py') |
| 196 | self.assertIn('op_type: "PyStringUpper"', str(onnx_model)) |
| 197 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 198 | input_1 = np.array([["Abc"]]) |
| 199 | txout = sess.run(None, {'input_1': input_1}) |
| 200 | self.assertEqual(txout[0].tolist(), np.array([["ABC"]]).tolist()) |
| 201 | |
| 202 | def test_string_upper_python_accent(self): |
| 203 | so = _ort.SessionOptions() |
| 204 | so.register_custom_ops_library(_get_library_path()) |
| 205 | onnx_model = _create_test_model_string_upper('Py') |
| 206 | self.assertIn('op_type: "PyStringUpper"', str(onnx_model)) |
| 207 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 208 | input_1 = np.array([["Abcé"]]) |
| 209 | txout = sess.run(None, {'input_1': input_1}) |
| 210 | self.assertEqual(txout[0].tolist(), |
| 211 | np.array([["ABCé".upper()]]).tolist()) |
| 212 | |
| 213 | def test_string_join_python(self): |
| 214 | so = _ort.SessionOptions() |
| 215 | so.register_custom_ops_library(_get_library_path()) |
| 216 | onnx_model = _create_test_model_string_join('Py') |
| 217 | self.assertIn('op_type: "PyStringJoin"', str(onnx_model)) |
| 218 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 219 | text = np.vstack([np.array([["a", "b", "c"]]), |
| 220 | np.array([["aa", "bb", ""]])]) |
| 221 | self.assertEqual(text.shape, (2, 3)) |
| 222 | sep = np.array([";"]) |
| 223 | axis = np.array([1], dtype=np.int64) |
| 224 | TestPythonOpString._string_join(text, sep, axis) |
| 225 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 226 | self.assertEqual( |
| 227 | txout[0].tolist(), np.array(["a;b;c", "aa;bb;"]).tolist()) |
| 228 | axis = np.array([0], dtype=np.int64) |
| 229 | TestPythonOpString._string_join(text, sep, axis) |
| 230 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 231 | self.assertEqual( |
| 232 | txout[0].tolist(), np.array(['a;aa', 'b;bb', 'c;']).tolist()) |
| 233 | |
| 234 | def test_string_join_python_3d(self): |
| 235 | so = _ort.SessionOptions() |
| 236 | so.register_custom_ops_library(_get_library_path()) |
| 237 | onnx_model = _create_test_model_string_join('Py') |
| 238 | self.assertIn('op_type: "PyStringJoin"', str(onnx_model)) |
| 239 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 240 | text = np.vstack([np.array([["a", "b", "c"]]), |
| 241 | np.array([["aa", "bb", ""]])]).reshape((2, 3, 1)) |
| 242 | sep = np.array([";"]) |
| 243 | axis = np.array([1], dtype=np.int64) |
| 244 | TestPythonOpString._string_join(text, sep, axis) |
| 245 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 246 | self.assertEqual( |
| 247 | txout[0].tolist(), np.array([['a;b;c'], ['aa;bb;']]).tolist()) |
| 248 | |
| 249 | def test_string_join_python_1d(self): |
| 250 | so = _ort.SessionOptions() |
| 251 | so.register_custom_ops_library(_get_library_path()) |
| 252 | onnx_model = _create_test_model_string_join('Py') |
| 253 | self.assertIn('op_type: "PyStringJoin"', str(onnx_model)) |
| 254 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 255 | text = np.array(["a", "b", "cc"]) |
| 256 | sep = np.array([";"]) |
| 257 | axis = np.array([0], dtype=np.int64) |
| 258 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 259 | self.assertEqual(txout[0].shape, (1, )) |
| 260 | self.assertEqual( |
| 261 | txout[0].tolist(), np.array(["a;b;cc"]).tolist()) |
| 262 | |
| 263 | def test_string_join_cc(self): |
| 264 | so = _ort.SessionOptions() |
| 265 | so.register_custom_ops_library(_get_library_path()) |
| 266 | onnx_model = _create_test_model_string_join('') |
| 267 | self.assertIn('op_type: "StringJoin"', str(onnx_model)) |
| 268 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 269 | text = np.vstack([np.array([["a", "b", "c"]]), |
| 270 | np.array([["aa", "bb", ""]])]) |
| 271 | sep = np.array([";"]) |
| 272 | axis = np.array([1], dtype=np.int64) |
| 273 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 274 | self.assertEqual( |
| 275 | txout[0].tolist(), np.array(["a;b;c", "aa;bb;"]).tolist()) |
| 276 | axis = np.array([0], dtype=np.int64) |
| 277 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 278 | self.assertEqual( |
| 279 | txout[0].tolist(), np.array(['a;aa', 'b;bb', 'c;']).tolist()) |
| 280 | |
| 281 | def test_string_join_cc_1d(self): |
| 282 | so = _ort.SessionOptions() |
| 283 | so.register_custom_ops_library(_get_library_path()) |
| 284 | onnx_model = _create_test_model_string_join('') |
| 285 | self.assertIn('op_type: "StringJoin"', str(onnx_model)) |
| 286 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 287 | text = np.array(["a", "b", "cc"]) |
| 288 | sep = np.array([";"]) |
| 289 | axis = np.array([0], dtype=np.int64) |
| 290 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 291 | self.assertEqual( |
| 292 | txout[0].tolist(), np.array(["a;b;cc"]).tolist()) |
| 293 | |
| 294 | def test_string_join_cc_3d(self): |
| 295 | so = _ort.SessionOptions() |
| 296 | so.register_custom_ops_library(_get_library_path()) |
| 297 | onnx_model = _create_test_model_string_join('') |
| 298 | self.assertIn('op_type: "StringJoin"', str(onnx_model)) |
| 299 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 300 | text = np.array(["a", "b", "c", "d", "e", "f", "g", "h"]).reshape(( |
| 301 | 2, 2, 2)) |
| 302 | sep = np.array([";"]) |
| 303 | axis = np.array([2], dtype=np.int64) |
| 304 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 305 | self.assertEqual( |
| 306 | txout[0].tolist(), |
| 307 | np.array([['a;b', 'c;d'], ['e;f', 'g;h']]).tolist()) |
| 308 | axis = np.array([1], dtype=np.int64) |
| 309 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 310 | self.assertEqual( |
| 311 | txout[0].tolist(), |
| 312 | np.array([['a;c', 'b;d'], ['e;g', 'f;h']]).tolist()) |
| 313 | axis = np.array([0], dtype=np.int64) |
| 314 | txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis}) |
| 315 | self.assertEqual( |
| 316 | txout[0].tolist(), |
| 317 | np.array([['a;e', 'b;f'], ['c;g', 'd;h']]).tolist()) |
| 318 | |
| 319 | def test_string_replace_cc(self): |
| 320 | so = _ort.SessionOptions() |
| 321 | so.register_custom_ops_library(_get_library_path()) |
| 322 | onnx_model = _create_test_model_string_replace('') |
| 323 | self.assertIn('op_type: "StringRegexReplace"', str(onnx_model)) |
| 324 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 325 | pattern = np.array([r'def\s+([a-zA-Z_][a-zA-Z_0-9]*)\s*\(\s*\):']) |
| 326 | rewrite = np.array([r'static PyObject* py_\1(void) {']) |
| 327 | text = np.array([['def myfunc():'], ['def dummy():']]) |
| 328 | txout = sess.run( |
| 329 | None, {'text': text, 'pattern': pattern, 'rewrite': rewrite}) |
| 330 | exp = [['static PyObject* py_myfunc(void) {'], |
| 331 | ['static PyObject* py_dummy(void) {']] |
| 332 | self.assertEqual(exp, txout[0].tolist()) |
| 333 | |
| 334 | def test_string_replace_python(self): |
| 335 | so = _ort.SessionOptions() |
| 336 | so.register_custom_ops_library(_get_library_path()) |
| 337 | onnx_model = _create_test_model_string_replace('Py') |
| 338 | self.assertIn('op_type: "PyStringRegexReplace"', str(onnx_model)) |
| 339 | sess = _ort.InferenceSession(onnx_model.SerializeToString(), so) |
| 340 | pattern = np.array([r'def\s+([a-zA-Z_][a-zA-Z_0-9]*)\s*\(\s*\):']) |
| 341 | rewrite = np.array([r'static PyObject*\npy_\1(void)\n{']) |
| 342 | text = np.array([['def myfunc():'], ['def dummy():']]) |
| 343 | txout = sess.run( |
| 344 | None, {'text': text, 'pattern': pattern, 'rewrite': rewrite}) |
| 345 | exp = [['static PyObject*\npy_myfunc(void)\n{'], |
| 346 | ['static PyObject*\npy_dummy(void)\n{']] |
| 347 | self.assertEqual(exp, txout[0].tolist()) |
| 348 | |
| 349 | |
| 350 | if __name__ == "__main__": |
| 351 | unittest.main() |