microsoft/onnxruntime-extensions
Publicmirrored fromhttps://github.com/microsoft/onnxruntime-extensionsAvailable
test/test_vector_to_string.py
73lines · modecode
| 1 | import os |
| 2 | from pathlib import Path |
| 3 | import unittest |
| 4 | import numpy as np |
| 5 | from onnx import helper, onnx_pb as onnx_proto |
| 6 | import onnxruntime as _ort |
| 7 | from onnxruntime_customops import ( |
| 8 | onnx_op, |
| 9 | enable_custom_op, |
| 10 | PyCustomOpDef, |
| 11 | expand_onnx_inputs, |
| 12 | get_library_path as _get_library_path) |
| 13 | |
| 14 | |
| 15 | def _create_test_model(map, unk, intput_dims, output_dims): |
| 16 | nodes = [] |
| 17 | nodes[0:] = [helper.make_node('Identity', ['input_1'], ['identity1'])] |
| 18 | nodes[1:] = [helper.make_node( |
| 19 | 'VectorToString', ['identity1'], ['customout'], map=_serialize_map(map), unk=unk, |
| 20 | domain='ai.onnx.contrib')] |
| 21 | |
| 22 | input0 = helper.make_tensor_value_info( |
| 23 | 'input_1', onnx_proto.TensorProto.INT64, [None] * intput_dims) |
| 24 | output0 = helper.make_tensor_value_info( |
| 25 | 'customout', onnx_proto.TensorProto.STRING, [None] * output_dims) |
| 26 | |
| 27 | graph = helper.make_graph(nodes, 'test0', [input0], [output0]) |
| 28 | model = helper.make_model( |
| 29 | graph, opset_imports=[helper.make_operatorsetid("", 12)]) |
| 30 | return model |
| 31 | |
| 32 | |
| 33 | def _serialize_map(map): |
| 34 | result = "" |
| 35 | for k in map: |
| 36 | result = result + k + "\t" + " ".join([str(i) for i in map[k]]) + "\n" |
| 37 | return result |
| 38 | |
| 39 | |
| 40 | def _run_vector_to_string(input, output, map, unk): |
| 41 | model = _create_test_model(map, unk, input.ndim, input.ndim) |
| 42 | |
| 43 | so = _ort.SessionOptions() |
| 44 | so.register_custom_ops_library(_get_library_path()) |
| 45 | sess = _ort.InferenceSession(model.SerializeToString(), so) |
| 46 | result = sess.run(None, {'input_1': input}) |
| 47 | np.testing.assert_array_equal(result, [output]) |
| 48 | |
| 49 | |
| 50 | class TestVectorToString(unittest.TestCase): |
| 51 | @classmethod |
| 52 | def setUpClass(cls): |
| 53 | pass |
| 54 | |
| 55 | def test_vector_to_(self): |
| 56 | _run_vector_to_string(input=np.array([0, 2, 3, 4], dtype=np.int64), |
| 57 | output=np.array(["a", "b", "c", "unknown_word"]), |
| 58 | map={"a": [0], "b": [2], "c": [3]}, |
| 59 | unk="unknown_word") |
| 60 | |
| 61 | _run_vector_to_string(input=np.array([[0, ], [2, ], [3, ], [4, ]], dtype=np.int64), |
| 62 | output=np.array(["a", "b", "c", "unknown_word"]), |
| 63 | map={"a": [0], "b": [2], "c": [3]}, |
| 64 | unk="unknown_word") |
| 65 | |
| 66 | _run_vector_to_string(input=np.array([[0, 1], [2, 3], [3, 4], [4, 5]], dtype=np.int64), |
| 67 | output=np.array(["a", "b", "c", "unknown_word"]), |
| 68 | map={"a": [0, 1], "b": [2, 3], "c": [3, 4]}, |
| 69 | unk="unknown_word") |
| 70 | |
| 71 | |
| 72 | if __name__ == "__main__": |
| 73 | unittest.main() |
| 74 | |