microsoft/onnxruntime-extensions
Publicmirrored from https://github.com/microsoft/onnxruntime-extensionsAvailable
test/test_bert_tokenizer.py
56lines · modecode
| 1 | from pathlib import Path |
| 2 | import unittest |
| 3 | import numpy as np |
| 4 | import transformers |
| 5 | from onnxruntime_extensions import PyOrtFunction, BertTokenizer |
| 6 | |
| 7 | |
| 8 | def _get_test_data_file(*sub_dirs): |
| 9 | test_dir = Path(__file__).parent |
| 10 | return str(test_dir.joinpath(*sub_dirs)) |
| 11 | |
| 12 | |
| 13 | bert_cased_tokenizer = transformers.BertTokenizer(_get_test_data_file('data', 'bert_basic_cased_vocab.txt'), False, |
| 14 | strip_accents=True) |
| 15 | |
| 16 | |
| 17 | def _run_basic_case(input, vocab_path): |
| 18 | t2stc = PyOrtFunction.from_customop(BertTokenizer, vocab_file=vocab_path, do_lower_case=0, strip_accents=1) |
| 19 | result = t2stc([input]) |
| 20 | expect_result = bert_cased_tokenizer.encode_plus(input) |
| 21 | np.testing.assert_array_equal(result[0], expect_result['input_ids']) |
| 22 | np.testing.assert_array_equal(result[1], expect_result['token_type_ids']) |
| 23 | np.testing.assert_array_equal(result[2], expect_result['attention_mask']) |
| 24 | |
| 25 | |
| 26 | def _run_combined_case(input, vocab_path): |
| 27 | t2stc = PyOrtFunction.from_customop(BertTokenizer, vocab_file=vocab_path, do_lower_case=0, strip_accents=1) |
| 28 | result = t2stc(input) |
| 29 | expect_result = bert_cased_tokenizer.encode_plus(input[0], input[1]) |
| 30 | np.testing.assert_array_equal(result[0], expect_result['input_ids']) |
| 31 | np.testing.assert_array_equal(result[1], expect_result['token_type_ids']) |
| 32 | np.testing.assert_array_equal(result[2], expect_result['attention_mask']) |
| 33 | |
| 34 | |
| 35 | class TestBertTokenizer(unittest.TestCase): |
| 36 | |
| 37 | def test_text_to_case1(self): |
| 38 | _run_basic_case(input="Input 'text' must not be empty.", |
| 39 | vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 40 | _run_basic_case( |
| 41 | input="ÀÁÂÃÄÅÇÈÉÊËÌÍÎÑÒÓÔÕÖÚÜ\t䗓𨖷虴𨀐辘𧄋脟𩑢𡗶镇伢𧎼䪱轚榶𢑌㺽𤨡!#$%&(Tom@microsoft.com)*+,-./:;<=>?@[\\]^_`{|}~", |
| 42 | vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 43 | _run_basic_case(input="网易云音乐", vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 44 | _run_basic_case(input="本想好好的伤感 想放任 但是没泪痕", vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 45 | _run_basic_case(input="网 易 云 音 乐", |
| 46 | vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 47 | _run_basic_case(input="cat is playing toys", |
| 48 | vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 49 | _run_basic_case(input="cat isnot playing toyssss", |
| 50 | vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 51 | _run_combined_case(["网 易 云 音 乐", "cat isnot playing toyssss"], |
| 52 | vocab_path=_get_test_data_file('data', 'bert_basic_cased_vocab.txt')) |
| 53 | |
| 54 | |
| 55 | if __name__ == "__main__": |
| 56 | unittest.main() |
| 57 | |