mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-29 23:37:28 +08:00
75 lines
2.9 KiB
Python
75 lines
2.9 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
from pathlib import Path
|
|
|
|
import pytest
|
|
from tokenizers import Tokenizer
|
|
from tokenizers.models import WordLevel
|
|
from tokenizers.pre_tokenizers import Whitespace
|
|
|
|
from deepinfer.hugging_face.tokenizer import get_hf_tokenizer
|
|
|
|
pytestmark = [pytest.mark.ci, pytest.mark.ci_npu]
|
|
|
|
|
|
@pytest.fixture
|
|
def checkpoint(tmp_path: Path) -> Path:
|
|
tokenizer = Tokenizer(WordLevel({"[UNK]": 0, "hello": 1, "world": 2}, unk_token="[UNK]"))
|
|
tokenizer.pre_tokenizer = Whitespace()
|
|
tokenizer.save(str(tmp_path / "tokenizer.json"))
|
|
(tmp_path / "config.json").write_text(json.dumps({"tokenizer_compatibility": "native"}))
|
|
(tmp_path / "tokenizer_config.json").write_text(
|
|
json.dumps(
|
|
{
|
|
"tokenizer_class": "CheckpointTokenizer",
|
|
"unk_token": "[UNK]",
|
|
"auto_map": {"AutoTokenizer": ["custom_tokenizer.CheckpointTokenizer", None]},
|
|
}
|
|
)
|
|
)
|
|
(tmp_path / "custom_tokenizer.py").write_text(
|
|
"from pathlib import Path\n"
|
|
"from transformers import PreTrainedTokenizerFast\n"
|
|
"Path(__file__).with_suffix('.loaded').write_text('imported')\n"
|
|
"class CheckpointTokenizer(PreTrainedTokenizerFast):\n"
|
|
" pass\n"
|
|
)
|
|
return tmp_path
|
|
|
|
|
|
@pytest.mark.parametrize("compatibility", ["native", "release"])
|
|
@pytest.mark.parametrize("explicit_false", [False, True], ids=["default", "false"])
|
|
def test_checkpoint_code_requires_trust(checkpoint: Path, compatibility: str, explicit_false: bool) -> None:
|
|
if compatibility == "release":
|
|
(checkpoint / "config.json").unlink() # Exercise the default release policy.
|
|
kwargs = {"trust_remote_code": False} if explicit_false else {}
|
|
|
|
with pytest.raises(ValueError, match="trust_remote_code=True"):
|
|
get_hf_tokenizer(str(checkpoint), **kwargs)
|
|
|
|
assert not (checkpoint / "custom_tokenizer.loaded").exists()
|
|
|
|
|
|
def test_trusted_checkpoint_tokenizer_loads(checkpoint: Path) -> None:
|
|
tokenizer = get_hf_tokenizer(str(checkpoint), trust_remote_code=True)
|
|
|
|
assert (checkpoint / "custom_tokenizer.loaded").read_text() == "imported"
|
|
assert type(tokenizer).__name__ == "CheckpointTokenizer"
|
|
assert tokenizer.encode("hello world", add_special_tokens=False) == [1, 2]
|
|
|
|
|
|
@pytest.mark.parametrize("explicit_false", [False, True], ids=["default", "false"])
|
|
def test_builtin_tokenizer_does_not_require_trust(checkpoint: Path, explicit_false: bool) -> None:
|
|
config_path = checkpoint / "tokenizer_config.json"
|
|
config = json.loads(config_path.read_text())
|
|
config.pop("auto_map")
|
|
config["tokenizer_class"] = "PreTrainedTokenizerFast"
|
|
config_path.write_text(json.dumps(config))
|
|
kwargs = {"trust_remote_code": False} if explicit_false else {}
|
|
|
|
tokenizer = get_hf_tokenizer(str(checkpoint), **kwargs)
|
|
|
|
assert tokenizer.encode("hello world", add_special_tokens=False) == [1, 2]
|
|
assert not (checkpoint / "custom_tokenizer.loaded").exists()
|