microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
ocos/kernels/string_join.cc
100lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | #include "string_join.hpp" |
| 5 | |
| 6 | KernelStringJoin::KernelStringJoin(OrtApi api) : BaseKernel(api) { |
| 7 | } |
| 8 | |
| 9 | void KernelStringJoin::Compute(OrtKernelContext* context) { |
| 10 | // Setup inputs |
| 11 | const OrtValue* input_X = ort_.KernelContext_GetInput(context, 0); |
| 12 | const std::string* X = ort_.GetTensorData<std::string>(input_X); |
| 13 | const OrtValue* input_sep = ort_.KernelContext_GetInput(context, 1); |
| 14 | const std::string* sep = ort_.GetTensorData<std::string>(input_sep); |
| 15 | const OrtValue* input_axis = ort_.KernelContext_GetInput(context, 2); |
| 16 | const int64_t* axis = ort_.GetTensorData<int64_t>(input_axis); |
| 17 | |
| 18 | // Setup output |
| 19 | OrtTensorDimensions dimensions_sep(ort_, input_sep); |
| 20 | if (dimensions_sep.size() != 1 || dimensions_sep[0] != 1) |
| 21 | throw std::runtime_error("Input 2 is the separator, it has 1 element."); |
| 22 | OrtTensorDimensions dimensions_axis(ort_, input_axis); |
| 23 | if (dimensions_axis.size() != 1 || dimensions_axis[0] != 1) |
| 24 | throw std::runtime_error("Input 3 is the axis, it has 1 element."); |
| 25 | OrtTensorDimensions dimensions(ort_, input_X); |
| 26 | if (*axis < 0 || *axis >= dimensions.size()) |
| 27 | throw std::runtime_error(MakeString("axis must be positive and smaller than the number of dimension but it is ", *axis)); |
| 28 | |
| 29 | std::vector<int64_t> dimensions_out(dimensions.size() > 1 ? dimensions.size() - 1 : 1); |
| 30 | if (dimensions.size() > 1) { |
| 31 | for (size_t i = 0, pos = 0; i < dimensions.size(); ++i) { |
| 32 | if (static_cast<int64_t>(i) == *axis) |
| 33 | continue; |
| 34 | dimensions_out[pos++] = dimensions[i]; |
| 35 | } |
| 36 | } else { |
| 37 | dimensions_out[0] = 1; |
| 38 | } |
| 39 | |
| 40 | OrtValue* output = ort_.KernelContext_GetOutput(context, 0, dimensions_out.data(), dimensions_out.size()); |
| 41 | std::string* out = ort_.GetTensorMutableData<std::string>(output); |
| 42 | |
| 43 | OrtTensorTypeAndShapeInfo* output_info = ort_.GetTensorTypeAndShape(output); |
| 44 | int64_t size = ort_.GetTensorShapeElementCount(output_info); |
| 45 | ort_.ReleaseTensorTypeAndShapeInfo(output_info); |
| 46 | |
| 47 | // Do computation |
| 48 | int64_t h = 1; |
| 49 | for (size_t i = *axis + 1; i < dimensions.size(); ++i) { |
| 50 | h *= dimensions[i]; |
| 51 | } |
| 52 | int64_t left_part = size / h; |
| 53 | int64_t right_part = size / left_part; |
| 54 | int64_t n_red = dimensions[*axis] - 1; |
| 55 | int64_t inc = right_part * (n_red + 1); |
| 56 | int64_t pos = 0; |
| 57 | for (int64_t li = 0; li < left_part; ++li) { |
| 58 | for (int64_t ri = 0; ri < right_part; ++ri, ++pos) { |
| 59 | std::ostringstream st; |
| 60 | int64_t index = ri + li * inc; |
| 61 | for (int64_t j = 0; j < n_red; ++j, index += h) { |
| 62 | st << X[index] << *sep; |
| 63 | } |
| 64 | st << X[index]; |
| 65 | out[pos] = st.str(); |
| 66 | } |
| 67 | } |
| 68 | } |
| 69 | |
| 70 | void* CustomOpStringJoin::CreateKernel(OrtApi api, const OrtKernelInfo* /* info */) { |
| 71 | return new KernelStringJoin(api); |
| 72 | }; |
| 73 | |
| 74 | const char* CustomOpStringJoin::GetName() const { |
| 75 | return "StringJoin"; |
| 76 | }; |
| 77 | |
| 78 | size_t CustomOpStringJoin::GetInputTypeCount() const { |
| 79 | return 3; |
| 80 | }; |
| 81 | |
| 82 | ONNXTensorElementDataType CustomOpStringJoin::GetInputType(size_t index) const { |
| 83 | switch (index) { |
| 84 | case 0: |
| 85 | case 1: |
| 86 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; |
| 87 | case 2: |
| 88 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64; |
| 89 | default: |
| 90 | throw std::runtime_error(MakeString("Unexpected input index ", index)); |
| 91 | } |
| 92 | }; |
| 93 | |
| 94 | size_t CustomOpStringJoin::GetOutputTypeCount() const { |
| 95 | return 1; |
| 96 | }; |
| 97 | |
| 98 | ONNXTensorElementDataType CustomOpStringJoin::GetOutputType(size_t /*index*/) const { |
| 99 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; |
| 100 | }; |
| 101 | |