From 3f3d8956056a03d126112a54d4461703806a26fa Mon Sep 17 00:00:00 2001 From: Daniel Song Date: Mon, 10 Aug 2026 17:24:51 -0700 Subject: [PATCH] Improve ONNX Runtime input validation --- operators/tokenizer/bert_tokenizer.cc | 24 +++++++++++++++++++----- test/test_bert_tokenizer.py | 25 +++++++++++++++++++++++++ 2 files changed, 44 insertions(+), 5 deletions(-) diff --git a/operators/tokenizer/bert_tokenizer.cc b/operators/tokenizer/bert_tokenizer.cc index aaa07f7e2..ed0b20d97 100644 --- a/operators/tokenizer/bert_tokenizer.cc +++ b/operators/tokenizer/bert_tokenizer.cc @@ -374,11 +374,18 @@ void KernelBertTokenizer::Compute(const ortc::Tensor& input, auto* p_out2 = output2.Allocate(output_dim); std::copy(attention_mask.begin(), attention_mask.end(), p_out2); - std::vector offset_dim{static_cast(input_ids.size()), 2}; // tuple of offsets for each input id - if (offset_mapping.has_value()) { + // Size the output from the same source that drives the write loop below. The + // number of offset pairs produced by tokenization does not necessarily equal + // input_ids.size() (e.g. truncation shrinks input_ids but not offset_map), so + // sizing from input_ids would under-allocate and overflow the buffer. + size_t offset_rows = 0; + for (auto& res : offset_map) { + offset_rows += res.size(); + } + std::vector offset_dim{static_cast(offset_rows), 2}; // tuple of offsets for each input id auto* offset = (*offset_mapping)->Allocate(offset_dim); - int idx2 = 0; + size_t idx2 = 0; for (auto& res : offset_map) { for (auto& mapping : res) { offset[idx2] = mapping.first; @@ -436,10 +443,17 @@ void KernelHfBertTokenizer::Compute(const ortc::Tensor& input, std::copy(token_type_ids.begin(), token_type_ids.end(), p_out2); } - std::vector offset_dim{static_cast(input_ids.size()), 2}; // tuple of offsets for each input id if (compute_offset_mapping) { + // Size the output from the same source that drives the write loop below, + // not from input_ids.size(); the two counts diverge and sizing from + // input_ids would under-allocate and overflow the buffer. + size_t offset_rows = 0; + for (auto& res : offset_map) { + offset_rows += res.size(); + } + std::vector offset_dim{static_cast(offset_rows), 2}; // tuple of offsets for each input id auto* offset = (*offset_mapping)->Allocate(offset_dim); - int idx2 = 0; + size_t idx2 = 0; for (auto& res : offset_map) { for (auto& mapping : res) { offset[idx2] = mapping.first; diff --git a/test/test_bert_tokenizer.py b/test/test_bert_tokenizer.py index e508679f4..68a96724b 100644 --- a/test/test_bert_tokenizer.py +++ b/test/test_bert_tokenizer.py @@ -34,6 +34,27 @@ def _run_combined_case(input, vocab_path): np.testing.assert_array_equal(result[1], expect_result["token_type_ids"]) np.testing.assert_array_equal(result[2], expect_result["attention_mask"]) +def _run_truncated_offset_bounds_check(input, vocab_path): + # Truncation shortens input_ids after offset_map has been generated. The + # output tensor must be allocated from offset_map so its write loop cannot + # run past a buffer sized from the truncated IDs. + t2stc = PyOrtFunction.from_customop( + BertTokenizer, + vocab_file=vocab_path, + do_lower_case=0, + strip_accents=1, + max_length=5, + ) + result = t2stc([input]) + offset_mapping = np.asarray(result[3]) + input_ids = np.asarray(result[0]) + assert offset_mapping.shape[1] == 2, offset_mapping.shape + assert offset_mapping.shape[0] > input_ids.shape[0], ( + offset_mapping.shape, + input_ids.shape, + ) + + def _run_basic_with_offset_check(input, vocab_path): t2stc = PyOrtFunction.from_customop( BertTokenizer, vocab_file=vocab_path, do_lower_case=0, strip_accents=1 @@ -80,6 +101,10 @@ def test_text_to_case1(self): ["网 易 云 音 乐", "cat isnot playing toyssss"], vocab_path=util.get_test_data_file("data", "bert_basic_cased_vocab.txt"), ) + _run_truncated_offset_bounds_check( + "cat isnot playing toyssss", + vocab_path=util.get_test_data_file("data", "bert_basic_cased_vocab.txt"), + ) print("\n****** Input ids, token type ids, and attention mask tests complete. ******\n\n\n") print("*** Starting offset mapping tests. ***\n")