Skip to content

Commit 70675c3

Browse files
larryliu0820facebook-github-bot
larryliu0820
authored and
facebook-github-bot
committed
Migrate users of llm tokenizer to use pytorch-labs/tokenizers
Summary: Finally migrate llm tokenizer usages to pytorch-labs/tokenizers. Differential Revision: D70932091
1 parent cf8ce89 commit 70675c3

File tree

13 files changed

+47
-27
lines changed

13 files changed

+47
-27
lines changed

‎examples/models/llama/runner/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -43,7 +43,7 @@ target_include_directories(
4343

4444
list(
4545
APPEND _llama_runner__srcs
46-
${CMAKE_CURRENT_SOURCE_DIR}/../../../../extension/llm/tokenizer/tiktoken.cpp
46+
${CMAKE_CURRENT_SOURCE_DIR}/../../../../extension/llm/tokenizers/src/tiktoken.cpp
4747
)
4848
list(APPEND _llama_runner__srcs
4949
${CMAKE_CURRENT_SOURCE_DIR}/../tokenizer/llama_tiktoken.cpp

‎examples/models/llama/runner/runner.cpp

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -16,7 +16,7 @@
1616
#include <executorch/extension/llm/runner/util.h>
1717

1818
#include <executorch/examples/models/llama/tokenizer/llama_tiktoken.h>
19-
#include <executorch/extension/llm/tokenizer/bpe_tokenizer.h>
19+
#include <pytorch/tokenizers/llama2c_tokenizer.h>
2020

2121
namespace example {
2222

@@ -78,16 +78,16 @@ Error Runner::load() {
7878
// load tokenizer. Assuming tiktoken is the default tokenizer
7979
tokenizer_ = nullptr;
8080
tokenizer_ = get_tiktoken_for_llama();
81-
Error err = tokenizer_->load(tokenizer_path_);
81+
::tokenizers::Error err = tokenizer_->load(tokenizer_path_);
8282
// Rely on tiktoken to throw error if the artifact is incompatible. Then we
8383
// fallback to BPE tokenizer.
84-
if (err == Error::InvalidArgument) {
84+
if (err == ::tokenizers::Error::LoadFailure) {
8585
ET_LOG(
8686
Info,
8787
"Failed to load %s as a Tiktoken artifact, trying BPE tokenizer",
8888
tokenizer_path_.c_str());
8989
tokenizer_.reset();
90-
tokenizer_ = std::make_unique<llm::BPETokenizer>();
90+
tokenizer_ = std::make_unique<::tokenizers::Llama2cTokenizer>();
9191
tokenizer_->load(tokenizer_path_);
9292
}
9393

@@ -201,12 +201,12 @@ Error Runner::generate(
201201
? seq_len
202202
: metadata_.at(kMaxSeqLen);
203203

204-
Result<std::vector<uint64_t>> encode_res = tokenizer_->encode(
204+
::tokenizers::Result<std::vector<uint64_t>> encode_res = tokenizer_->encode(
205205
prompt,
206206
/* bos */ 0,
207207
/* eos */ 0);
208208

209-
ET_CHECK_OK_OR_RETURN_ERROR(
209+
ET_CHECK_TK_OK_OR_RETURN_ERROR(
210210
encode_res.error(), "Failed to encode prompt %s", prompt.c_str());
211211

212212
// encode the (string) prompt into tokens sequence
@@ -242,7 +242,7 @@ Error Runner::generate(
242242
uint64_t cur_token = prefill_res.get();
243243

244244
// print the first token from prefill. No prev_token so use cur_token for it.
245-
wrapped_callback(ET_UNWRAP(tokenizer_->decode(cur_token, cur_token)));
245+
wrapped_callback(ET_UNWRAP_TOKENIZER(tokenizer_->decode(cur_token, cur_token)));
246246
RUNNER_ET_LOG(
247247
warmup,
248248
"RSS after prompt prefill: %f MiB (0 if unsupported)",

‎examples/models/llama/runner/runner.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@
2323
#include <executorch/extension/llm/runner/text_decoder_runner.h>
2424
#include <executorch/extension/llm/runner/text_prefiller.h>
2525
#include <executorch/extension/llm/runner/text_token_generator.h>
26-
#include <executorch/extension/llm/tokenizer/tokenizer.h>
26+
#include <pytorch/tokenizers/tokenizer.h>
2727
#include <executorch/extension/module/module.h>
2828

2929
namespace example {
@@ -58,7 +58,7 @@ class ET_EXPERIMENTAL Runner : public executorch::extension::llm::IRunner {
5858
// model
5959
std::unique_ptr<::executorch::extension::Module> module_;
6060
std::string tokenizer_path_;
61-
std::unique_ptr<::executorch::extension::llm::Tokenizer> tokenizer_;
61+
std::unique_ptr<::tokenizers::Tokenizer> tokenizer_;
6262
std::unordered_map<std::string, int64_t> metadata_;
6363
std::unique_ptr<::executorch::extension::llm::TextDecoderRunner>
6464
text_decoder_runner_;

‎examples/models/llama/runner/targets.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -48,7 +48,7 @@ def define_common_targets():
4848
"//executorch/runtime/core/exec_aten:lib" + aten_suffix,
4949
"//executorch/runtime/core/exec_aten/util:tensor_util" + aten_suffix,
5050
"//executorch/examples/models/llama/tokenizer:tiktoken",
51-
"//executorch/extension/llm/tokenizer:bpe_tokenizer",
51+
"//pytorch/tokenizers:llama2c_tokenizer",
5252
] + (_get_operator_lib(aten)) + ([
5353
# Vulkan API currently cannot build on some platforms (e.g. Apple, FBCODE)
5454
# Therefore enable it explicitly for now to avoid failing tests

‎examples/models/llama/tokenizer/llama_tiktoken.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,7 @@
1010

1111
namespace example {
1212

13-
using ::executorch::extension::llm::Tiktoken;
13+
using ::tokenizers::Tiktoken;
1414

1515
namespace {
1616
static constexpr int32_t kSpecialTokensSize = 256;

‎examples/models/llama/tokenizer/llama_tiktoken.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -8,7 +8,7 @@
88

99
#pragma once
1010

11-
#include <executorch/extension/llm/tokenizer/tiktoken.h>
11+
#include <pytorch/tokenizers/tiktoken.h>
1212

1313
namespace example {
1414

@@ -17,7 +17,7 @@ enum class Version {
1717
Multimodal,
1818
};
1919

20-
std::unique_ptr<::executorch::extension::llm::Tiktoken> get_tiktoken_for_llama(
20+
std::unique_ptr<::tokenizers::Tiktoken> get_tiktoken_for_llama(
2121
Version version = Version::Default);
2222

2323
} // namespace example

‎examples/models/llama/tokenizer/targets.bzl

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ def define_common_targets():
1515
"llama_tiktoken.h",
1616
],
1717
exported_deps = [
18-
"//executorch/extension/llm/tokenizer:tiktoken",
18+
"//pytorch/tokenizers:tiktoken",
1919
],
2020
visibility = [
2121
"@EXECUTORCH_CLIENTS",

‎examples/models/llava/runner/CMakeLists.txt

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,7 @@ set(_common_include_directories ${EXECUTORCH_ROOT}/..)
2929
set(_llava_runner__srcs
3030
"${CMAKE_CURRENT_SOURCE_DIR}/llava_runner.cpp"
3131
"${EXECUTORCH_ROOT}/extension/llm/sampler/sampler.cpp"
32-
"${EXECUTORCH_ROOT}/extension/llm/tokenizer/bpe_tokenizer.cpp"
32+
"${EXECUTORCH_ROOT}/extension/llm/tokenizers/src/llama2c_tokenizer.cpp"
3333
)
3434

3535
# extension llm runner lib
@@ -47,5 +47,6 @@ set(llava_runner_deps executorch extension_data_loader extension_llm_runner
4747
target_link_libraries(llava_runner PUBLIC ${llava_runner_deps})
4848

4949
target_include_directories(
50-
llava_runner INTERFACE ${_common_include_directories} ${EXECUTORCH_ROOT}
50+
llava_runner INTERFACE ${_common_include_directories}
51+
${EXECUTORCH_ROOT}/extension/llm/tokenizers/include
5152
)

‎extension/llm/runner/CMakeLists.txt

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -51,5 +51,5 @@ target_link_libraries(extension_llm_runner PUBLIC ${runner_deps})
5151

5252
target_include_directories(
5353
extension_llm_runner INTERFACE ${_common_include_directories}
54-
${EXECUTORCH_ROOT}
54+
${EXECUTORCH_ROOT}/extension/llm/tokenizers/include
5555
)

‎extension/llm/runner/multimodal_runner.h

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,8 +26,8 @@
2626
#include <executorch/extension/llm/runner/text_prefiller.h>
2727
#include <executorch/extension/llm/runner/text_token_generator.h>
2828
#include <executorch/extension/llm/sampler/sampler.h>
29-
#include <executorch/extension/llm/tokenizer/tokenizer.h>
3029
#include <executorch/extension/module/module.h>
30+
#include <pytorch/tokenizers/tokenizer.h>
3131

3232
namespace executorch {
3333
namespace extension {
@@ -129,7 +129,7 @@ class ET_EXPERIMENTAL MultimodalRunner {
129129
std::unique_ptr<ImagePrefiller> image_prefiller_;
130130
std::unique_ptr<TextTokenGenerator> text_token_generator_;
131131
std::string tokenizer_path_;
132-
std::unique_ptr<Tokenizer> tokenizer_;
132+
std::unique_ptr<::tokenizers::Tokenizer> tokenizer_;
133133

134134
// stats
135135
Stats stats_;

‎extension/llm/runner/targets.bzl

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,7 @@ def define_common_targets():
4949
],
5050
exported_deps = [
5151
":text_decoder_runner" + aten_suffix,
52-
"//executorch/extension/llm/tokenizer:tokenizer_header",
52+
"//pytorch/tokenizers:headers",
5353
"//executorch/extension/module:module" + aten_suffix,
5454
"//executorch/extension/tensor:tensor" + aten_suffix,
5555
],
@@ -63,7 +63,7 @@ def define_common_targets():
6363
],
6464
exported_deps = [
6565
":text_decoder_runner" + aten_suffix,
66-
"//executorch/extension/llm/tokenizer:tokenizer_header",
66+
"//pytorch/tokenizers:headers",
6767
"//executorch/extension/module:module" + aten_suffix,
6868
"//executorch/extension/tensor:tensor" + aten_suffix,
6969
],

‎extension/llm/runner/text_token_generator.h

Lines changed: 24 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -11,17 +11,37 @@
1111

1212
#include <executorch/extension/llm/runner/stats.h>
1313
#include <executorch/extension/llm/runner/text_decoder_runner.h>
14-
#include <executorch/extension/llm/tokenizer/tokenizer.h>
14+
#include <pytorch/tokenizers/tokenizer.h>
1515
#include <executorch/extension/tensor/tensor.h>
1616

17+
#define ET_UNWRAP_TOKENIZER(result__) \
18+
({ \
19+
auto tk_result__ = (result__); \
20+
if (!tk_result__.ok()) { \
21+
ET_LOG(Error, "Tokenizers error code %d", static_cast<uint32_t>(tk_result__.error())); \
22+
return ::executorch::runtime::Error::InvalidArgument; \
23+
} \
24+
std::move(*tk_result__); \
25+
})
26+
27+
#define ET_CHECK_TK_OK_OR_RETURN_ERROR(result__, ...) \
28+
({ \
29+
auto tk_result__ = (result__); \
30+
if (tk_result__ != ::tokenizers::Error::Ok) { \
31+
ET_LOG(Error, "Tokenizer error: %d", static_cast<uint32_t>(tk_result__)); \
32+
ET_LOG(Error, __VA_ARGS__); \
33+
return ::executorch::runtime::Error::InvalidArgument; \
34+
} \
35+
})
36+
1737
namespace executorch {
1838
namespace extension {
1939
namespace llm {
2040

2141
class ET_EXPERIMENTAL TextTokenGenerator {
2242
public:
2343
TextTokenGenerator(
24-
Tokenizer* tokenizer,
44+
::tokenizers::Tokenizer* tokenizer,
2545
TextDecoderRunner* text_decoder_runner,
2646
bool use_kv_cache,
2747
std::unique_ptr<std::unordered_set<uint64_t>>&& eos_ids,
@@ -106,7 +126,7 @@ class ET_EXPERIMENTAL TextTokenGenerator {
106126
}
107127

108128
// print the token as string, decode it with the Tokenizer object
109-
token_callback(ET_UNWRAP(tokenizer_->decode(prev_token, cur_token)));
129+
token_callback(ET_UNWRAP_TOKENIZER(tokenizer_->decode(prev_token, cur_token)));
110130

111131
if (should_stop_) {
112132
break;
@@ -130,7 +150,7 @@ class ET_EXPERIMENTAL TextTokenGenerator {
130150
}
131151

132152
private:
133-
Tokenizer* tokenizer_;
153+
::tokenizers::Tokenizer* tokenizer_;
134154
TextDecoderRunner* text_decoder_runner_;
135155
std::unique_ptr<std::unordered_set<uint64_t>> eos_ids_;
136156
bool use_kv_cache_;

‎shim_et/xplat/executorch/build/env_interface.bzl

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -46,7 +46,6 @@ _EXTERNAL_DEPS = {
4646
"re2": "//extension/llm/tokenizers/third-party:re2",
4747
"sentencepiece": [], # Intentionally not supporting OSS buck build of sentencepiece.
4848
"sentencepiece-py": [],
49-
"tiktoken": "//extension/llm/tokenizers:tiktoken",
5049
# Core C++ PyTorch functionality like Tensor and ScalarType.
5150
"torch-core-cpp": "//third-party:libtorch",
5251
"torchgen": "//third-party:torchgen",

0 commit comments

Comments
 (0)