diff --git a/src/megatron/bridge/inference/_tokenizer.py b/src/megatron/bridge/inference/_tokenizer.py index b1e3873d73..5ec33be81f 100644 --- a/src/megatron/bridge/inference/_tokenizer.py +++ b/src/megatron/bridge/inference/_tokenizer.py @@ -46,6 +46,11 @@ def __init__( ) -> None: self._tokenizer = tokenizer self._force_skip_special_tokens = force_skip_special_tokens + # MCore's incremental detokenizer treats ``_tokenizer`` as its Hugging Face wrapper. + if not hasattr(tokenizer, "tokenizer"): + tokenizer.tokenizer = tokenizer + if not hasattr(tokenizer, "include_special_tokens"): + tokenizer.include_special_tokens = force_skip_special_tokens is False if set_pad_token and tokenizer.pad_token is None and tokenizer.eos_token is not None: tokenizer.pad_token = tokenizer.eos_token self.eod = tokenizer.eos_token_id diff --git a/tests/unit_tests/inference/test_tokenizer.py b/tests/unit_tests/inference/test_tokenizer.py index d0f626e745..68af3f3a72 100644 --- a/tests/unit_tests/inference/test_tokenizer.py +++ b/tests/unit_tests/inference/test_tokenizer.py @@ -16,6 +16,12 @@ from __future__ import annotations +from megatron.core.inference.text_generation_server.dynamic_text_gen_server.incremental_detokenizer import ( + HuggingFaceFastIncrementalDetokenizer, +) +from tokenizers import Tokenizer, models +from transformers import PreTrainedTokenizerFast + from megatron.bridge.inference._tokenizer import HFTokenizerAdapter @@ -119,3 +125,14 @@ def test_openai_server_preserves_hf_chat_template_contract(): }, ) ] + + +def test_openai_streaming_accepts_hf_fast_tokenizer_adapter(): + backend = Tokenizer(models.WordLevel(vocab={"[UNK]": 0, "hello": 1}, unk_token="[UNK]")) + tokenizer = PreTrainedTokenizerFast(tokenizer_object=backend, unk_token="[UNK]", eos_token="[UNK]") + adapter = HFTokenizerAdapter(tokenizer) + + detokenizer = HuggingFaceFastIncrementalDetokenizer(adapter, []) + + assert detokenizer.text == "" + assert detokenizer.update([tokenizer.eos_token_id]) == ""