bugfix: setup the tokenizer config function of bos and eos to fast tokenizer. (#367)

Co-authored-by: phantomlei <phantomlei3@gmail.com>

Author: 2 people

xllm/core/framework/hf_model_loader.cpp

@@ -259,6 +259,15 @@ bool HFModelLoader::load_tokenizer_args(const std::string& model_weights_path) {
   JsonReader tokenizer_reader;
   const std::string tokenizer_args_file_path =
       model_weights_path_ + "/tokenizer_config.json";
+
+  // check if tokenizer.json exists, if exists, set the tokenizer type to fast
+  const std::string tokenizer_json_path =
+      model_weights_path + "/tokenizer.json";
+  if (std::filesystem::exists(tokenizer_json_path)) {
+    tokenizer_args_.tokenizer_type() = "fast";
+    tokenizer_args_.vocab_file() = tokenizer_json_path;
+  }
+
   if (tokenizer_reader.parse(tokenizer_args_file_path)) {
     // read chat template if exists
     if (auto v = load_chat_template_file(model_weights_path_)) {

xllm/core/framework/tokenizer/CMakeLists.txt

@@ -30,3 +30,14 @@ cc_library(
     re2::re2
 )
 
+cc_test(
+  NAME
+    fast_tokenizer_test
+  SRCS
+    tests/fast_tokenizer_tests.cpp
+  DEPS
+    :tokenizer
+    glog::glog
+    GTest::gtest_main
+)
+

xllm/core/framework/tokenizer/fast_tokenizer.cpp

@@ -20,19 +20,55 @@ limitations under the License.
 
 namespace xllm {
 
-FastTokenizer::FastTokenizer(const std::string& tokenizer_json_path)
-    : tokenizer_json_path_(tokenizer_json_path) {
-  handle_ = tokenizers_new_from_path(tokenizer_json_path.c_str());
+FastTokenizer::FastTokenizer(const TokenizerArgs& tokenizer_args)
+    : tokenizer_args_(tokenizer_args) {
+  handle_ = tokenizers_new_from_path(tokenizer_args.vocab_file().c_str());
   CHECK(handle_ != nullptr)
-      << "Failed to load tokenizer from file: " << tokenizer_json_path;
+      << "Failed to load tokenizer from file: " << tokenizer_args.vocab_file();
 }
 
 std::unique_ptr<Tokenizer> FastTokenizer::clone() const {
-  return std::make_unique<FastTokenizer>(tokenizer_json_path_);
+  return std::make_unique<FastTokenizer>(tokenizer_args_);
 }
 
 FastTokenizer::~FastTokenizer() { tokenizers_free(handle_); }
 
+namespace {
+// Helper function to add a special token to the beginning or end of ids
+// Checks if token already exists before adding to avoid duplication
+// Returns true on success, false if token is not found, empty, or already
+// exists
+bool add_special_token_id(const std::string& token,
+                          std::optional<int32_t> token_id,
+                          std::vector<int32_t>* ids,
+                          bool prepend) {
+  if (token.empty() || !token_id.has_value()) {
+    if (!token.empty() && !token_id.has_value()) {
+      LOG(WARNING) << "Failed to find token ID for token: " << token;
+    }
+    return false;
+  }
+
+  const int32_t id = token_id.value();
+
+  // Check if token already exists at the expected position
+  if (prepend) {
+    // For BOS: check if already at the beginning
+    if (!ids->empty() && ids->front() == id) {
+      return false;  // Already exists, skip adding
+    }
+    ids->insert(ids->begin(), id);
+  } else {
+    // For EOS: check if already at the end
+    if (!ids->empty() && ids->back() == id) {
+      return false;  // Already exists, skip adding
+    }
+    ids->push_back(id);
+  }
+  return true;
+}
+}  // namespace
+
 bool FastTokenizer::encode(const std::string_view& text,
                            std::vector<int32_t>* ids,
                            bool add_special_tokens) const {
@@ -43,6 +79,31 @@ bool FastTokenizer::encode(const std::string_view& text,
   std::vector<int32_t> ret(result.token_ids, result.token_ids + result.len);
   *ids = std::move(ret);
 
+  // Free the memory allocated by Rust tokenizer
+  // The token_ids pointer is allocated by Rust's Box::into_raw and must be
+  // freed
+  if (result.token_ids != nullptr && result.len > 0) {
+    tokenizers_free_encode_results(&result, 1);
+  }
+
+  // Add BOS token if configured
+  if (tokenizer_args_.add_bos_token() && !tokenizer_args_.bos_token().empty()) {
+    const auto bos_id = token_to_id(tokenizer_args_.bos_token());
+    add_special_token_id(tokenizer_args_.bos_token(),
+                         bos_id,
+                         ids,
+                         /*prepend=*/true);
+  }
+
+  // Add EOS token if configured
+  if (tokenizer_args_.add_eos_token() && !tokenizer_args_.eos_token().empty()) {
+    const auto eos_id = token_to_id(tokenizer_args_.eos_token());
+    add_special_token_id(tokenizer_args_.eos_token(),
+                         eos_id,
+                         ids,
+                         /*prepend=*/false);
+  }
+
   return true;
 }
 

xllm/core/framework/tokenizer/fast_tokenizer.h

@@ -17,13 +17,14 @@ limitations under the License.
 #pragma once
 
 #include "tokenizer.h"
+#include "tokenizer_args.h"
 #include "tokenizers/tokenizers.h"
 
 namespace xllm {
 
 class FastTokenizer : public Tokenizer {
  public:
-  FastTokenizer(const std::string& tokenizer_json_path);
+  FastTokenizer(const TokenizerArgs& tokenizer_args);
 
   ~FastTokenizer() override;
 
@@ -44,8 +45,7 @@ class FastTokenizer : public Tokenizer {
   std::unique_ptr<Tokenizer> clone() const override;
 
  private:
-  std::string tokenizer_json_path_;
-
+  TokenizerArgs tokenizer_args_;
   TokenizerHandle handle_ = nullptr;
 };