Skip to content

Commit 8ccbb8c

Browse files
sayanshaw24wenbingl
sayanshaw24
authored and
wenbingl
committed
Initial RobertaTokenizer implementation (#365)
* Added initial RobertaTokenizer implementation * Added offset mapping to output * Updates for new custom op changes --------- Authored-by: Sayan Shaw <sayanshaw@microsoft.com>
1 parent 1051543 commit 8ccbb8c

7 files changed

Lines changed: 367 additions & 2 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -322,7 +322,7 @@ set(_HAS_TOKENIZER OFF)
322322
if(OCOS_ENABLE_GPT2_TOKENIZER)
323323
# GPT2
324324
set(_HAS_TOKENIZER ON)
325-
file(GLOB tok_TARGET_SRC "operators/tokenizer/gpt*.cc" "operators/tokenizer/unicode*.*" "operators/tokenizer/clip*.cc")
325+
file(GLOB tok_TARGET_SRC "operators/tokenizer/gpt*.cc" "operators/tokenizer/unicode*.*" "operators/tokenizer/clip*.cc" "operators/tokenizer/roberta*.cc")
326326
list(APPEND TARGET_SRC ${tok_TARGET_SRC})
327327
endif()
328328

‎operators/tokenizer/bpetokenizer.hpp‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -20,6 +20,7 @@
2020
#include "nlohmann/json.hpp"
2121
#include "clip_tokenizer.hpp"
2222
#include "gpt2_tokenizer.hpp"
23+
#include "roberta_tokenizer.hpp"
2324
#include "string_tensor.h"
2425
#include "unicode.h"
2526

‎operators/tokenizer/roberta_tokenizer.cc‎

Lines changed: 247 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,247 @@
1+
// Copyright (c) Microsoft Corporation. All rights reserved.
2+
// Licensed under the MIT License.
3+
// Partial code comes from other Microsoft employee.
4+
5+
#include <string>
6+
#include <vector>
7+
#include <fstream>
8+
#include <sstream>
9+
#include <iostream>
10+
#include <algorithm>
11+
#include <list>
12+
#include <memory>
13+
#include <regex>
14+
#include <sstream>
15+
#include <stdexcept>
16+
#include <unordered_map>
17+
#include <functional>
18+
#include <codecvt>
19+
#include <mutex>
20+
21+
#include "nlohmann/json.hpp"
22+
#include "bpetokenizer.hpp"
23+
#include "string_tensor.h"
24+
#include "unicode.h"
25+
26+
// Note: the following logic comes from CPython: unicodetype_db.h (_PyUnicode_IsWhitespace)
27+
bool IsWithinUnicodeSpace(char32_t ch) {
28+
switch (ch) {
29+
case 0x0009:
30+
case 0x000A:
31+
case 0x000B:
32+
case 0x000C:
33+
case 0x000D:
34+
case 0x001C:
35+
case 0x001D:
36+
case 0x001E:
37+
case 0x001F:
38+
case 0x0020:
39+
case 0x0085:
40+
case 0x00A0:
41+
case 0x1680:
42+
case 0x2000:
43+
case 0x2001:
44+
case 0x2002:
45+
case 0x2003:
46+
case 0x2004:
47+
case 0x2005:
48+
case 0x2006:
49+
case 0x2007:
50+
case 0x2008:
51+
case 0x2009:
52+
case 0x200A:
53+
case 0x2028:
54+
case 0x2029:
55+
case 0x202F:
56+
case 0x205F:
57+
case 0x3000:
58+
return true;
59+
}
60+
return false;
61+
}
62+
63+
bool IsEmptyuString(const ustring& str) {
64+
return std::all_of(str.begin(), str.end(), [](char32_t ch) { return IsWithinUnicodeSpace(ch); });
65+
}
66+
67+
KernelRobertaBpeTokenizer::KernelRobertaBpeTokenizer(const OrtApi& api, const OrtKernelInfo& info)
68+
: BaseKernel(api, info) {
69+
std::string vocab = ort_.KernelInfoGetAttribute<std::string>(&info, "vocab");
70+
if (vocab.empty()) {
71+
ORTX_CXX_API_THROW("vocabulary shouldn't be empty.", ORT_INVALID_ARGUMENT);
72+
}
73+
74+
std::string merges = ort_.KernelInfoGetAttribute<std::string>(&info, "merges");
75+
if (merges.empty()) {
76+
ORTX_CXX_API_THROW("merges shouldn't be empty.", ORT_INVALID_ARGUMENT);
77+
}
78+
79+
if (!TryToGetAttribute<int64_t>("padding_length", padding_length_)) {
80+
padding_length_ = -1;
81+
}
82+
83+
if (padding_length_ != -1 && padding_length_ <= 0) {
84+
ORTX_CXX_API_THROW("padding_length should be more than 0 or equal -1", ORT_INVALID_ARGUMENT);
85+
}
86+
87+
std::stringstream vocabu_stream(vocab);
88+
std::stringstream merges_stream(merges);
89+
bbpe_tokenizer_ = std::make_shared<VocabData>();
90+
bbpe_tokenizer_->Load(vocabu_stream, merges_stream, "<|endoftext|>", "<|endoftext|>");
91+
}
92+
93+
std::vector<int64_t> KernelRobertaBpeTokenizer::Tokenize(ustring& input, int64_t max_length, std::list<std::list<std::pair<int, int>>>& offset_map) {
94+
std::vector<int64_t> res;
95+
96+
if (IsEmptyuString(input)) {
97+
return res;
98+
}
99+
// Add BOS token to result
100+
res.push_back(bbpe_tokenizer_->GetEncoding("<s>"));
101+
102+
// Parse input
103+
auto special_token_split_res = bbpe_tokenizer_->SplitBySpecialTokens(input);
104+
TokenWithRegularExp regcmp;
105+
106+
for (auto& seg_id : special_token_split_res) {
107+
if (static_cast<int64_t>(res.size()) >= max_length) break;
108+
109+
if (seg_id.second != -1) {
110+
res.push_back(seg_id.second);
111+
continue;
112+
}
113+
114+
auto cur_input = std::move(seg_id.first);
115+
// Note: keep ptr to make sure the string_view is valid in the following process
116+
const char32_t* ptr = cur_input.c_str();
117+
regcmp.Set(ptr);
118+
119+
int offset = 0;
120+
std::list<std::pair<int, int>> offset_mapping;
121+
122+
// Add offset mapping for BOS token
123+
offset_mapping.push_back(std::make_pair(0, 0));
124+
125+
while (static_cast<int64_t>(res.size()) < max_length) {
126+
auto [b, tok] = regcmp.GetNextToken();
127+
if (!b) break;
128+
129+
std::string utf8_token = std::string(ustring(tok));
130+
131+
// Handle offset mapping and special cases
132+
if (utf8_token.at(0) == ' ') {
133+
offset_mapping.push_back(std::make_pair(offset + 1, offset + utf8_token.size()));
134+
} else {
135+
offset_mapping.push_back(std::make_pair(offset, offset + utf8_token.size()));
136+
}
137+
offset += utf8_token.size();
138+
139+
// Get byte encodings prior to performing BPE
140+
byte_list_.clear();
141+
for (char& cp : utf8_token) {
142+
byte_list_.push_back(bbpe_tokenizer_->ByteEncoder()[static_cast<unsigned char>(cp)]);
143+
}
144+
145+
// Perform BPE
146+
bbpe_tokenizer_->bpe(byte_list_);
147+
148+
// Add output to result
149+
for (auto p : byte_list_) {
150+
if (static_cast<int64_t>(res.size()) >= max_length) {
151+
break;
152+
}
153+
154+
res.push_back(p);
155+
}
156+
}
157+
// Add offset mapping for EOS token
158+
offset_mapping.push_back(std::make_pair(0, 0));
159+
160+
// Add offset mappings for input in this instance to list of offset mappings for all inputs
161+
offset_map.push_back(offset_mapping);
162+
}
163+
// Add EOS token to result
164+
res.push_back(bbpe_tokenizer_->GetEncoding("</s>"));
165+
return res;
166+
}
167+
168+
void KernelRobertaBpeTokenizer::Compute(OrtKernelContext* context) {
169+
// Setup inputs
170+
const OrtValue* input = ort_.KernelContext_GetInput(context, 0);
171+
std::vector<std::string> str_input;
172+
std::list<std::list<std::pair<int, int>>> offset_map;
173+
GetTensorMutableDataString(api_, ort_, context, input, str_input);
174+
OrtTensorDimensions input_dim(ort_, input);
175+
176+
std::vector<std::vector<int64_t>> tokenize_results;
177+
for (auto& str : str_input) {
178+
ustring ustr = ustring(str);
179+
tokenize_results.emplace_back(Tokenize(ustr, padding_length_ < 0 ? INT64_MAX : padding_length_, offset_map));
180+
}
181+
182+
size_t max_length = 0;
183+
if (padding_length_ == -1) {
184+
for (auto& res : tokenize_results) {
185+
max_length = std::max(max_length, res.size());
186+
}
187+
} else {
188+
max_length = static_cast<size_t>(padding_length_);
189+
}
190+
191+
OrtTensorDimensions output_dim = input_dim;
192+
output_dim.push_back(max_length);
193+
194+
OrtTensorDimensions offset_dim = output_dim;
195+
offset_dim.push_back(2); // tuple of offsets for each input id
196+
197+
OrtValue* tokenize_output = ort_.KernelContext_GetOutput(context, 0, output_dim.data(), output_dim.size());
198+
OrtValue* attention_mask = ort_.KernelContext_GetOutput(context, 1, output_dim.data(), output_dim.size());
199+
OrtValue* offset_mapping = ort_.KernelContext_GetOutput(context, 2, offset_dim.data(), offset_dim.size());
200+
auto* token = ort_.GetTensorMutableData<int64_t>(tokenize_output);
201+
auto* mask = ort_.GetTensorMutableData<int64_t>(attention_mask);
202+
auto* offset = ort_.GetTensorMutableData<int64_t>(offset_mapping);
203+
204+
int idx = 0;
205+
for (auto& res : tokenize_results) {
206+
for (int64_t id : res) {
207+
token[idx] = id;
208+
mask[idx] = 1;
209+
idx++;
210+
}
211+
212+
for (size_t i = res.size(); i < max_length; i++) {
213+
token[idx] = 0;
214+
mask[idx] = 0;
215+
idx++;
216+
}
217+
}
218+
219+
int idx2 = 0;
220+
for (auto& res : offset_map) {
221+
for (auto& mapping : res) {
222+
offset[idx2] = mapping.first;
223+
idx2++;
224+
offset[idx2] = mapping.second;
225+
idx2++;
226+
}
227+
}
228+
}
229+
230+
const char* CustomOpRobertaBpeTokenizer::GetName() const {
231+
return "RobertaTokenizer";
232+
}
233+
234+
size_t CustomOpRobertaBpeTokenizer::GetInputTypeCount() const {
235+
return 1;
236+
}
237+
238+
ONNXTensorElementDataType CustomOpRobertaBpeTokenizer::GetInputType(size_t /*index*/) const {
239+
return ONNX_TENSOR_ELEMENT_DATA_TYPE_STRING;
240+
}
241+
size_t CustomOpRobertaBpeTokenizer::GetOutputTypeCount() const {
242+
return 3;
243+
}
244+
245+
ONNXTensorElementDataType CustomOpRobertaBpeTokenizer::GetOutputType(size_t /*index*/) const {
246+
return ONNX_TENSOR_ELEMENT_DATA_TYPE_INT64;
247+
}

‎operators/tokenizer/roberta_tokenizer.hpp‎

Lines changed: 26 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,26 @@
1+
#include <list>
2+
#include "ocos.h"
3+
#include "ustring.h"
4+
#include "string_utils.h"
5+
6+
class VocabData;
7+
8+
struct KernelRobertaBpeTokenizer : BaseKernel {
9+
KernelRobertaBpeTokenizer(const OrtApi& api, const OrtKernelInfo& info);
10+
void Compute(OrtKernelContext* context);
11+
12+
private:
13+
std::vector<int64_t> Tokenize(ustring& input, int64_t max_length, std::list<std::list<std::pair<int, int>>>& offset_map);
14+
15+
int64_t padding_length_;
16+
std::list<int> byte_list_;
17+
std::shared_ptr<VocabData> bbpe_tokenizer_;
18+
};
19+
20+
struct CustomOpRobertaBpeTokenizer : OrtW::CustomOpBase<CustomOpRobertaBpeTokenizer, KernelRobertaBpeTokenizer> {
21+
const char* GetName() const;
22+
size_t GetInputTypeCount() const;
23+
ONNXTensorElementDataType GetInputType(size_t index) const;
24+
size_t GetOutputTypeCount() const;
25+
ONNXTensorElementDataType GetOutputType(size_t index) const;
26+
};

‎operators/tokenizer/tokenizers.cc‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
#ifdef ENABLE_GPT2_TOKENIZER
55
#include "gpt2_tokenizer.hpp"
66
#include "clip_tokenizer.hpp"
7+
#include "roberta_tokenizer.hpp"
78
#endif
89

910
#ifdef ENABLE_SPM_TOKENIZER
@@ -31,6 +32,7 @@ FxLoadCustomOpFactory LoadCustomOpClasses_Tokenizer = LoadCustomOpClasses<
3132
#ifdef ENABLE_GPT2_TOKENIZER
3233
, CustomOpBpeTokenizer
3334
, CustomOpClipBpeTokenizer
35+
, CustomOpRobertaBpeTokenizer
3436
#endif
3537

3638
#ifdef ENABLE_SPM_TOKENIZER

‎test/test_robertatok.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import unittest
2+
import numpy as np
3+
import numpy.lib.recfunctions as nlr
4+
import onnxruntime as _ort
5+
6+
from pathlib import Path
7+
from onnx import helper, onnx_pb as onnx_proto
8+
from transformers import RobertaTokenizerFast
9+
from onnxruntime_extensions import (
10+
make_onnx_model,
11+
get_library_path as _get_library_path)
12+
13+
def _get_file_content(path):
14+
with open(path, "rb") as file:
15+
return file.read()
16+
17+
18+
def _create_test_model(**kwargs):
19+
vocab_file = kwargs["vocab_file"]
20+
merges_file = kwargs["merges_file"]
21+
max_length = kwargs["max_length"]
22+
23+
node = [helper.make_node(
24+
'RobertaTokenizer', ['string_input'], ['input_ids', 'attention_mask', 'offset_mapping'], vocab=_get_file_content(vocab_file),
25+
merges=_get_file_content(merges_file), name='bpetok', padding_length=max_length,
26+
domain='ai.onnx.contrib')]
27+
input1 = helper.make_tensor_value_info(
28+
'string_input', onnx_proto.TensorProto.STRING, [None])
29+
output1 = helper.make_tensor_value_info(
30+
'input_ids', onnx_proto.TensorProto.INT64, ["batch_size", "num_input_ids"])
31+
output2 = helper.make_tensor_value_info(
32+
'attention_mask', onnx_proto.TensorProto.INT64, ["batch_size", "num_attention_masks"])
33+
output3 = helper.make_tensor_value_info(
34+
'offset_mapping', onnx_proto.TensorProto.INT64, ["batch_size", "num_offsets", 2])
35+
36+
graph = helper.make_graph(node, 'test0', [input1], [output1, output2, output3])
37+
model = make_onnx_model(graph)
38+
return model
39+
40+
41+
class TestRobertaTokenizer(unittest.TestCase):
42+
@classmethod
43+
def setUpClass(cls):
44+
cls.tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
45+
files = cls.tokenizer.save_vocabulary(".")
46+
cls.tokjson = files[0]
47+
cls.merges = files[1]
48+
49+
def _run_tokenizer(self, test_sentence, padding_length=-1):
50+
model = _create_test_model(vocab_file=self.tokjson, merges_file=self.merges, max_length=padding_length)
51+
so = _ort.SessionOptions()
52+
so.register_custom_ops_library(_get_library_path())
53+
sess = _ort.InferenceSession(model.SerializeToString(), so)
54+
input_text = np.array(test_sentence)
55+
input_ids, attention_mask, offset_mapping = sess.run(None, {'string_input': input_text})
56+
print("\nTest Sentence: " + str(test_sentence))
57+
print("\nInput IDs: " + str(input_ids))
58+
print("Attention Mask: " + str(attention_mask))
59+
# Reformat offset mapping from 3d array to 2d array of tuples before printing for readability
60+
reformatted_offset_mapping = nlr.unstructured_to_structured(np.array(offset_mapping)).astype('O')
61+
print("Offset Mapping: " + str(reformatted_offset_mapping))
62+
roberta_out = self.tokenizer(test_sentence, return_offsets_mapping=True)
63+
expect_input_ids = roberta_out['input_ids']
64+
expect_attention_mask = roberta_out['attention_mask']
65+
expect_offset_mapping = roberta_out['offset_mapping']
66+
print("\nExpected Input IDs: " + str(expect_input_ids))
67+
print("Expected Attention Mask: " + str(expect_attention_mask))
68+
print("Expected Offset Mapping: " + str(expect_offset_mapping) + "\n")
69+
np.testing.assert_array_equal(expect_input_ids, input_ids)
70+
np.testing.assert_array_equal(expect_attention_mask, attention_mask)
71+
np.testing.assert_array_equal(expect_offset_mapping, offset_mapping)
72+
73+
del sess
74+
del so
75+
76+
def test_tokenizer(self):
77+
self._run_tokenizer(["I can feel the magic, can you?"])
78+
self._run_tokenizer(["Hey Cortana"])
79+
self._run_tokenizer(["lower newer"])
80+
self._run_tokenizer(["a diagram", "a dog", "a cat"])
81+
self._run_tokenizer(["a photo of a cat", "a photo of a dog"])
82+
self._run_tokenizer(["one + two = three"])
83+
self._run_tokenizer(["9 8 7 6 5 4 3 2 1 0"])
84+
self._run_tokenizer(["9 8 7 - 6 5 4 - 3 2 1 0"])
85+
self._run_tokenizer(["One Microsoft Way, Redmond, WA"])
86+
87+
88+
if __name__ == "__main__":
89+
unittest.main()

0 commit comments

Comments
 (0)