Accept custom pattern string and special tokens

Differential Revision: D74264910

Pull Request resolved: #69

Author: sxu

include/pytorch/tokenizers/tiktoken.h

@@ -32,20 +32,33 @@ static constexpr size_t kEOSTokenIndex = 1;
 class Tiktoken : public detail::BPETokenizerBase {
  public:
   explicit Tiktoken(
+      std::string pattern,
       std::unique_ptr<std::vector<std::string>> special_tokens,
       size_t bos_token_index,
       size_t eos_token_index)
-      : _special_tokens(std::move(special_tokens)),
+      : _pattern(std::move(pattern)),
+        _special_tokens(std::move(special_tokens)),
         _bos_token_index(bos_token_index),
         _eos_token_index(eos_token_index) {
     if (_bos_token_index >= _special_tokens->size() ||
         _eos_token_index >= _special_tokens->size()) {
       abort();
     }
-  };
+  }
+
+  explicit Tiktoken(
+      std::unique_ptr<std::vector<std::string>> special_tokens,
+      size_t bos_token_index,
+      size_t eos_token_index)
+      : Tiktoken(
+            _get_default_patern(),
+            std::move(special_tokens),
+            bos_token_index,
+            eos_token_index) {}
 
   explicit Tiktoken()
-      : _special_tokens(_get_default_special_tokens()),
+      : _pattern(_get_default_patern()),
+        _special_tokens(_get_default_special_tokens()),
         _bos_token_index(kBOSTokenIndex),
         _eos_token_index(kEOSTokenIndex){};
 
@@ -77,6 +90,11 @@ class Tiktoken : public detail::BPETokenizerBase {
     return special_tokens;
   }
 
