From ac7d9c4e5c81c2da426d8e3e2bcf8d7be18f1dd1 Mon Sep 17 00:00:00 2001 From: Haokai Ding Date: Wed, 16 Sep 2026 14:11:30 +0400 Subject: [PATCH] Fix custom prefixes in BertWordPieceTokenizer --- .../implementations/bert_wordpiece.py | 6 ++-- .../implementations/test_bert_wordpiece.py | 32 +++++++++++++++++++ 2 files changed, 36 insertions(+), 2 deletions(-) diff --git a/bindings/python/py_src/tokenizers/implementations/bert_wordpiece.py b/bindings/python/py_src/tokenizers/implementations/bert_wordpiece.py index 1f34e3ca8a..f2bd9aedb6 100644 --- a/bindings/python/py_src/tokenizers/implementations/bert_wordpiece.py +++ b/bindings/python/py_src/tokenizers/implementations/bert_wordpiece.py @@ -27,9 +27,11 @@ def __init__( wordpieces_prefix: str = "##", ): if vocab is not None: - tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(unk_token))) + tokenizer = Tokenizer( + WordPiece(vocab, unk_token=str(unk_token), continuing_subword_prefix=wordpieces_prefix) + ) else: - tokenizer = Tokenizer(WordPiece(unk_token=str(unk_token))) + tokenizer = Tokenizer(WordPiece(unk_token=str(unk_token), continuing_subword_prefix=wordpieces_prefix)) # Let the tokenizer know about special tokens if they are part of the vocab if tokenizer.token_to_id(str(unk_token)) is not None: diff --git a/bindings/python/tests/implementations/test_bert_wordpiece.py b/bindings/python/tests/implementations/test_bert_wordpiece.py index a7fefc1d30..7aafb95b7a 100644 --- a/bindings/python/tests/implementations/test_bert_wordpiece.py +++ b/bindings/python/tests/implementations/test_bert_wordpiece.py @@ -5,6 +5,38 @@ class TestBertWordPieceTokenizer: + @pytest.mark.parametrize("wordpieces_prefix", ["##", "@@"]) + def test_wordpieces_prefix(self, wordpieces_prefix): + vocab = {"[UNK]": 0, "[CLS]": 1, "[SEP]": 2, "h": 3, f"{wordpieces_prefix}ello": 4} + tokenizer = BertWordPieceTokenizer(vocab, wordpieces_prefix=wordpieces_prefix) + + output = tokenizer.encode("hello", add_special_tokens=False) + assert output.tokens == ["h", f"{wordpieces_prefix}ello"] + assert output.ids == [3, 4] + assert output.offsets == [(0, 1), (1, 5)] + assert tokenizer.decode(output.ids) == "hello" + + @pytest.mark.parametrize("wordpieces_prefix", ["##", "@@"]) + def test_train_save_reload_wordpieces_prefix(self, wordpieces_prefix, tmp_path): + tokenizer = BertWordPieceTokenizer(wordpieces_prefix=wordpieces_prefix) + tokenizer.train_from_iterator( + ["hello", "help", "hello", "help"], + vocab_size=11, + min_frequency=2, + wordpieces_prefix=wordpieces_prefix, + show_progress=False, + ) + output = tokenizer.encode("hello", add_special_tokens=False) + assert len(output.tokens) > 1 + assert all(token.startswith(wordpieces_prefix) for token in output.tokens[1:]) + + vocab_path = tokenizer.save_model(str(tmp_path))[0] + reloaded = BertWordPieceTokenizer.from_file(vocab_path, wordpieces_prefix=wordpieces_prefix) + reloaded_output = reloaded.encode("hello", add_special_tokens=False) + assert reloaded_output.tokens == output.tokens + assert reloaded_output.ids == output.ids + assert reloaded.decode(reloaded_output.ids) == "hello" + @pytest.mark.network def test_basic_encode(self, bert_files): tokenizer = BertWordPieceTokenizer.from_file(bert_files["vocab"])