microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
operators/text/masked_fill.cc
76lines · modecode
| 1 | // Copyright (c) Microsoft Corporation. All rights reserved. |
| 2 | // Licensed under the MIT License. |
| 3 | |
| 4 | #include "masked_fill.hpp" |
| 5 | #include "string_tensor.h" |
| 6 | #include <vector> |
| 7 | #include <locale> |
| 8 | #include <codecvt> |
| 9 | #include <algorithm> |
| 10 | |
| 11 | KernelMaskedFill::KernelMaskedFill(const OrtApi& api, const OrtKernelInfo& info) : BaseKernel(api, info) { |
| 12 | } |
| 13 | |
| 14 | void KernelMaskedFill::Compute(OrtKernelContext* context) { |
| 15 | // Setup inputs |
| 16 | const OrtValue* input_value = ort_.KernelContext_GetInput(context, 0); |
| 17 | const OrtValue* input_mask = ort_.KernelContext_GetInput(context, 1); |
| 18 | |
| 19 | OrtTensorDimensions value_dimensions(ort_, input_value); |
| 20 | OrtTensorDimensions mask_dimensions(ort_, input_mask); |
| 21 | |
| 22 | if (!(value_dimensions.IsScalar() || value_dimensions.IsVector())) { |
| 23 | ORTX_CXX_API_THROW("[MaskedFill]: the dimension of input value should be vector or scalar.", ORT_INVALID_ARGUMENT); |
| 24 | } |
| 25 | |
| 26 | if (value_dimensions != mask_dimensions) { |
| 27 | ORTX_CXX_API_THROW("[MaskedFill]: the dimension of input value and mask should be same.", ORT_INVALID_ARGUMENT); |
| 28 | } |
| 29 | |
| 30 | std::vector<std::string> value; |
| 31 | const bool* mask = nullptr; |
| 32 | |
| 33 | GetTensorMutableDataString(api_, ort_, context, input_value, value); |
| 34 | mask = ort_.GetTensorData<bool>(input_mask); |
| 35 | |
| 36 | std::vector<std::string> result; |
| 37 | std::vector<int64_t> result_dimension; |
| 38 | |
| 39 | for (size_t i = 0; i < value.size(); i++) { |
| 40 | if (!mask[i]) { |
| 41 | continue; |
| 42 | } |
| 43 | |
| 44 | result.push_back(value[i]); |
| 45 | } |
| 46 | result_dimension.push_back(result.size()); |
| 47 | |
| 48 | OrtValue* output = ort_.KernelContext_GetOutput(context, 0, result_dimension.data(), result_dimension.size()); |
| 49 | |
| 50 | FillTensorDataString(api_, ort_, context, result, output); |
| 51 | } |
| 52 | |
| 53 | const char* CustomOpMaskedFill::GetName() const { return "MaskedFill"; }; |
| 54 | |
| 55 | size_t CustomOpMaskedFill::GetInputTypeCount() const { |
| 56 | return 2; |
| 57 | }; |
| 58 | |
| 59 | ONNXTensorElementDataType CustomOpMaskedFill::GetInputType(size_t index) const { |
| 60 | switch (index) { |
| 61 | case 0: |
| 62 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; |
| 63 | case 1: |
| 64 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_BOOL; |
| 65 | default: |
| 66 | ORTX_CXX_API_THROW(MakeString("Unexpected input index ", index), ORT_INVALID_ARGUMENT); |
| 67 | } |
| 68 | }; |
| 69 | |
| 70 | size_t CustomOpMaskedFill::GetOutputTypeCount() const { |
| 71 | return 1; |
| 72 | }; |
| 73 | |
| 74 | ONNXTensorElementDataType CustomOpMaskedFill::GetOutputType(size_t /*index*/) const { |
| 75 | return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING; |
| 76 | }; |
| 77 | |