+  static inline std::string _get_default_patern() {
+    // Removed negative lookahead \s+(?!\S) since it's not supported by RE2.
+    return R"((?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+)";
+  }
+
   Error _encode(
       const std::string& input,
       std::vector<uint64_t>& ret,
@@ -86,14 +104,11 @@ class Tiktoken : public detail::BPETokenizerBase {
 
   detail::TokenMap _build_special_token_map(ssize_t num_base_tokens) const;
 
+  std::string _pattern;
   std::unique_ptr<std::vector<std::string>> _special_tokens;
   size_t _bos_token_index;
   size_t _eos_token_index;
 
-  // Removed negative lookahead \s+(?!\S) since it's not supported by RE2.
-  const std::string _pattern =
-      R"((?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+)";
-
   std::unique_ptr<IRegex> _regex;
 };
 

pytorch_tokenizers/constants.py

@@ -0,0 +1,31 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+# @lint-ignore-every LICENSELINT
+
+CL100K_PAT_STR = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"  # noqa: E501
+
+LLAMA_BASIC_SPECIAL_TOKENS = [
+    "<|begin_of_text|>",
+    "<|end_of_text|>",
+    "<|reserved_special_token_0|>",
+    "<|reserved_special_token_1|>",
+    "<|finetune_right_pad_id|>",
+    "<|step_id|>",
+    "<|start_header_id|>",
+    "<|end_header_id|>",
+    "<|eom_id|>",  # end of message
+    "<|eot_id|>",  # end of turn
+    "<|python_tag|>",
+    "<|image|>",
+]
+
+LLAMA_NUM_RESERVED_SPECIAL_TOKENS = 256
+LLAMA_RESERVED_SPECIAL_TOKENS = [
+    f"<|reserved_special_token_{2 + i}|>"
+    for i in range(LLAMA_NUM_RESERVED_SPECIAL_TOKENS - len(LLAMA_BASIC_SPECIAL_TOKENS))
+]
+
+LLAMA_SPECIAL_TOKENS = LLAMA_BASIC_SPECIAL_TOKENS + LLAMA_RESERVED_SPECIAL_TOKENS

pytorch_tokenizers/targets.bzl

@@ -10,6 +10,7 @@ def define_common_targets():
         name = "tokenizers",
         srcs = [
             "__init__.py",
+            "constants.py",
             "llama2c.py",
             "tiktoken.py",
             "hf_tokenizer.py",

pytorch_tokenizers/tiktoken.py

@@ -25,6 +25,8 @@
 
 from tiktoken.load import load_tiktoken_bpe
 
+from .constants import CL100K_PAT_STR, LLAMA_SPECIAL_TOKENS
+
 logger = getLogger(__name__)
 
 
@@ -47,12 +49,6 @@ class TiktokenTokenizer:
     WARNING: The regex and special tokens are hardcoded from Llama 3+.
     """
 
-    special_tokens: Dict[str, int]
-
-    num_reserved_special_tokens = 256
-
-    pat_str = r"(?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+(?!\S)|\s+"  # noqa: E501
-
     @classmethod
     def get_instance(cls):
         global _INSTANCE
@@ -63,7 +59,12 @@ def get_instance(cls):
             )
         return _INSTANCE
 
-    def __init__(self, model_path: str):
+    def __init__(
+        self,
+        model_path: str,
+        pat_str: str = CL100K_PAT_STR,
+        special_tokens: List[str] = LLAMA_SPECIAL_TOKENS,
+    ):
         """
         Initializes the Tokenizer with a Tiktoken model.
 
@@ -74,32 +75,13 @@ def __init__(self, model_path: str):
 
         mergeable_ranks = load_tiktoken_bpe(model_path)
         num_base_tokens = len(mergeable_ranks)
-        special_tokens = [
-            "<|begin_of_text|>",
-            "<|end_of_text|>",
-            "<|reserved_special_token_0|>",
-            "<|reserved_special_token_1|>",
-            "<|finetune_right_pad_id|>",
-            "<|step_id|>",
-            "<|start_header_id|>",
-            "<|end_header_id|>",
-            "<|eom_id|>",  # end of message
-            "<|eot_id|>",  # end of turn
-            "<|python_tag|>",
-            "<|image|>",
-        ]
-        reserved_tokens = [
-            f"<|reserved_special_token_{2 + i}|>"
-            for i in range(self.num_reserved_special_tokens - len(special_tokens))
-        ]
-        special_tokens = special_tokens + reserved_tokens
 
         self.special_tokens = {
             token: num_base_tokens + i for i, token in enumerate(special_tokens)
         }
         self.model = tiktoken.Encoding(
             name=Path(model_path).name,
-            pat_str=self.pat_str,
+            pat_str=pat_str,
             mergeable_ranks=mergeable_ranks,
             special_tokens=self.special_tokens,
         )
@@ -108,15 +90,6 @@ def __init__(self, model_path: str):
         # BOS / EOS token IDs
         self.bos_id: int = self.special_tokens["<|begin_of_text|>"]
         self.eos_id: int = self.special_tokens["<|end_of_text|>"]
-        self.eot_id: int = self.special_tokens["<|eot_id|>"]
-        self.eom_id: int = self.special_tokens["<|eom_id|>"]
-        self.python_tag_id = self.special_tokens["<|python_tag|>"]
-        self.pad_id: int = self.special_tokens["<|finetune_right_pad_id|>"]
-        self.stop_tokens = [
-            self.eos_id,
-            self.special_tokens["<|eom_id|>"],
-            self.special_tokens["<|eot_id|>"],
-        ]
 
     def encode(
         self,

test/targets.bzl

@@ -108,3 +108,16 @@ def define_common_targets():
         src = "resources/test_tiktoken_tokenizer.model",
         visibility = ["@EXECUTORCH_CLIENTS", "//pytorch/tokenizers/..."],
     )
+
+    runtime.python_test(
+        name = "test_tiktoken_py",
+        srcs = [
+            "test_tiktoken.py",
+        ],
+        deps = [
+            "//pytorch/tokenizers/pytorch_tokenizers:tokenizers",
+        ],
+        resources = {
+            ":test_tiktoken_tokenizer_model": "test_tiktoken_tokenizer.model",
+        },
+    )

test/test_tiktoken.cpp

@@ -16,6 +16,8 @@ namespace tokenizers {
 
 namespace {
 // Test case based on Llama 2
+const std::string kPattern =
+    R"((?i:'s|'t|'re|'ve|'m|'ll|'d)|[^\r\n\p{L}\p{N}]?\p{L}+|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n]*|\s*[\r\n]+|\s+)";
 static constexpr int32_t kSpecialTokensSize = 256;
 static inline std::unique_ptr<std::vector<std::string>> _get_special_tokens() {
   auto special_tokens =
@@ -50,7 +52,8 @@ static inline std::string _get_resource_path(const std::string& name) {
 class TiktokenTest : public Test {
  public:
   void SetUp() override {
-    tokenizer_ = std::make_unique<Tiktoken>(_get_special_tokens(), 0, 1);
+    tokenizer_ =
+        std::make_unique<Tiktoken>(kPattern, _get_special_tokens(), 0, 1);
     modelPath_ = _get_resource_path("test_tiktoken_tokenizer.model");
   }
 

test/test_tiktoken.py

@@ -0,0 +1,48 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+# @lint-ignore-every LICENSELINT
+
+import unittest
+
+import pkg_resources
+
+from pytorch_tokenizers.tiktoken import TiktokenTokenizer
+
+
+class TestTiktokenTokenizer(unittest.TestCase):
+    def test_default(self):
+        model_path = pkg_resources.resource_filename(
+            "pytorch.tokenizers.test", "test_tiktoken_tokenizer.model"
+        )
+        tiktoken = TiktokenTokenizer(model_path)
+        s = "<|begin_of_text|> hellow world."
+        self.assertEqual(s, tiktoken.decode(tiktoken.encode(s, bos=False, eos=False)))
+
+    def test_custom_pattern_and_special_tokens(self):
+        o220k_pattern = r"""[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]*[\p{Ll}\p{Lm}\p{Lo}\p{M}]+(?i:'s|'t|'re|'ve|'m|'ll|'d)?|[^\r\n\p{L}\p{N}]?[\p{Lu}\p{Lt}\p{Lm}\p{Lo}\p{M}]+[\p{Ll}\p{Lm}\p{Lo}\p{M}]*(?i:'s|'t|'re|'ve|'m|'ll|'d)?|\p{N}{1,3}| ?[^\s\p{L}\p{N}]+[\r\n/]*|\s*[\r\n]+|\s+(?!\S)|\s+"""
+        model_path = pkg_resources.resource_filename(
+            "pytorch.tokenizers.test", "test_tiktoken_tokenizer.model"
+        )
+        tiktoken = TiktokenTokenizer(
+            model_path,
+            pat_str=o220k_pattern,
+            special_tokens=[
+                "<|begin_of_text|>",
+                "<|end_of_text|>",
+                "<|custom_token|>",
+            ],
+        )
+        custom_token_id = tiktoken.special_tokens["<|custom_token|>"]
+
+        s = "<|begin_of_text|> hellow world, this is a custom token: <|custom_token|>."
+        encoding = tiktoken.encode(
+            s,
+            bos=False,
+            eos=False,
+            allowed_special="all",
+        )
+        self.assertTrue(custom_token_id in encoding)
+        self.assertEqual(s, tiktoken.decode(encoding))