Merge f0b869cf0af7f709b5556cd9cdd6e6de51a94378 into 553741557a57df44a46fb3ab72b6aafeaf945337

deepinfer/hugging_face/tokenizer.py

@@ -101,6 +101,8 @@ def get_hf_tokenizer(
     tokenizer_config = get_tokenizer_config(str(tokenizer_name), **kwargs)
 
     if "auto_map" in tokenizer_config:
+        if not trust_remote_code:
+            raise ValueError("Loading tokenizer code from the checkpoint requires trust_remote_code=True")
         class_ref = tokenizer_config["auto_map"]["AutoTokenizer"][0]
         module_file, class_name = class_ref.split(".")
         file_path = os.path.join(str(tokenizer_name), module_file + ".py")

tests/test_tokenizer.py

@@ -0,0 +1,74 @@
+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()