| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import pickle |
| 4 | from copy import deepcopy |
| 5 | |
| 6 | import pytest |
| 7 | from transformers import AutoTokenizer |
| 8 | |
| 9 | from vllm.tokenizers import TokenizerLike |
| 10 | from vllm.tokenizers.hf import ( |
| 11 | ThreadSafeHFTokenizerMixin, |
| 12 | get_cached_tokenizer, |
| 13 | maybe_make_thread_pool, |
| 14 | ) |
| 15 | |
| 16 | |
| 17 | @pytest.mark.parametrize("model_id", ["openai-community/gpt2", "zai-org/chatglm3-6b"]) |
| 18 | def test_cached_tokenizer(model_id: str): |
| 19 | reference_tokenizer = AutoTokenizer.from_pretrained( |
| 20 | model_id, trust_remote_code=True |
| 21 | ) |
| 22 | reference_tokenizer.add_special_tokens({"cls_token": "<CLS>"}) |
| 23 | reference_tokenizer.add_special_tokens({"additional_special_tokens": ["<SEP>"]}) |
| 24 | |
| 25 | cached_tokenizer = get_cached_tokenizer(deepcopy(reference_tokenizer)) |
| 26 | _check_consistency(cached_tokenizer, reference_tokenizer) |
| 27 | |
| 28 | pickled_tokenizer = pickle.dumps(cached_tokenizer) |
| 29 | unpickled_tokenizer = pickle.loads(pickled_tokenizer) |
| 30 | _check_consistency(unpickled_tokenizer, reference_tokenizer) |
| 31 | |
| 32 | |
| 33 | def _check_consistency(target: TokenizerLike, expected: TokenizerLike): |
| 34 | assert isinstance(target, type(expected)) |
| 35 | |
| 36 | # Cached attributes |
| 37 | assert target.all_special_ids == expected.all_special_ids |
| 38 | assert target.all_special_tokens == expected.all_special_tokens |
| 39 | assert target.get_vocab() == expected.get_vocab() |
| 40 | assert len(target) == len(expected) |
| 41 | |
| 42 | # Other attributes |
| 43 | assert getattr(target, "padding_side", None) == getattr( |
| 44 | expected, "padding_side", None |
| 45 | ) |
| 46 | |
| 47 | assert target.encode("prompt") == expected.encode("prompt") |
| 48 | |
| 49 | |
| 50 | @pytest.mark.parametrize("model_id", ["openai-community/gpt2"]) |
| 51 | def test_thread_pool_tokenizer_pickle(model_id: str): |
| 52 | """Regression test for issue #45433: the thread-pool tokenizer wrapper |
| 53 | reconstructs through maybe_make_thread_pool on unpickling, which used to |
| 54 | fall off the end and return None.""" |
| 55 | reference_tokenizer = AutoTokenizer.from_pretrained(model_id) |
| 56 | |
| 57 | pooled_tokenizer = maybe_make_thread_pool(deepcopy(reference_tokenizer)) |
| 58 | assert pooled_tokenizer is not None |
| 59 | assert isinstance(pooled_tokenizer, ThreadSafeHFTokenizerMixin) |
| 60 | |
| 61 | unpickled_tokenizer = pickle.loads(pickle.dumps(pooled_tokenizer)) |
| 62 | assert unpickled_tokenizer is not None |
| 63 | assert isinstance(unpickled_tokenizer, ThreadSafeHFTokenizerMixin) |
| 64 | assert unpickled_tokenizer.encode("prompt") == reference_tokenizer.encode("prompt") # type: ignore[attr-defined] |
| 65 | |
| 66 | # Idempotence: wrapping an already-pooled tokenizer returns it unchanged. |
| 67 | assert maybe_make_thread_pool(pooled_tokenizer) is pooled_tokenizer # type: ignore[type-var] |
| 68 | |
| 69 | |
| 70 | @pytest.mark.parametrize("model_id", ["openai-community/gpt2"]) |
| 71 | def test_cached_tokenizer_max_token_id_is_a_valid_id(model_id: str): |
| 72 | """max_token_id is the largest valid id, one less than vocab_size.""" |
| 73 | reference_tokenizer = AutoTokenizer.from_pretrained(model_id) |
| 74 | cached_tokenizer = get_cached_tokenizer(deepcopy(reference_tokenizer)) |
| 75 | |
| 76 | expected = max( |
| 77 | max(reference_tokenizer.get_vocab().values()), |
| 78 | reference_tokenizer.vocab_size - 1, |
| 79 | ) |
| 80 | assert cached_tokenizer.max_token_id == expected |
| 81 | assert cached_tokenizer.max_token_id == max( |
| 82 | reference_tokenizer.get_vocab().values() |
| 83 | ) |