-
Notifications
You must be signed in to change notification settings - Fork 3k
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
### Description <!-- Describe your changes. --> ### Motivation and Context Closes #17596
- Loading branch information
1 parent
48d163a
commit 7b6a615
Showing
7 changed files
with
248 additions
and
7 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,98 @@ | ||
// Copyright (c) Microsoft Corporation. All rights reserved. | ||
// Licensed under the MIT License. | ||
|
||
#include "string_split.h" | ||
#include <algorithm> | ||
#include <limits> | ||
#include <string> | ||
#include "core/common/common.h" | ||
namespace onnxruntime { | ||
|
||
ONNX_CPU_OPERATOR_KERNEL(StringSplit, 20, | ||
KernelDefBuilder() | ||
.TypeConstraint("T1", DataTypeImpl::GetTensorType<std::string>()) | ||
.TypeConstraint("T2", DataTypeImpl::GetTensorType<std::string>()) | ||
.TypeConstraint("T3", DataTypeImpl::GetTensorType<int64_t>()), | ||
StringSplit); | ||
|
||
/// Calculate substrings in ``str`` delimited by ``delimiter``. A maximum of ``max_splits`` splits are permitted. | ||
/// Returns a vector of string slices into ``str`` representing the substrings as string views. The user must ensure | ||
/// the returned views' lifetime does not exceed ``str``'s. | ||
void ComputeSubstrings(std::string_view str, std::string_view delimiter, int64_t max_splits, InlinedVector<std::string_view>& out) { | ||
if (str.empty()) { | ||
return; | ||
} | ||
if (delimiter.empty()) { | ||
// Count consecutive whitespace as one delimiter. Preceding and trailing whitespace is meant to be ignored. | ||
size_t pos = str.find_first_not_of(" "); | ||
int64_t token_count = 0; | ||
while (pos != std::string::npos) { | ||
if (token_count++ == max_splits) { | ||
// Trim down last substring as required in specification | ||
size_t next_pos = str.length() - 1; | ||
while (str[next_pos] == ' ') { | ||
next_pos--; | ||
} | ||
out.push_back(str.substr(pos, next_pos - pos + 1)); | ||
break; | ||
} else { | ||
auto next_pos = str.find_first_of(" ", pos); | ||
out.push_back(str.substr(pos, next_pos - pos)); | ||
pos = str.find_first_not_of(" ", next_pos); | ||
} | ||
} | ||
} else { | ||
size_t pos = 0; | ||
int64_t token_count = 0; | ||
while (pos != std::string::npos) { | ||
auto next_pos = str.find(delimiter, pos); | ||
if (token_count++ == max_splits || next_pos == std::string::npos) { | ||
out.push_back(str.substr(pos)); | ||
break; | ||
} | ||
out.push_back(str.substr(pos, next_pos - pos)); | ||
pos = next_pos + delimiter.size(); | ||
} | ||
} | ||
} | ||
|
||
StringSplit::StringSplit(const OpKernelInfo& info) : OpKernel(info) { | ||
info.GetAttrOrDefault("maxsplit", &maxsplit_, std::numeric_limits<int64_t>::max() - 1); | ||
info.GetAttrOrDefault("delimiter", &delimiter_, std::string()); | ||
} | ||
|
||
Status StringSplit::Compute(OpKernelContext* context) const { | ||
const Tensor* input = context->Input<Tensor>(0); | ||
auto input_data = input->template DataAsSpan<std::string>(); | ||
|
||
// Set up number of tokens output | ||
auto num_tokens_data = context->Output(1, input->Shape())->template MutableDataAsSpan<int64_t>(); | ||
auto num_tokens_iter = num_tokens_data.begin(); | ||
|
||
InlinedVector<InlinedVector<std::string_view>> input_slices; | ||
input_slices.reserve(input_data.size()); | ||
size_t last_dim = 0; | ||
|
||
for (const auto& s : input_data) { | ||
auto& substrs = input_slices.emplace_back(); | ||
ComputeSubstrings(s, delimiter_, maxsplit_, substrs); | ||
auto substr_count = substrs.size(); | ||
last_dim = std::max(last_dim, substr_count); | ||
*num_tokens_iter = static_cast<int64_t>(substr_count); | ||
++num_tokens_iter; | ||
} | ||
|
||
// Set up splits output | ||
auto splits_shape = input->Shape().AsShapeVector(); | ||
splits_shape.push_back(last_dim); | ||
|
||
auto splits_data = context->Output(0, splits_shape)->template MutableDataAsSpan<std::string>(); | ||
auto slices_iter = input_slices.begin(); | ||
for (auto output_splits_iter = splits_data.begin(); output_splits_iter != splits_data.end(); output_splits_iter += last_dim, ++slices_iter) { | ||
std::copy(slices_iter->begin(), slices_iter->end(), output_splits_iter); | ||
} | ||
|
||
return Status::OK(); | ||
} | ||
|
||
} // namespace onnxruntime |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,20 @@ | ||
// Copyright (c) Microsoft Corporation. All rights reserved. | ||
// Licensed under the MIT License. | ||
|
||
#pragma once | ||
|
||
#include "core/framework/op_kernel.h" | ||
|
||
namespace onnxruntime { | ||
|
||
class StringSplit final : public OpKernel { | ||
public: | ||
explicit StringSplit(const OpKernelInfo& info); | ||
Status Compute(OpKernelContext* context) const override; | ||
|
||
private: | ||
std::string delimiter_; | ||
int64_t maxsplit_; | ||
}; | ||
|
||
} // namespace onnxruntime |
126 changes: 126 additions & 0 deletions
126
onnxruntime/test/providers/cpu/text/string_split_test.cc
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,126 @@ | ||
#include "gtest/gtest.h" | ||
#include "test/providers/provider_test_utils.h" | ||
|
||
namespace onnxruntime { | ||
namespace test { | ||
|
||
TEST(StringSplit, BasicSplitTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {3}, {"hello world", "hello", "world"}); | ||
test.AddAttribute<std::string>("delimiter", " "); | ||
test.AddOutput<std::string>("Y", {3, 2}, {"hello", "world", "hello", "", "world", ""}); | ||
test.AddOutput<int64_t>("Z", {3}, {2, 1, 1}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, MaxSplitTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {2, 2}, {"eggs;milk;chesse", "pepper;salt", "chicken;fish;pork", "spinach"}); | ||
test.AddAttribute<std::string>("delimiter", ";"); | ||
test.AddAttribute<int64_t>("maxsplit", 1); | ||
test.AddOutput<std::string>("Y", {2, 2, 2}, | ||
{"eggs", "milk;chesse", "pepper", "salt", "chicken", "fish;pork", "spinach", ""}); | ||
test.AddOutput<int64_t>("Z", {2, 2}, {2, 2, 2, 1}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, EmptyStringDelimiterTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 4}, {"hello world", "hello world", " hello world", "hello world "}); | ||
test.AddAttribute<std::string>("delimiter", ""); | ||
test.AddOutput<std::string>("Y", {1, 4, 2}, {"hello", "world", "hello", "world", "hello", "world", "hello", "world"}); | ||
test.AddOutput<int64_t>("Z", {1, 4}, {2, 2, 2, 2}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, SubsequentWhitespaceDefaultTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 4}, {"hello world", "hello world", " hello world", "hello world "}); | ||
test.AddOutput<std::string>("Y", {1, 4, 2}, {"hello", "world", "hello", "world", "hello", "world", "hello", "world"}); | ||
test.AddOutput<int64_t>("Z", {1, 4}, {2, 2, 2, 2}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, SubsequentWhitespaceWithLimitTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 4}, | ||
{"lorem ipsum doler", " Open Neural Network Exchange (ONNX)", "onnx", "ONNX runtime "}); | ||
test.AddAttribute<int64_t>("maxsplit", 1); | ||
test.AddOutput<std::string>( | ||
"Y", {1, 4, 2}, | ||
{"lorem", "ipsum doler", "Open", "Neural Network Exchange (ONNX)", "onnx", "", "ONNX", "runtime"}); | ||
test.AddOutput<int64_t>("Z", {1, 4}, {2, 2, 1, 2}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, SingleTokenTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddAttribute<std::string>("delimiter", "*"); | ||
test.AddInput<std::string>("X", {1, 1, 1}, {"lorem"}); | ||
test.AddOutput<std::string>("Y", {1, 1, 1, 1}, {"lorem"}); | ||
test.AddOutput<int64_t>("Z", {1, 1, 1}, {1}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, SingleTokenWhitespaceTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 1, 1}, {"lorem"}); | ||
test.AddOutput<std::string>("Y", {1, 1, 1, 1}, {"lorem"}); | ||
test.AddOutput<int64_t>("Z", {1, 1, 1}, {1}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, EdgeWhitespaceTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 1, 1}, {" lorem "}); | ||
test.AddOutput<std::string>("Y", {1, 1, 1, 1}, {"lorem"}); | ||
test.AddOutput<int64_t>("Z", {1, 1, 1}, {1}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, EmptyInputTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 3, 1}, {"", "+", "*"}); | ||
test.AddAttribute<std::string>("delimiter", "*"); | ||
test.AddOutput<std::string>("Y", {1, 3, 1, 2}, {"", "", "+", "", "", ""}); | ||
test.AddOutput<int64_t>("Z", {1, 3, 1}, {0, 1, 2}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, OnlyEmptyInputTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddAttribute<std::string>("delimiter", "*"); | ||
test.AddInput<std::string>("X", {1, 2, 1}, {"", ""}); | ||
test.AddOutput<std::string>("Y", {1, 2, 1, 0}, {}); | ||
test.AddOutput<int64_t>("Z", {1, 2, 1}, {0, 0}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, OnlyEmptyNoDelimiterInputTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", {1, 2, 1}, {"", ""}); | ||
test.AddOutput<std::string>("Y", {1, 2, 1, 0}, {}); | ||
test.AddOutput<int64_t>("Z", {1, 2, 1}, {0, 0}); | ||
test.Run(); | ||
} | ||
|
||
TEST(StringSplit, NoInputTest) { | ||
OpTester test("StringSplit", 20); | ||
test.AddInput<std::string>("X", { | ||
0, | ||
}, | ||
{}); | ||
test.AddOutput<std::string>("Y", { | ||
0, | ||
0, | ||
}, | ||
{}); | ||
test.AddOutput<int64_t>("Z", { | ||
0, | ||
}, | ||
{}); | ||
test.Run(); | ||
} | ||
|
||
} // namespace test | ||
} // namespace onnxruntime |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters