Move tokenizer.py into extension/llm/tokenizer

helunwencser · helunwencser · commit cf838ae28f78 · 2024-07-15T11:13:33.000-07:00
Pull Request resolved: #4255 as title ghstack-source-id: 233778092 Differential Revision: [D59698943](https://our.internmc.facebook.com/intern/diff/D59698943/)
diff --git a/.ci/scripts/test_llama.sh b/.ci/scripts/test_llama.sh
@@ -171,7 +171,7 @@ $PYTHON_EXECUTABLE -m examples.models.llama2.export_llama ${EXPORT_ARGS}
 
 # Create tokenizer.bin.
 echo "Creating tokenizer.bin"
-$PYTHON_EXECUTABLE -m examples.models.llama2.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
+$PYTHON_EXECUTABLE -m extension.llm.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
 
 
 RUNTIME_ARGS="--model_path=${EXPORTED_MODEL_NAME} --tokenizer_path=tokenizer.bin --prompt=Once --temperature=0 --seq_len=10"
diff --git a/backends/vulkan/docs/android_demo.md b/backends/vulkan/docs/android_demo.md
@@ -65,7 +65,7 @@ echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps":
 Then, create a tokenizer binary file:
 
 ```shell
-python -m examples.models.llama2.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
+python -m extension.llm.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
 ```
 
 Finally, export the `stories110M.pt` file into an ExecuTorch program:
diff --git a/examples/models/llama2/README.md b/examples/models/llama2/README.md
@@ -90,7 +90,7 @@ You can export and run the original Llama 2 7B model.
 4. Create tokenizer.bin.
 
     ```
-    python -m examples.models.llama2.tokenizer.tokenizer -t <tokenizer.model> -o tokenizer.bin
+    python -m extension.llm.tokenizer.tokenizer -t <tokenizer.model> -o tokenizer.bin
     ```
 
 ### Option B: Download and export stories110M model
@@ -113,7 +113,7 @@ If you want to deploy and run a smaller model for educational purposes. From `ex
 4. Create tokenizer.bin.
 
     ```
