From fcb372a2f10580c0b8f9a5cd5c15814b281ea16a Mon Sep 17 00:00:00 2001
From: Chen Lai <chenlai@fb.com>
Date: Mon, 29 Jul 2024 15:43:15 -0700
Subject: [PATCH 1/3] move get tokenizer to export_llama_lib

Differential Revision: [D60419366](https://our.internmc.facebook.com/intern/diff/D60419366/)

[ghstack-poisoned]
---
 examples/models/llama2/eval_llama_lib.py   | 11 ++---------
 examples/models/llama2/export_llama_lib.py | 14 ++++++++++++++
 2 files changed, 16 insertions(+), 9 deletions(-)

diff --git a/examples/models/llama2/eval_llama_lib.py b/examples/models/llama2/eval_llama_lib.py
index 5e23a250026..ee3b8ddab18 100644
--- a/examples/models/llama2/eval_llama_lib.py
+++ b/examples/models/llama2/eval_llama_lib.py
@@ -13,13 +13,10 @@
 from executorch.examples.models.llama2.evaluate import EagerEvalWrapper, evaluate_model
 from executorch.examples.models.llama2.export_llama_lib import (
     get_quantizer_and_quant_params,
+    get_tokenizer,
 )
-from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
 
 from executorch.extension.llm.export import LLMEdgeManager
-from executorch.extension.llm.tokenizer.tokenizer import (
-    Tokenizer as SentencePieceTokenizer,
-)
 
 from lm_eval.api.model import LM
 
@@ -103,11 +100,7 @@ def gen_eval_wrapper(
     Returns:
         eval_wrapper (LM): A wrapper interface for the lm-evaluation-harness library.
     """
-    try:
-        tokenizer = SentencePieceTokenizer(model_path=str(args.tokenizer_path))
-    except Exception:
-        print("Using Tiktokenizer")
-        tokenizer = Tiktoken(model_path=str(args.tokenizer_path))
+    tokenizer = get_tokenizer(args.tokenizer_path)
 
     # ExecuTorch Binary Evaluation
     if (model := args.pte) is not None:
diff --git a/examples/models/llama2/export_llama_lib.py b/examples/models/llama2/export_llama_lib.py
index cf5d1b6e6f9..75c91335a9b 100644
--- a/examples/models/llama2/export_llama_lib.py
+++ b/examples/models/llama2/export_llama_lib.py
@@ -23,6 +23,7 @@
 import torch
 
 from executorch.examples.models.llama2.llama_transformer import ModelArgs
+from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
 
 from executorch.extension.llm.export.builder import DType, LLMEdgeManager
 
@@ -40,6 +41,10 @@
     get_qnn_quantizer,
 )
 
+from executorch.extension.llm.tokenizer.tokenizer import (
+    Tokenizer as SentencePieceTokenizer,
+)
+
 from executorch.sdk.etrecord import generate_etrecord
 from executorch.util.activation_memory_profiler import generate_memory_trace
 
@@ -86,6 +91,15 @@ def verbose_export():
     return verbosity_setting
 
 
+def get_tokenizer(tokenizer_path):
+    try:
+        tokenizer = SentencePieceTokenizer(model_path=str(tokenizer_path))
+    except Exception:
+        print("Using Tiktokenizer")
+        tokenizer = Tiktoken(model_path=str(tokenizer_path))
+    return tokenizer
+
+
 def build_model(
     modelname: str = "model",
     extra_opts: str = "",

From 82fecf6178dc73510abd305419e0a391b889eeb0 Mon Sep 17 00:00:00 2001
From: Chen Lai <chenlai@fb.com>
Date: Tue, 30 Jul 2024 11:51:41 -0700
Subject: [PATCH 2/3] Update on "move get tokenizer to export_llama_lib"

Differential Revision: [D60419366](https://our.internmc.facebook.com/intern/diff/D60419366/)

[ghstack-poisoned]
---
 examples/models/llama2/eval_llama_lib.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/examples/models/llama2/eval_llama_lib.py b/examples/models/llama2/eval_llama_lib.py
index ee3b8ddab18..365c78f28cd 100644
--- a/examples/models/llama2/eval_llama_lib.py
+++ b/examples/models/llama2/eval_llama_lib.py
@@ -15,9 +15,12 @@
     get_quantizer_and_quant_params,
     get_tokenizer,
 )
+from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
 
 from executorch.extension.llm.export import LLMEdgeManager
-
+from executorch.extension.llm.tokenizer.tokenizer import (
+    Tokenizer as SentencePieceTokenizer,
+)
 from lm_eval.api.model import LM
 
 from .export_llama_lib import (

From 0177b5377659183609642869fa6a7cf4bc2e68d9 Mon Sep 17 00:00:00 2001
From: Chen Lai <chenlai@fb.com>
Date: Tue, 30 Jul 2024 12:25:34 -0700
Subject: [PATCH 3/3] Update on "move get tokenizer to export_llama_lib"

Differential Revision: [D60419366](https://our.internmc.facebook.com/intern/diff/D60419366/)

[ghstack-poisoned]
---
 examples/models/llama2/eval_llama_lib.py   |  2 +-
 examples/models/llama2/export_llama_lib.py | 14 --------------
 extension/llm/tokenizer/targets.bzl        | 18 ++++++++++++++++++
 extension/llm/tokenizer/utils.py           | 19 +++++++++++++++++++
 4 files changed, 38 insertions(+), 15 deletions(-)
 create mode 100644 extension/llm/tokenizer/utils.py

diff --git a/examples/models/llama2/eval_llama_lib.py b/examples/models/llama2/eval_llama_lib.py
index 365c78f28cd..d44820e1873 100644
--- a/examples/models/llama2/eval_llama_lib.py
+++ b/examples/models/llama2/eval_llama_lib.py
@@ -13,7 +13,6 @@
 from executorch.examples.models.llama2.evaluate import EagerEvalWrapper, evaluate_model
 from executorch.examples.models.llama2.export_llama_lib import (
     get_quantizer_and_quant_params,
-    get_tokenizer,
 )
 from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
 
@@ -21,6 +20,7 @@
 from executorch.extension.llm.tokenizer.tokenizer import (
     Tokenizer as SentencePieceTokenizer,
 )
+from executorch.extension.llm.tokenizer.utils import get_tokenizer
 from lm_eval.api.model import LM
 
 from .export_llama_lib import (
diff --git a/examples/models/llama2/export_llama_lib.py b/examples/models/llama2/export_llama_lib.py
index 75c91335a9b..cf5d1b6e6f9 100644
--- a/examples/models/llama2/export_llama_lib.py
+++ b/examples/models/llama2/export_llama_lib.py
@@ -23,7 +23,6 @@
 import torch
 
 from executorch.examples.models.llama2.llama_transformer import ModelArgs
-from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
 
 from executorch.extension.llm.export.builder import DType, LLMEdgeManager
 
@@ -41,10 +40,6 @@
     get_qnn_quantizer,
 )
 
-from executorch.extension.llm.tokenizer.tokenizer import (
-    Tokenizer as SentencePieceTokenizer,
-)
-
 from executorch.sdk.etrecord import generate_etrecord
 from executorch.util.activation_memory_profiler import generate_memory_trace
 
@@ -91,15 +86,6 @@ def verbose_export():
     return verbosity_setting
 
 
-def get_tokenizer(tokenizer_path):
-    try:
-        tokenizer = SentencePieceTokenizer(model_path=str(tokenizer_path))
-    except Exception:
-        print("Using Tiktokenizer")
-        tokenizer = Tiktoken(model_path=str(tokenizer_path))
-    return tokenizer
-
-
 def build_model(
     modelname: str = "model",
     extra_opts: str = "",
diff --git a/extension/llm/tokenizer/targets.bzl b/extension/llm/tokenizer/targets.bzl
index 023968165bb..f2476468b84 100644
--- a/extension/llm/tokenizer/targets.bzl
+++ b/extension/llm/tokenizer/targets.bzl
@@ -25,6 +25,24 @@ def define_common_targets():
         ],
     )
 
+    runtime.python_library(
+        name = "utils",
+        srcs = [
+            "utils.py",
+        ],
+        base_module = "executorch.extension.llm.utils",
+        visibility = [
+            "//executorch/examples/...",
+            "//executorch/extension/llm/tokenizer/...",
+            "//bento/...",
+            "//bento_kernels/...",
+        ],
+        _is_external_target = True,
+        external_deps = [
+            "sentencepiece-py",
+        ],
+    )
+
     runtime.python_binary(
         name = "tokenizer_py",
         main_module = "executorch.extension.llm.tokenizer.tokenizer",
diff --git a/extension/llm/tokenizer/utils.py b/extension/llm/tokenizer/utils.py
new file mode 100644
index 00000000000..763febdf478
--- /dev/null
+++ b/extension/llm/tokenizer/utils.py
@@ -0,0 +1,19 @@
+# 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.
+
+from executorch.examples.models.llama2.tokenizer.tiktoken import Tokenizer as Tiktoken
+from executorch.extension.llm.tokenizer.tokenizer import (
+    Tokenizer as SentencePieceTokenizer,
+)
+
+
+def get_tokenizer(tokenizer_path):
+    try:
+        tokenizer = SentencePieceTokenizer(model_path=str(tokenizer_path))
+    except Exception:
+        print("Using Tiktokenizer")
+        tokenizer = Tiktoken(model_path=str(tokenizer_path))
+    return tokenizer