Migrate to PyTorch tokenizers

packages/react-native-executorch/android/CMakeLists.txt

@@ -21,6 +21,7 @@ string(APPEND CMAKE_CXX_FLAGS " -DRCT_NEW_ARCH_ENABLED")
 set(ANDROID_CPP_DIR "${CMAKE_SOURCE_DIR}/src/main/cpp")
 set(COMMON_CPP_DIR "${CMAKE_SOURCE_DIR}/../common")
 set(LIBS_DIR "${CMAKE_SOURCE_DIR}/../third-party/android/libs")
+set(TOKENIZERS_DIR "${CMAKE_SOURCE_DIR}/../third-party/include/executorch/extension/llm/tokenizers/include")
 set(INCLUDE_DIR "${CMAKE_SOURCE_DIR}/../third-party/include")
 
 # Treat third-party headers as system headers to suppress deprecation warnings

packages/react-native-executorch/android/src/main/cpp/CMakeLists.txt

@@ -17,6 +17,7 @@ target_include_directories(
   "${COMMON_CPP_DIR}"
   "${ANDROID_CPP_DIR}"
   "${INCLUDE_DIR}"
+  "${TOKENIZERS_DIR}"
   "${REACT_NATIVE_DIR}/ReactCommon"
   "${REACT_NATIVE_DIR}/ReactAndroid/src/main/jni/react/turbomodule"
   "${REACT_NATIVE_DIR}/ReactCommon/callinvoker"
@@ -84,13 +85,6 @@ elseif(ANDROID_ABI STREQUAL "x86_64")
   set(OPENCV_THIRD_PARTY_LIBS "")
 endif()
 
-# ------- tokenizers-cpp -------
-
-set(TOKENIZERS_LIBS
-  "${LIBS_DIR}/tokenizers-cpp/${ANDROID_ABI}/libtokenizers_c.a"
-  "${LIBS_DIR}/tokenizers-cpp/${ANDROID_ABI}/libtokenizers_cpp.a"
-  "${LIBS_DIR}/tokenizers-cpp/${ANDROID_ABI}/libsentencepiece.a"
-)
 
 # ------- phonemis -------
 
@@ -108,8 +102,6 @@ target_link_libraries(
   ${RN_VERSION_LINK_LIBRARIES}
   ${OPENCV_LIBS}
   ${OPENCV_THIRD_PARTY_LIBS}
-  ${TOKENIZERS_LIBS}
-  ${TOKENIZERS_THIRD_PARTY_LIBS}
   ${PHONEMIS_LIBS}
   executorch
   ${EXECUTORCH_LIBS}

packages/react-native-executorch/common/rnexecutorch/ErrorCodes.h

@@ -67,14 +67,18 @@ enum class RnExecutorchErrorCode : int32_t {
   WrongDimensions = 116,
   /**
    * Thrown when the input passed to our APIs is invalid, for example when
-   * passing an empty message aray to LLM's generate().
+   * passing an empty message array to LLM's generate().
    */
   InvalidUserInput = 117,
   /**
    * Thrown when the number of downloaded files is unexpected, due to download
    * interruptions.
    */
   DownloadInterrupted = 118,
+  /**
+   * Thrown when an error occurs with the tokenizer or tokenization process.
+   */
+  TokenizerError = 167,
   /**
    * Thrown when there's a configuration mismatch between multilingual and
    * language settings in Speech-to-Text models.

packages/react-native-executorch/common/rnexecutorch/TokenizerModule.cpp

@@ -1,18 +1,28 @@
 #include "TokenizerModule.h"
 #include "Error.h"
 #include "ErrorCodes.h"
+#include <cstdint>
 #include <executorch/extension/module/module.h>
 #include <filesystem>
-#include <rnexecutorch/data_processing/FileUtils.h>
+#include <pytorch/tokenizers/error.h>
+#include <runner/constants.h>
 
 namespace rnexecutorch {
 using namespace facebook;
+using namespace executorch::extension::constants;
 
 TokenizerModule::TokenizerModule(
     std::string source, std::shared_ptr<react::CallInvoker> callInvoker)
-    : tokenizer(tokenizers::Tokenizer::FromBlobJSON(
-          file_utils::loadBytesFromFile(source))),
-      memorySizeLowerBound(std::filesystem::file_size(source)) {}
+    : tokenizer(std::make_unique<tokenizers::HFTokenizer>()),
+      memorySizeLowerBound(std::filesystem::file_size(source)) {
+
+  auto status = tokenizer->load(source);
+
+  if (status != tokenizers::Error::Ok) {
+    throw RnExecutorchError(RnExecutorchErrorCode::TokenizerError,
+                            "Unexpected issue occured while loading tokenizer");
+  };
+}
 
 void TokenizerModule::ensureTokenizerLoaded(
     const std::string &methodName) const {
@@ -23,31 +33,69 @@ void TokenizerModule::ensureTokenizerLoaded(
   }
 }
 
-std::vector<int32_t> TokenizerModule::encode(std::string s) const {
+std::vector<uint64_t> TokenizerModule::encode(std::string s) const {
   ensureTokenizerLoaded("encode");
-  return tokenizer->Encode(s);
+
+  // If the used tokenizer.json has defined post_processor field,
+  // setting any of bos or eos arguments to value other than provided constant
+  // ( which is 0) will result in running the post_processor with
+  // 'add_special_token' flag
+  auto encodeResult =
+      tokenizer->encode(s, numOfAddedBoSTokens, numOfAddedEoSTokens);
+  if (!encodeResult.ok()) {
+    throw rnexecutorch::RnExecutorchError(
+        rnexecutorch::RnExecutorchErrorCode::TokenizerError,
+        "Unexpected issue occured while encoding: " +
+            std::to_string(static_cast<int32_t>(encodeResult.error())));
+  }
+  return encodeResult.get();
 }
 
-std::string TokenizerModule::decode(std::vector<int32_t> vec,
+std::string TokenizerModule::decode(std::vector<uint64_t> vec,
                                     bool skipSpecialTokens) const {
   ensureTokenizerLoaded("decode");
-  return tokenizer->Decode(vec, skipSpecialTokens);
+
+  auto decodeResult = tokenizer->decode(vec, skipSpecialTokens);
+  if (!decodeResult.ok()) {
+    throw RnExecutorchError(
+        RnExecutorchErrorCode::TokenizerError,
+        "Unexpected issue occured while decoding: " +
+            std::to_string(static_cast<int32_t>(decodeResult.error())));
+  }
+
+  return decodeResult.get();
 }
 
 size_t TokenizerModule::getVocabSize() const {
   ensureTokenizerLoaded("getVocabSize");
-  return tokenizer->GetVocabSize();
+  return static_cast<size_t>(tokenizer->vocab_size());
 }
 
-std::string TokenizerModule::idToToken(int32_t tokenId) const {
+std::string TokenizerModule::idToToken(uint64_t tokenId) const {
   ensureTokenizerLoaded("idToToken");
-  return tokenizer->IdToToken(tokenId);
+  auto result = tokenizer->id_to_piece(tokenId);
+  if (!result.ok()) {
+    throw rnexecutorch::RnExecutorchError(
+        rnexecutorch::RnExecutorchErrorCode::TokenizerError,
+        "Unexpected issue occured while trying to convert id to token: " +
+            std::to_string(static_cast<int32_t>(result.error())));
+  }
+  return result.get();
 }
 
-int32_t TokenizerModule::tokenToId(std::string token) const {
+uint64_t TokenizerModule::tokenToId(std::string token) const {
   ensureTokenizerLoaded("tokenToId");
-  return tokenizer->TokenToId(token);
+
+  auto result = tokenizer->piece_to_id(token);
+  if (!result.ok()) {
+    throw rnexecutorch::RnExecutorchError(
+        rnexecutorch::RnExecutorchErrorCode::TokenizerError,
+        "Unexpected issue occured while trying to convert token to id: " +
+            std::to_string(static_cast<int32_t>(result.error())));
+  }
+  return result.get();
 }
+
 std::size_t TokenizerModule::getMemoryLowerBound() const noexcept {
   return memorySizeLowerBound;
 }

packages/react-native-executorch/common/rnexecutorch/TokenizerModule.h

@@ -2,33 +2,33 @@
 
 #include "rnexecutorch/metaprogramming/ConstructorHelpers.h"
 #include <ReactCommon/CallInvoker.h>
+#include <pytorch/tokenizers/hf_tokenizer.h>
 #include <string>
-#include <tokenizers-cpp/tokenizers_cpp.h>
 namespace rnexecutorch {
 using namespace facebook;
 
 class TokenizerModule {
 public:
   explicit TokenizerModule(std::string source,
                            std::shared_ptr<react::CallInvoker> callInvoker);
-  [[nodiscard("Registered non-void function")]] std::vector<int32_t>
+  [[nodiscard("Registered non-void function")]] std::vector<uint64_t>
   encode(std::string s) const;
   [[nodiscard("Registered non-void function")]] std::string
-  decode(std::vector<int32_t> vec, bool skipSpecialTokens) const;
+  decode(std::vector<uint64_t> vec, bool skipSpecialTokens) const;
   [[nodiscard("Registered non-void function")]] std::string
-  idToToken(int32_t tokenId) const;
-  [[nodiscard("Registered non-void function")]] int32_t
+  idToToken(uint64_t tokenId) const;
+  [[nodiscard("Registered non-void function")]] uint64_t
   tokenToId(std::string token) const;
   [[nodiscard("Registered non-void function")]] std::size_t
   getVocabSize() const;
   std::size_t getMemoryLowerBound() const noexcept;
 
 private:
   void ensureTokenizerLoaded(const std::string &methodName) const;
-  std::unique_ptr<tokenizers::Tokenizer> tokenizer;
+  std::unique_ptr<tokenizers::HFTokenizer> tokenizer;
   const std::size_t memorySizeLowerBound{0};
 };
 
 REGISTER_CONSTRUCTOR(TokenizerModule, std::string,
                      std::shared_ptr<react::CallInvoker>);
-} // namespace rnexecutorch
+} // namespace rnexecutorch

packages/react-native-executorch/common/rnexecutorch/host_objects/JsiConversions.h

@@ -218,6 +218,12 @@ getValue<std::vector<int64_t>>(const jsi::Value &val, jsi::Runtime &runtime) {
   return getArrayAsVector<int64_t>(val, runtime);
 }
 
+template <>
+inline std::vector<uint64_t>
+getValue<std::vector<uint64_t>>(const jsi::Value &val, jsi::Runtime &runtime) {
+  return getArrayAsVector<uint64_t>(val, runtime);
+}
+
 // Template specializations for std::span<T> types
 template <>
 inline std::span<float> getValue<std::span<float>>(const jsi::Value &val,
@@ -273,6 +279,12 @@ inline std::span<int64_t> getValue<std::span<int64_t>>(const jsi::Value &val,
   return getTypedArrayAsSpan<int64_t>(val, runtime);
 }
 
+template <>
+inline std::span<uint64_t>
+getValue<std::span<uint64_t>>(const jsi::Value &val, jsi::Runtime &runtime) {
+  return getTypedArrayAsSpan<uint64_t>(val, runtime);
+}
+
 // Conversion from C++ types to jsi --------------------------------------------
 
 // Implementation functions might return any type, but in a promise we can only
@@ -293,6 +305,15 @@ inline jsi::Value getJsiValue(const std::vector<int32_t> &vec,
   return {runtime, array};
 }
 
+inline jsi::Value getJsiValue(const std::vector<uint64_t> &vec,
+                              jsi::Runtime &runtime) {
+  jsi::Array array(runtime, vec.size());
+  for (size_t i = 0; i < vec.size(); i++) {
+    array.setValueAtIndex(runtime, i, jsi::Value(static_cast<double>(vec[i])));
+  }
+  return {runtime, array};
+}
+
 inline jsi::Value getJsiValue(const std::vector<float> &vec,
                               jsi::Runtime &runtime) {
   jsi::Array array(runtime, vec.size());
@@ -302,6 +323,16 @@ inline jsi::Value getJsiValue(const std::vector<float> &vec,
   return {runtime, array};
 }
 
+inline jsi::Value getJsiValue(const std::vector<std::string> &vec,
+                              jsi::Runtime &runtime) {
+  jsi::Array array(runtime, vec.size());
+  for (size_t i = 0; i < vec.size(); i++) {
+    array.setValueAtIndex(runtime, i,
+                          jsi::String::createFromUtf8(runtime, vec[i]));
+  }
+  return {runtime, array};
+}
+
 inline jsi::Value getJsiValue(const std::vector<char> &vec,
                               jsi::Runtime &runtime) {
   jsi::Array array(runtime, vec.size());
@@ -311,10 +342,28 @@ inline jsi::Value getJsiValue(const std::vector<char> &vec,
   return {runtime, array};
 }
 
+// Conditional as on android, size_t and uint64_t reduce to the same type,
+// introducing ambiguity
+template <typename T,
+          typename = std::enable_if_t<std::is_same_v<T, size_t> &&
+                                      !std::is_same_v<size_t, uint64_t>>>
+inline jsi::Value getJsiValue(T val, jsi::Runtime &runtime) {
+  return jsi::Value(static_cast<double>(val));
+}
+
+inline jsi::Value getJsiValue(uint64_t val, jsi::Runtime &runtime) {
+  jsi::BigInt bigInt = jsi::BigInt::fromUint64(runtime, val);
+  return {runtime, bigInt};
+}
+
 inline jsi::Value getJsiValue(int val, jsi::Runtime &runtime) {
   return {runtime, val};
 }
 
+inline jsi::Value getJsiValue(bool val, jsi::Runtime &runtime) {
+  return jsi::Value(val);
+}
+
 inline jsi::Value getJsiValue(const std::shared_ptr<OwningArrayBuffer> &buf,
                               jsi::Runtime &runtime) {
   jsi::ArrayBuffer arrayBuffer(runtime, buf);

packages/react-native-executorch/common/rnexecutorch/models/speech_to_text/SpeechToText.cpp

@@ -37,7 +37,7 @@ SpeechToText::encode(std::span<float> waveform) const {
 }
 
 std::shared_ptr<OwningArrayBuffer>
-SpeechToText::decode(std::span<int32_t> tokens,
+SpeechToText::decode(std::span<uint64_t> tokens,
                      std::span<float> encoderOutput) const {
   std::vector<float> decoderOutput = this->asr->decode(tokens, encoderOutput);
   return std::make_shared<OwningArrayBuffer>(decoderOutput);

packages/react-native-executorch/common/rnexecutorch/models/speech_to_text/SpeechToText.h

@@ -22,7 +22,7 @@ class SpeechToText {
   encode(std::span<float> waveform) const;
   [[nodiscard(
       "Registered non-void function")]] std::shared_ptr<OwningArrayBuffer>
-  decode(std::span<int32_t> tokens, std::span<float> encoderOutput) const;
+  decode(std::span<uint64_t> tokens, std::span<float> encoderOutput) const;
   [[nodiscard("Registered non-void function")]] std::vector<char>
   transcribe(std::span<float> waveform, std::string languageOption) const;