Add optional Hugging Face .tok tokenizer backend

CMakeLists.txt

@@ -19,6 +19,11 @@ project(Tokenizers)
 option(TOKENIZERS_BUILD_TEST "Build tests" OFF)
 option(TOKENIZERS_BUILD_TOOLS "Build tools" OFF)
 option(TOKENIZERS_BUILD_PYTHON "Build Python bindings" OFF)
+option(
+  TOKENIZERS_BUILD_HF_RUST_TOKENIZER
+  "Build the opt-in Hugging Face .tok tokenizer backend" OFF
+)
+option(TOKENIZERS_OPTIMIZE_SIZE "Optimize optional tokenizer backends for size" OFF)
 option(SUPPORT_REGEX_LOOKAHEAD
        "Support regex lookahead patterns (requires PCRE2)" OFF
 )
@@ -201,6 +206,10 @@ endif()
 # Installation rules
 include(GNUInstallDirs)
 
+if(TOKENIZERS_BUILD_HF_RUST_TOKENIZER)
+  add_subdirectory(rust_tokenizer)
+endif()
+
 if(NOT TOKENIZERS_BUILD_PYTHON)
   # Install the library and its dependencies
   install(

cmake/tokenizers-config.cmake.in

@@ -25,6 +25,25 @@ endif()
 find_dependency(re2 REQUIRED)
 find_dependency(absl REQUIRED)
 
+# The optional .tok backend is a C++ archive with a Rust static-library
+# dependency. Define the latter before importing the exported targets when the
+# archive is present. Merely finding the package does not link either target.
+if(WIN32)
+  set(_TOKENIZERS_HF_FFI_NAME "tokenizers_hf_ffi.lib")
+else()
+  set(_TOKENIZERS_HF_FFI_NAME "libtokenizers_hf_ffi.a")
+endif()
+set(_TOKENIZERS_HF_FFI "${TOKENIZERS_LIBDIR}/${_TOKENIZERS_HF_FFI_NAME}")
+if(EXISTS "${_TOKENIZERS_HF_FFI}" AND NOT TARGET tokenizers_hf_ffi)
+  if(UNIX AND NOT APPLE)
+    find_dependency(Threads)
+  endif()
+  add_library(tokenizers_hf_ffi STATIC IMPORTED)
+  set_target_properties(
+    tokenizers_hf_ffi PROPERTIES IMPORTED_LOCATION "${_TOKENIZERS_HF_FFI}"
+  )
+endif()
+
 # Include the exported targets file
 include("${CMAKE_CURRENT_LIST_DIR}/tokenizers-targets.cmake")
 

include/pytorch/tokenizers/rust_hf_tokenizer.h

@@ -0,0 +1,56 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include <cstdint>
+#include <memory>
+#include <optional>
+#include <string>
+#include <unordered_set>
+#include <vector>
+
+#include <pytorch/tokenizers/string_integer_map.h>
+#include <pytorch/tokenizers/tokenizer.h>
+
+namespace tokenizers {
+
+class RustHFTokenizer final : public Tokenizer {
+ public:
+  RustHFTokenizer();
+  ~RustHFTokenizer() override;
+
+  Error load(const std::string& tokenizer_path) override;
+  Result<std::string> id_to_piece(uint64_t token) const override;
+  Result<uint64_t> piece_to_id(const std::string& text) const override;
+  Result<std::vector<uint64_t>> encode(
+      const std::string& input,
+      int8_t bos = 0,
+      int8_t eos = 0) const override;
+  Result<std::string> decode(
+      uint64_t prev_token,
+      uint64_t token,
+      bool skip_special_tokens = false) const override;
+
+ private:
+  using TokenMap = detail::StringIntegerMap<>;
+
+  struct RustHandleDeleter {
+    void operator()(void* handle) const;
+  };
+
+  Error load_metadata(const void* handle);
+
+  std::unique_ptr<void, RustHandleDeleter> handle_;
+  std::optional<TokenMap> token_map_;
+  std::optional<TokenMap> added_token_map_;
+  std::unordered_set<uint64_t> special_token_ids_;
+  bool byte_level_ = false;
+};
+
+} // namespace tokenizers

rust_tokenizer/.gitignore

@@ -0,0 +1 @@
+/target/

rust_tokenizer/CMakeLists.txt

@@ -0,0 +1,79 @@
+# 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.
+
+find_program(CARGO_EXECUTABLE cargo REQUIRED)
+
+set(_rust_profile release)
+if(TOKENIZERS_OPTIMIZE_SIZE)
+  set(_rust_profile minsize)
+endif()
+
+set(_cargo_target_dir ${CMAKE_CURRENT_BINARY_DIR}/cargo-target)
+string(
+  CONCAT
+  _rust_library_name
+  ${CMAKE_STATIC_LIBRARY_PREFIX}
+  tokenizers_hf_ffi
+  ${CMAKE_STATIC_LIBRARY_SUFFIX}
+)
+set(_rust_library ${_cargo_target_dir}/${_rust_profile}/${_rust_library_name})
+file(GLOB_RECURSE _rust_sources CONFIGURE_DEPENDS
+     ${CMAKE_CURRENT_SOURCE_DIR}/src/*.rs
+)
+
+add_custom_command(
+  OUTPUT ${_rust_library}
+  COMMAND
+    ${CMAKE_COMMAND} -E env CARGO_TARGET_DIR=${_cargo_target_dir}
+    CARGO_ENCODED_RUSTFLAGS=-Crelocation-model=pic
+    ${CARGO_EXECUTABLE} build --locked --manifest-path
+    ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml --profile ${_rust_profile}
+  DEPENDS ${_rust_sources} ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.toml
+          ${CMAKE_CURRENT_SOURCE_DIR}/Cargo.lock
+  WORKING_DIRECTORY ${CMAKE_CURRENT_SOURCE_DIR}
+  VERBATIM
+)
+
+add_custom_target(tokenizers_hf_ffi_build DEPENDS ${_rust_library})
+add_library(tokenizers_hf_ffi STATIC IMPORTED GLOBAL)
+set_target_properties(
+  tokenizers_hf_ffi PROPERTIES IMPORTED_LOCATION ${_rust_library}
+)
+add_dependencies(tokenizers_hf_ffi tokenizers_hf_ffi_build)
+
+add_library(tokenizers_hf_rust_tokenizer STATIC ../src/rust_hf_tokenizer.cpp)
+add_library(tokenizers::hf_rust_tokenizer ALIAS tokenizers_hf_rust_tokenizer)
+set_target_properties(
+  tokenizers_hf_rust_tokenizer PROPERTIES EXPORT_NAME hf_rust_tokenizer
+)
+target_include_directories(
+  tokenizers_hf_rust_tokenizer
+  PUBLIC
+    $<BUILD_INTERFACE:${CMAKE_CURRENT_SOURCE_DIR}/../include>
+    $<INSTALL_INTERFACE:${CMAKE_INSTALL_INCLUDEDIR}>
+)
+target_compile_features(tokenizers_hf_rust_tokenizer PUBLIC cxx_std_17)
+target_link_libraries(
+  tokenizers_hf_rust_tokenizer
+  PRIVATE tokenizers_hf_ffi
+)
+
+if(APPLE)
+  target_link_libraries(tokenizers_hf_rust_tokenizer PRIVATE iconv)
+elseif(UNIX)
+  find_package(Threads REQUIRED)
+  target_link_libraries(
+    tokenizers_hf_rust_tokenizer
+    PRIVATE Threads::Threads ${CMAKE_DL_LIBS} m rt util
+  )
+endif()
+
+install(
+  TARGETS tokenizers_hf_rust_tokenizer
+  EXPORT tokenizers-targets
+  ARCHIVE DESTINATION ${CMAKE_INSTALL_LIBDIR}
+)
+install(FILES ${_rust_library} DESTINATION ${CMAKE_INSTALL_LIBDIR})