-    python -m examples.models.llama2.tokenizer.tokenizer -t <tokenizer.model> -o tokenizer.bin
+    python -m extension.llm.tokenizer.tokenizer -t <tokenizer.model> -o tokenizer.bin
     ```
 
 ### Option C: Download and export Llama 3 8B instruct model
diff --git a/examples/models/llama2/eval_llama_lib.py b/examples/models/llama2/eval_llama_lib.py
@@ -15,11 +15,11 @@
     get_quantizer_and_quant_params,
 )
 from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
-from executorch.examples.models.llama2.tokenizer.tokenizer import (
-    Tokenizer as SentencePieceTokenizer,
-)
 
 from executorch.extension.llm.export import LLMEdgeManager
+from executorch.extension.llm.tokenizer.tokenizer import (
+    Tokenizer as SentencePieceTokenizer,
+)
 
 from lm_eval.api.model import LM
 
diff --git a/examples/models/llama2/evaluate/eager_eval.py b/examples/models/llama2/evaluate/eager_eval.py
@@ -10,7 +10,7 @@
 import lm_eval
 import torch
 from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
-from executorch.examples.models.llama2.tokenizer.tokenizer import (
+from executorch.extension.llm.tokenizer.tokenizer import (
     Tokenizer as SentencePieceTokenizer,
 )
 
diff --git a/examples/models/llama2/tokenizer/targets.bzl b/examples/models/llama2/tokenizer/targets.bzl
@@ -42,32 +42,3 @@ def define_common_targets():
             "re2",
         ],
     )
-
-    runtime.python_library(
-        name = "tokenizer_py_lib",
-        srcs = [
-            "__init__.py",
-            "tokenizer.py",
-        ],
-        base_module = "executorch.examples.models.llama2.tokenizer",
-        visibility = [
-            "//executorch/examples/...",
-            "//bento/...",
-            "//bento_kernels/...",
-        ],
-        _is_external_target = True,
-        deps = [] if runtime.is_oss else ["fbsource//third-party/pypi/sentencepiece:sentencepiece"],
-    )
-
-    runtime.python_binary(
-        name = "tokenizer_py",
-        main_module = "executorch.examples.models.llama2.tokenizer.tokenizer",
-        visibility = [
-            "//executorch/examples/...",
-            "fbsource//xplat/executorch/examples/...",
-        ],
-        _is_external_target = True,
-        deps = [
-            ":tokenizer_py_lib",
-        ],
-    )
diff --git a/examples/models/llama2/tokenizer/test/__init__.py b/examples/models/llama2/tokenizer/test/__init__.py
diff --git a/examples/models/llama2/tokenizer/test/targets.bzl b/examples/models/llama2/tokenizer/test/targets.bzl
@@ -42,16 +42,3 @@ def define_common_targets():
             "resources/**",
         ]),
     )
-
-    runtime.python_test(
-        name = "test_bpe_tokenizer_py",
-        srcs = [
-            "test_bpe_tokenizer.py",
-        ],
-        visibility = [
-            "//executorch/examples/...",
-        ],
-        deps = [
-            "//executorch/examples/models/llama2/tokenizer:tokenizer_py_lib",
-        ],
-    )
diff --git a/examples/qualcomm/llama2/README.md b/examples/qualcomm/llama2/README.md
@@ -23,7 +23,7 @@ wget "https://huggingface.co/karpathy/tinyllamas/resolve/main/stories110M.pt"
 wget "https://raw.githubusercontent.com/karpathy/llama2.c/master/tokenizer.model"
 
 # tokenizer.bin:
-python -m examples.models.llama2.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
+python -m extension.llm.tokenizer.tokenizer -t tokenizer.model -o tokenizer.bin
 
 # params.json:
 echo '{"dim": 768, "multiple_of": 32, "n_heads": 12, "n_layers": 12, "norm_eps": 1e-05, "vocab_size": 32000}' > params.json
diff --git a/extension/llm/tokenizer/TARGETS b/extension/llm/tokenizer/TARGETS
@@ -0,0 +1,8 @@
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()
diff --git a/extension/llm/tokenizer/__init__.py b/extension/llm/tokenizer/__init__.py
diff --git a/extension/llm/tokenizer/targets.bzl b/extension/llm/tokenizer/targets.bzl
@@ -0,0 +1,38 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+def define_common_targets():
+    """Defines targets that should be shared between fbcode and xplat.
+
+    The directory containing this targets.bzl file should also contain both
+    TARGETS and BUCK files that call this function.
+    """
+    runtime.python_library(
+        name = "tokenizer_py_lib",
+        srcs = [
+            "__init__.py",
+            "tokenizer.py",
+        ],
+        base_module = "executorch.extension.llm.tokenizer",
+        visibility = [
+            "//executorch/examples/...",
+            "//executorch/extension/llm/tokenizer/...",
+            "//bento/...",
+            "//bento_kernels/...",
+        ],
+        _is_external_target = True,
+        # TODO: Define an external_deps entry for sentencepiece instead of pointing to an fbsource path.
+        deps = [] if runtime.is_oss else ["fbsource//third-party/pypi/sentencepiece:sentencepiece"],
+    )
+
+    runtime.python_binary(
+        name = "tokenizer_py",
+        main_module = "executorch.extension.llm.tokenizer.tokenizer",
+        visibility = [
+            "//executorch/examples/...",
+            "fbsource//xplat/executorch/examples/...",
+        ],
+        _is_external_target = True,
+        deps = [
+            ":tokenizer_py_lib",
+        ],
+    )
diff --git a/extension/llm/tokenizer/test/TARGETS b/extension/llm/tokenizer/test/TARGETS
@@ -0,0 +1,8 @@
+# Any targets that should be shared between fbcode and xplat must be defined in
+# targets.bzl. This file can contain fbcode-only targets.
+
+load(":targets.bzl", "define_common_targets")
+
+oncall("executorch")
+
+define_common_targets()
diff --git a/extension/llm/tokenizer/test/targets.bzl b/extension/llm/tokenizer/test/targets.bzl
@@ -0,0 +1,18 @@
+load("@fbsource//xplat/executorch/build:runtime_wrapper.bzl", "runtime")
+
+def define_common_targets():
+    """Defines targets that should be shared between fbcode and xplat.
+
+    The directory containing this targets.bzl file should also contain both
+    TARGETS and BUCK files that call this function.
+    """
+
+    runtime.python_test(
+        name = "test_tokenizer_py",
+        srcs = [
+            "test_tokenizer.py",
+        ],
+        deps = [
+            "//executorch/extension/llm/tokenizer:tokenizer_py_lib",
+        ],
+    )
diff --git a/extension/llm/tokenizer/test/test_tokenizer.py b/extension/llm/tokenizer/test/test_tokenizer.py
@@ -10,13 +10,11 @@
 import unittest
 from unittest.mock import patch
 
-from executorch.examples.models.llama2.tokenizer.tokenizer import Tokenizer
+from executorch.extension.llm.tokenizer.tokenizer import Tokenizer
 
 
 class TestTokenizer(unittest.TestCase):
-    @patch(
-        "executorch.examples.models.llama2.tokenizer.tokenizer.SentencePieceProcessor"
-    )
+    @patch("executorch.extension.llm.tokenizer.tokenizer.SentencePieceProcessor")
     def test_export(self, mock_sp):
         # Set up the mock SentencePieceProcessor
         mock_sp.return_value.vocab_size.return_value = 0
diff --git a/extension/llm/tokenizer/tokenizer.py b/extension/llm/tokenizer/tokenizer.py

Original file line number	Diff line number	Diff line change
`@@ -10,7 +10,7 @@`
`10`	`10`	`import lm_eval`
`11`	`11`	`import torch`
`12`	`12`	`from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken`
`13`		`-from executorch.examples.models.llama2.tokenizer.tokenizer import (`
	`13`	`+from executorch.extension.llm.tokenizer.tokenizer import (`
`14`	`14`	`Tokenizer as SentencePieceTokenizer,`
`15`	`15`	`)`
`16`	`16`