microsoft/onnxruntime-extensions

Public

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

CodeCommitsIssuesPull requestsActionsInsightsSecurity
fadcf2ab89ed9c28edb1f0ae4b529ceeac55f375

Branches

Tags

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

Clone

HTTPS

Download ZIP

test/test_string_ops.py

581lines · modecode

1# coding: utf-8
2import unittest
3import re
4from binascii import crc32
5import numpy as np
6from onnx import helper, onnx_pb as onnx_proto
7import onnxruntime as _ort
8from ortcustomops import (
9 onnx_op, PyCustomOpDef,
10 get_library_path as _get_library_path,
11 hash_64)
12
13NUM_BUCKETS = 23
14
15
16def _create_test_model_string_upper(prefix, domain='ai.onnx.contrib'):
17 nodes = []
18 nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])]
19 nodes[1:] = [helper.make_node('%sStringUpper' % prefix,
20 ['identity1'], ['customout'],
21 domain=domain)]
22
23 input0 = helper.make_tensor_value_info(
24 'input_1', onnx_proto.TensorProto.STRING, [None, None])
25 output0 = helper.make_tensor_value_info(
26 'customout', onnx_proto.TensorProto.STRING, [None, None])
27
28 graph = helper.make_graph(nodes, 'test0', [input0], [output0])
29 model = helper.make_model(
30 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
31 return model
32
33
34def _create_test_model_string_join(prefix, domain='ai.onnx.contrib'):
35 nodes = []
36 nodes.append(
37 helper.make_node('Identity', ['text'], ['identity1']))
38 nodes.append(
39 helper.make_node('Identity', ['sep'], ['identity2']))
40 nodes.append(
41 helper.make_node('Identity', ['axis'], ['identity3']))
42 nodes.append(
43 helper.make_node(
44 '%sStringJoin' % prefix, ['identity1', 'identity2', 'identity3'],
45 ['customout'], domain=domain))
46
47 input0 = helper.make_tensor_value_info(
48 'text', onnx_proto.TensorProto.STRING, None)
49 input1 = helper.make_tensor_value_info(
50 'sep', onnx_proto.TensorProto.STRING, [1])
51 input2 = helper.make_tensor_value_info(
52 'axis', onnx_proto.TensorProto.INT64, [1])
53 output0 = helper.make_tensor_value_info(
54 'customout', onnx_proto.TensorProto.STRING, None)
55
56 graph = helper.make_graph(
57 nodes, 'test0', [input0, input1, input2], [output0])
58 model = helper.make_model(
59 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
60 return model
61
62
63def _create_test_model_string_replace(prefix, domain='ai.onnx.contrib'):
64 nodes = []
65 nodes.append(
66 helper.make_node('Identity', ['text'], ['id1']))
67 nodes.append(
68 helper.make_node('Identity', ['pattern'], ['id2']))
69 nodes.append(
70 helper.make_node('Identity', ['rewrite'], ['id3']))
71 nodes.append(
72 helper.make_node(
73 '%sStringRegexReplace' % prefix, ['id1', 'id2', 'id3'],
74 ['customout'], domain=domain))
75
76 input0 = helper.make_tensor_value_info(
77 'text', onnx_proto.TensorProto.STRING, [None, 1])
78 input1 = helper.make_tensor_value_info(
79 'pattern', onnx_proto.TensorProto.STRING, [1])
80 input2 = helper.make_tensor_value_info(
81 'rewrite', onnx_proto.TensorProto.STRING, [1])
82 output0 = helper.make_tensor_value_info(
83 'customout', onnx_proto.TensorProto.STRING, [None, 1])
84
85 graph = helper.make_graph(
86 nodes, 'test0', [input0, input1, input2], [output0])
87 model = helper.make_model(
88 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
89 return model
90
91
92def _create_test_model_string_to_hash(
93 prefix, domain='ai.onnx.contrib', kind=None):
94 if kind == 'crc32':
95 op_type = 'StringToCRC32'
96 out_type = onnx_proto.TensorProto.UINT32
97 in_type = out_type
98 elif kind == 'hash_bucket':
99 op_type = 'StringToHashBucket'
100 out_type = onnx_proto.TensorProto.INT64
101 in_type = out_type
102 elif kind == 'hash_bucket_fast':
103 op_type = 'StringToHashBucketFast'
104 out_type = onnx_proto.TensorProto.INT64
105 in_type = out_type
106 else:
107 raise ValueError('Unknown value %r.' % kind)
108 nodes = []
109 nodes.append(
110 helper.make_node('Identity', ['text'], ['id1']))
111 nodes.append(
112 helper.make_node('Identity', ['num_buckets'], ['id2']))
113 nodes.append(
114 helper.make_node(
115 '%s%s' % (prefix, op_type), ['id1', 'id2'],
116 ['customout'], domain=domain))
117
118 input0 = helper.make_tensor_value_info(
119 'text', onnx_proto.TensorProto.STRING, [None, None])
120 input1 = helper.make_tensor_value_info(
121 'num_buckets', in_type, [1])
122 output0 = helper.make_tensor_value_info(
123 'customout', out_type, [None, None])
124
125 graph = helper.make_graph(
126 nodes, 'test0', [input0, input1], [output0])
127 model = helper.make_model(
128 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
129 return model
130
131
132def _create_test_model_string_equal(prefix, domain='ai.onnx.contrib'):
133 nodes = []
134 nodes.append(helper.make_node('Identity', ['x'], ['id1']))
135 nodes.append(helper.make_node('Identity', ['y'], ['id2']))
136 nodes.append(
137 helper.make_node(
138 '%sStringEqual' % prefix, ['id1', 'id2'], ['z'], domain=domain))
139
140 input0 = helper.make_tensor_value_info(
141 'x', onnx_proto.TensorProto.STRING, [])
142 input1 = helper.make_tensor_value_info(
143 'y', onnx_proto.TensorProto.STRING, [])
144 output0 = helper.make_tensor_value_info(
145 'z', onnx_proto.TensorProto.BOOL, [])
146
147 graph = helper.make_graph(nodes, 'test0', [input0, input1], [output0])
148 model = helper.make_model(
149 graph, opset_imports=[helper.make_operatorsetid(domain, 1)])
150 return model
151
152
153class TestPythonOpString(unittest.TestCase):
154
155 _string_join = None
156 _string_to_crc32 = None
157
158 @classmethod
159 def setUpClass(cls):
160
161 @onnx_op(op_type="PyStringUpper",
162 inputs=[PyCustomOpDef.dt_string],
163 outputs=[PyCustomOpDef.dt_string])
164 def string_upper(x):
165 # The user custom op implementation here.
166 return np.array([s.upper() for s in x.ravel()]).reshape(x.shape)
167
168 @onnx_op(op_type="PyStringJoin",
169 inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_string,
170 PyCustomOpDef.dt_int64],
171 outputs=[PyCustomOpDef.dt_string])
172 def string_join(x, sep, axis):
173 # The user custom op implementation here.
174 if sep.shape != (1, ):
175 raise RuntimeError(
176 "Unexpected shape {} for 'sep'.".format(sep.shape))
177 if axis.shape != (1, ):
178 raise RuntimeError(
179 "Unexpected shape {} for 'axis'.".format(axis.shape))
180 sp = sep[0]
181 ax = axis[0]
182 if ax < 0 or ax >= len(x.shape):
183 raise RuntimeError(
184 "axis must be in [%r,%r] but is %r" % (
185 0, len(x.shape), ax))
186 if len(x.shape) == 1:
187 return np.array([sp.join(x)])
188 dims = np.arange(len(x.shape))
189 dims[ax], dims[-1] = dims[-1], dims[ax]
190 x2 = np.transpose(x, dims)
191 res_shape = x2.shape[:-1]
192 x2 = x2.reshape((-1, x2.shape[-1]))
193 res = np.empty(x2.shape[0], dtype=x.dtype)
194 for i in range(x2.shape[0]):
195 res[i] = sp.join(x2[i, :])
196 return res.reshape(res_shape)
197
198 @onnx_op(op_type="PyStringRegexReplace",
199 inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_string,
200 PyCustomOpDef.dt_string],
201 outputs=[PyCustomOpDef.dt_string])
202 def string_replace(x, pattern, rewrite):
203 # The user custom op implementation here.
204 if pattern.shape != (1, ):
205 raise RuntimeError(
206 "Unexpected shape {} for 'pattern'.".format(pattern.shape))
207 if rewrite.shape != (1, ):
208 raise RuntimeError(
209 "Unexpected shape {} for 'rewrite'.".format(rewrite.shape))
210 reg = re.compile(pattern[0])
211 res = np.array(
212 list(map(lambda t: reg.sub(rewrite[0], t), x.ravel())))
213 return res.reshape(x.shape)
214
215 @onnx_op(op_type="PyStringToCRC32",
216 inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_uint32],
217 outputs=[PyCustomOpDef.dt_uint32])
218 def string_to_crc32(x, num_buckets):
219 if num_buckets.shape != (1, ):
220 raise RuntimeError(
221 "Unexpected shape {} for 'num_buckets'.".format(
222 num_buckets.shape))
223 nb = num_buckets[0]
224 res = np.array(
225 list(map(
226 lambda x: crc32(x.encode('iso-8859-15')) % nb,
227 x.ravel())))
228 return res.reshape(x.shape)
229
230 @onnx_op(op_type="PyStringToHashBucket",
231 inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_int64],
232 outputs=[PyCustomOpDef.dt_int64])
233 def string_to_hash_bucket(x, num_buckets):
234 if num_buckets.shape != (1, ):
235 raise RuntimeError(
236 "Unexpected shape {} for 'num_buckets'.".format(
237 num_buckets.shape))
238 nb = num_buckets[0]
239 res = np.array(
240 list(map(lambda x: hash_64(x, nb, True), x.ravel())))
241 return res.reshape(x.shape).astype(np.int64)
242
243 @onnx_op(op_type="PyStringEqual",
244 inputs=[PyCustomOpDef.dt_string, PyCustomOpDef.dt_string],
245 outputs=[PyCustomOpDef.dt_bool])
246 def string_equal(x, y):
247 return x == y
248
249 cls._string_join = string_join
250 cls._string_to_crc32 = string_to_crc32
251
252 def test_check_types(self):
253 def_list = set(dir(PyCustomOpDef))
254 type_list = [
255 # 'dt_bfloat16',
256 'dt_bool',
257 'dt_complex128',
258 'dt_complex64',
259 'dt_double',
260 'dt_float',
261 'dt_float16',
262 'dt_int16',
263 'dt_int32',
264 'dt_int64',
265 'dt_int8',
266 'dt_string',
267 'dt_uint16',
268 'dt_uint32',
269 'dt_uint64',
270 'dt_uint8']
271 for t in type_list:
272 self.assertIn(t, def_list)
273
274 def test_string_upper_cc(self):
275 so = _ort.SessionOptions()
276 so.register_custom_ops_library(_get_library_path())
277 onnx_model = _create_test_model_string_upper('')
278 self.assertIn('op_type: "StringUpper"', str(onnx_model))
279 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
280 input_1 = np.array([["Abc"]])
281 txout = sess.run(None, {'input_1': input_1})
282 self.assertEqual(txout[0].tolist(), np.array([["ABC"]]).tolist())
283
284 def test_string_upper_cc_accent(self):
285 so = _ort.SessionOptions()
286 so.register_custom_ops_library(_get_library_path())
287 onnx_model = _create_test_model_string_upper('')
288 self.assertIn('op_type: "StringUpper"', str(onnx_model))
289 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
290 input_1 = np.array([["Abcé"]])
291 txout = sess.run(None, {'input_1': input_1})
292 self.assertEqual(txout[0].tolist(), np.array([["ABCé"]]).tolist())
293
294 def test_string_upper_python(self):
295 so = _ort.SessionOptions()
296 so.register_custom_ops_library(_get_library_path())
297 onnx_model = _create_test_model_string_upper('Py')
298 self.assertIn('op_type: "PyStringUpper"', str(onnx_model))
299 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
300 input_1 = np.array([["Abc"]])
301 txout = sess.run(None, {'input_1': input_1})
302 self.assertEqual(txout[0].tolist(), np.array([["ABC"]]).tolist())
303
304 def test_string_upper_python_accent(self):
305 so = _ort.SessionOptions()
306 so.register_custom_ops_library(_get_library_path())
307 onnx_model = _create_test_model_string_upper('Py')
308 self.assertIn('op_type: "PyStringUpper"', str(onnx_model))
309 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
310 input_1 = np.array([["Abcé"]])
311 txout = sess.run(None, {'input_1': input_1})
312 self.assertEqual(txout[0].tolist(),
313 np.array([["ABCé".upper()]]).tolist())
314
315 def test_string_join_python(self):
316 so = _ort.SessionOptions()
317 so.register_custom_ops_library(_get_library_path())
318 onnx_model = _create_test_model_string_join('Py')
319 self.assertIn('op_type: "PyStringJoin"', str(onnx_model))
320 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
321 text = np.vstack([np.array([["a", "b", "c"]]),
322 np.array([["aa", "bb", ""]])])
323 self.assertEqual(text.shape, (2, 3))
324 sep = np.array([";"])
325 axis = np.array([1], dtype=np.int64)
326 TestPythonOpString._string_join(text, sep, axis)
327 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
328 self.assertEqual(
329 txout[0].tolist(), np.array(["a;b;c", "aa;bb;"]).tolist())
330 axis = np.array([0], dtype=np.int64)
331 TestPythonOpString._string_join(text, sep, axis)
332 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
333 self.assertEqual(
334 txout[0].tolist(), np.array(['a;aa', 'b;bb', 'c;']).tolist())
335
336 def test_string_join_python_3d(self):
337 so = _ort.SessionOptions()
338 so.register_custom_ops_library(_get_library_path())
339 onnx_model = _create_test_model_string_join('Py')
340 self.assertIn('op_type: "PyStringJoin"', str(onnx_model))
341 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
342 text = np.vstack([np.array([["a", "b", "c"]]),
343 np.array([["aa", "bb", ""]])]).reshape((2, 3, 1))
344 sep = np.array([";"])
345 axis = np.array([1], dtype=np.int64)
346 TestPythonOpString._string_join(text, sep, axis)
347 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
348 self.assertEqual(
349 txout[0].tolist(), np.array([['a;b;c'], ['aa;bb;']]).tolist())
350
351 def test_string_join_python_1d(self):
352 so = _ort.SessionOptions()
353 so.register_custom_ops_library(_get_library_path())
354 onnx_model = _create_test_model_string_join('Py')
355 self.assertIn('op_type: "PyStringJoin"', str(onnx_model))
356 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
357 text = np.array(["a", "b", "cc"])
358 sep = np.array([";"])
359 axis = np.array([0], dtype=np.int64)
360 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
361 self.assertEqual(txout[0].shape, (1, ))
362 self.assertEqual(
363 txout[0].tolist(), np.array(["a;b;cc"]).tolist())
364
365 def test_string_join_cc(self):
366 so = _ort.SessionOptions()
367 so.register_custom_ops_library(_get_library_path())
368 onnx_model = _create_test_model_string_join('')
369 self.assertIn('op_type: "StringJoin"', str(onnx_model))
370 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
371 text = np.vstack([np.array([["a", "b", "c"]]),
372 np.array([["aa", "bb", ""]])])
373 sep = np.array([";"])
374 axis = np.array([1], dtype=np.int64)
375 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
376 self.assertEqual(
377 txout[0].tolist(), np.array(["a;b;c", "aa;bb;"]).tolist())
378 axis = np.array([0], dtype=np.int64)
379 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
380 self.assertEqual(
381 txout[0].tolist(), np.array(['a;aa', 'b;bb', 'c;']).tolist())
382
383 def test_string_join_cc_1d(self):
384 so = _ort.SessionOptions()
385 so.register_custom_ops_library(_get_library_path())
386 onnx_model = _create_test_model_string_join('')
387 self.assertIn('op_type: "StringJoin"', str(onnx_model))
388 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
389 text = np.array(["a", "b", "cc"])
390 sep = np.array([";"])
391 axis = np.array([0], dtype=np.int64)
392 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
393 self.assertEqual(
394 txout[0].tolist(), np.array(["a;b;cc"]).tolist())
395
396 def test_string_join_cc_3d(self):
397 so = _ort.SessionOptions()
398 so.register_custom_ops_library(_get_library_path())
399 onnx_model = _create_test_model_string_join('')
400 self.assertIn('op_type: "StringJoin"', str(onnx_model))
401 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
402 text = np.array(["a", "b", "c", "d", "e", "f", "g", "h"]).reshape((
403 2, 2, 2))
404 sep = np.array([";"])
405 axis = np.array([2], dtype=np.int64)
406 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
407 self.assertEqual(
408 txout[0].tolist(),
409 np.array([['a;b', 'c;d'], ['e;f', 'g;h']]).tolist())
410 axis = np.array([1], dtype=np.int64)
411 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
412 self.assertEqual(
413 txout[0].tolist(),
414 np.array([['a;c', 'b;d'], ['e;g', 'f;h']]).tolist())
415 axis = np.array([0], dtype=np.int64)
416 txout = sess.run(None, {'text': text, 'sep': sep, 'axis': axis})
417 self.assertEqual(
418 txout[0].tolist(),
419 np.array([['a;e', 'b;f'], ['c;g', 'd;h']]).tolist())
420
421 def test_string_replace_cc(self):
422 so = _ort.SessionOptions()
423 so.register_custom_ops_library(_get_library_path())
424 onnx_model = _create_test_model_string_replace('')
425 self.assertIn('op_type: "StringRegexReplace"', str(onnx_model))
426 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
427 pattern = np.array([r'def\s+([a-zA-Z_][a-zA-Z_0-9]*)\s*\(\s*\):'])
428 rewrite = np.array([r'static PyObject* py_\1(void) {'])
429 text = np.array([['def myfunc():'], ['def dummy():']])
430 txout = sess.run(
431 None, {'text': text, 'pattern': pattern, 'rewrite': rewrite})
432 exp = [['static PyObject* py_myfunc(void) {'],
433 ['static PyObject* py_dummy(void) {']]
434 self.assertEqual(exp, txout[0].tolist())
435
436 def test_string_replace_python(self):
437 so = _ort.SessionOptions()
438 so.register_custom_ops_library(_get_library_path())
439 onnx_model = _create_test_model_string_replace('Py')
440 self.assertIn('op_type: "PyStringRegexReplace"', str(onnx_model))
441 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
442 pattern = np.array([r'def\s+([a-zA-Z_][a-zA-Z_0-9]*)\s*\(\s*\):'])
443 rewrite = np.array([r'static PyObject*\npy_\1(void)\n{'])
444 text = np.array([['def myfunc():'], ['def dummy():']])
445 txout = sess.run(
446 None, {'text': text, 'pattern': pattern, 'rewrite': rewrite})
447 exp = [['static PyObject*\npy_myfunc(void)\n{'],
448 ['static PyObject*\npy_dummy(void)\n{']]
449 self.assertEqual(exp, txout[0].tolist())
450
451 def test_string_to_crc32_python(self):
452 so = _ort.SessionOptions()
453 so.register_custom_ops_library(_get_library_path())
454 onnx_model = _create_test_model_string_to_hash('Py', kind='crc32')
455 self.assertIn('op_type: "PyStringToCRC32"', str(onnx_model))
456 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
457 text = np.array([["abc", "abcdé"], ["$$^l!%*ù", ""]])
458 num_buckets = np.array([44], dtype=np.uint32)
459 res = self._string_to_crc32(text, num_buckets)
460 self.assertEqual(res.shape, text.shape)
461 exp = np.array([[10, 38], [29, 0]], dtype=np.uint32)
462 self.assertEqual(exp.tolist(), res.tolist())
463 txout = sess.run(
464 None, {'text': text, 'num_buckets': num_buckets})
465 self.assertEqual(exp.tolist(), txout[0].tolist())
466
467 def test_string_to_hash_bucket_cc(self):
468 so = _ort.SessionOptions()
469 so.register_custom_ops_library(_get_library_path())
470 onnx_model = _create_test_model_string_to_hash(
471 '', kind='hash_bucket')
472 self.assertIn('op_type: "StringToHashBucket"', str(onnx_model))
473 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
474 raw = ["abc", "abcdé", "$$^l!%*ù", "", "a", "A"]
475 text = np.array(raw).reshape((3, 2))
476 num_buckets = np.array([NUM_BUCKETS], dtype=np.int64)
477 txout = sess.run(
478 None, {'text': text, 'num_buckets': num_buckets})
479 try:
480 from tensorflow.raw_ops import StringToHashBucket
481 dotf = True
482 except ImportError:
483 dotf = False
484 if dotf:
485 tfres = StringToHashBucket(
486 string_tensor=text, num_buckets=num_buckets[0])
487 self.assertEqual(tfres.shape, txout[0].shape)
488 self.assertEqual(tfres.numpy().tolist(), txout[0].tolist())
489 exp = np.array([[15, 11], [10, 21], [20, 21]], dtype=np.int64)
490 self.assertEqual(exp.shape, txout[0].shape)
491 self.assertEqual(exp.tolist(), txout[0].tolist())
492
493 def test_string_to_hash_bucket_fast_cc(self):
494 so = _ort.SessionOptions()
495 so.register_custom_ops_library(_get_library_path())
496 onnx_model = _create_test_model_string_to_hash(
497 '', kind='hash_bucket_fast')
498 self.assertIn('op_type: "StringToHashBucketFast"', str(onnx_model))
499 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
500 raw = ["abc", "abcdé", "$$^l!%*ù", "", "a", "A"]
501 text = np.array(raw).reshape((3, 2))
502 num_buckets = np.array([NUM_BUCKETS], dtype=np.int64)
503 txout = sess.run(
504 None, {'text': text, 'num_buckets': num_buckets})
505 try:
506 from tensorflow.raw_ops import StringToHashBucketFast
507 dotf = True
508 except ImportError:
509 dotf = False
510 if dotf:
511 tfres = StringToHashBucketFast(
512 input=text, num_buckets=num_buckets[0])
513 self.assertEqual(tfres.shape, txout[0].shape)
514 self.assertEqual(tfres.numpy().tolist(), txout[0].tolist())
515 exp = np.array([[9, 17], [4, 21], [14, 12]], dtype=np.int64)
516 self.assertEqual(exp.shape, txout[0].shape)
517 self.assertEqual(exp.tolist(), txout[0].tolist())
518
519 def test_string_to_hash_bucket_python(self):
520 so = _ort.SessionOptions()
521 so.register_custom_ops_library(_get_library_path())
522 onnx_model = _create_test_model_string_to_hash(
523 'Py', kind='hash_bucket')
524 self.assertIn('op_type: "PyStringToHashBucket"', str(onnx_model))
525 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
526 raw = ["abc", "abcdé", "$$^l!%*ù", "", "a", "A"]
527 text = np.array(raw).reshape((3, 2))
528 num_buckets = np.array([NUM_BUCKETS], dtype=np.int64)
529 exp = np.array([[9, 17], [4, 21], [14, 12]], dtype=np.int64)
530 txout = sess.run(
531 None, {'text': text, 'num_buckets': num_buckets})
532 self.assertEqual(exp.shape, txout[0].shape)
533 self.assertEqual(exp.tolist(), txout[0].tolist())
534
535 def enumerate_matrix_couples(self):
536 for i in range(1, 5):
537 shape = (3,) * i
538 a = (np.random.rand(*shape) * 10).astype(np.int32).astype(np.str)
539 yield a, a
540 for j in range(i):
541 shape2 = list(shape)
542 shape2[j] = 1
543 b = (np.random.rand(*shape2) * 10).astype(
544 np.int32).astype(np.str)
545 yield a, b
546 for k in range(j+1, i):
547 shape3 = list(shape2)
548 shape3[k] = 1
549 b = (np.random.rand(*shape3) * 10).astype(
550 np.int32).astype(np.str)
551 yield a, b
552
553 def test_string_equal_python(self):
554 so = _ort.SessionOptions()
555 so.register_custom_ops_library(_get_library_path())
556 onnx_model = _create_test_model_string_equal('Py')
557 self.assertIn('op_type: "PyStringEqual"', str(onnx_model))
558 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
559
560 for x, y in self.enumerate_matrix_couples():
561 txout = sess.run(None, {'x': x, 'y': y})
562 self.assertEqual(txout[0].tolist(), (x == y).tolist())
563 txout = sess.run(None, {'x': y, 'y': x})
564 self.assertEqual(txout[0].tolist(), (y == x).tolist())
565
566 def test_string_equal_cc(self):
567 so = _ort.SessionOptions()
568 so.register_custom_ops_library(_get_library_path())
569 onnx_model = _create_test_model_string_equal('')
570 self.assertIn('op_type: "StringEqual"', str(onnx_model))
571 sess = _ort.InferenceSession(onnx_model.SerializeToString(), so)
572
573 for x, y in self.enumerate_matrix_couples():
574 txout = sess.run(None, {'x': x, 'y': y})
575 self.assertEqual(txout[0].tolist(), (x == y).tolist())
576 txout = sess.run(None, {'x': y, 'y': x})
577 self.assertEqual(txout[0].tolist(), (y == x).tolist())
578
579
580if __name__ == "__main__":
581 unittest.main()
582