Native tokenizer

BUILD.bazel

@@ -220,6 +220,30 @@ cc_test(
     ],
 )
 
+genrule(
+    name = "packed_tokenizer",
+    srcs = ["tokenizer/testdata/tokenizer.json"],
+    outs = ["tokenizer/testdata/tokenizer.packed"],
+    cmd = "$(location //python:convert_from_safetensors) --pack_tokenizer --tokenizer_file $(location tokenizer/testdata/tokenizer.json) --output_packed $@",
+    tools = ["//python:convert_from_safetensors"],
+)
+
+cc_test(
+    name = "bpe_tokenizer_test",
+    srcs = ["tokenizer/bpe_tokenizer_test.cc"],
+    data = [
+        "testdata/frankenstein.txt",
+        "tokenizer/testdata/tokenizer.json",
+        "tokenizer/testdata/tokenizer.model",
+        ":packed_tokenizer",
+    ],
+    deps = [
+        ":tokenizer",
+        "//testing/base/public:gunit_main",  # buildcleaner: keep
+        "//io",
+    ],
+)
+
 cc_library(
     name = "tensor_info",
     srcs = [
@@ -252,10 +276,21 @@ cc_library(
 
 cc_library(
     name = "tokenizer",
-    srcs = ["gemma/tokenizer.cc"],
-    hdrs = ["gemma/tokenizer.h"],
+    srcs = [
+        "gemma/tokenizer.cc",
+        "tokenizer/bpe_tokenizer.cc",
+        "tokenizer/sentencepiece_tokenizer.cc",
+    ],
+    hdrs = [
+        "gemma/tokenizer.h",
+        "tokenizer/bpe_tokenizer.h",
+        "tokenizer/sentencepiece_tokenizer.h",
+        "tokenizer/tokenizer.h",
+    ],
     deps = [
+        ":basics",
         ":configs",
+        "//io:fields",
         "@highway//:hwy",
         "@highway//:profiler",
         "@com_google_sentencepiece//:sentencepiece_processor",

CMakeLists.txt

@@ -192,6 +192,11 @@ set(SOURCES
   paligemma/image.h
   paligemma/paligemma_helper.cc
   paligemma/paligemma_helper.h
+  tokenizer/bpe_tokenizer.cc
+  tokenizer/bpe_tokenizer.h
+  tokenizer/sentencepiece_tokenizer.cc
+  tokenizer/sentencepiece_tokenizer.h
+  tokenizer/tokenizer.h
   util/allocator.cc
   util/allocator.h
   util/basics.cc

compression/python/compression_clif_aux.cc

@@ -45,6 +45,7 @@
 #include "hwy/highway.h"
 // After highway.h
 #include "compression/compress-inl.h"
+#include "tokenizer/bpe_tokenizer.h"
 
 // SIMD code, compiled once per target.
 HWY_BEFORE_NAMESPACE();
@@ -118,11 +119,18 @@ class SbsWriterImpl : public ISbsWriter {
   }
 
   void Write(const ModelConfig& config,
-             const std::string& tokenizer_path) override {
-    const GemmaTokenizer tokenizer(
-        tokenizer_path.empty() ? kMockTokenizer
-                               : ReadFileToString(Path(tokenizer_path)));
-    WriteSingleFile(config, tokenizer, serialized_mat_ptrs_, writer_);
+             const std::string& tokenizer_blob) override {
+    WriteSingleFile(config, BuildTokenizer(config, tokenizer_blob),
+                    serialized_mat_ptrs_, writer_);
+  }
+
+  static GemmaTokenizer BuildTokenizer(const ModelConfig& config,
+                                       const std::string& tokenizer_blob) {
+    if (tokenizer_blob.empty()) return GemmaTokenizer(kMockTokenizer);
+    if (config.tokenizer_kind == TokenizerKind::kHfBpe) {
+      return GemmaTokenizer(CreateBpeTokenizer(tokenizer_blob));
+    }
+    return GemmaTokenizer(tokenizer_blob);
   }
 
   ThreadingContext ctx_;

compression/python/compression_clif_aux.h

@@ -44,7 +44,7 @@ class ISbsWriter {
                       const TensorInfo& tensor_info) = 0;
 
   virtual void Write(const ModelConfig& config,
-                     const std::string& tokenizer_path) = 0;
+                     const std::string& tokenizer_blob) = 0;
 };
 
 // Non-virtual class used by pybind that calls the interface's virtual methods.
@@ -58,8 +58,8 @@ class SbsWriter {
     impl_->Insert(name, weights, type, tensor_info);
   }
 
-  void Write(const ModelConfig& config, const std::string& tokenizer_path) {
-    impl_->Write(config, tokenizer_path);
+  void Write(const ModelConfig& config, const std::string& tokenizer_blob) {
+    impl_->Write(config, tokenizer_blob);
   }
 
  private:

compression/python/compression_extension.cc

@@ -46,7 +46,7 @@ PYBIND11_MODULE(compression, m) {
   class_<SbsWriter>(m, "SbsWriter")
       .def(init<std::string>())
       .def("insert", CallWithF32Span<&SbsWriter::Insert>)
-      .def("write", &SbsWriter::Write, arg("config"), arg("tokenizer_path"));
+      .def("write", &SbsWriter::Write, arg("config"), arg("tokenizer_blob"));
 
   class_<MatPtr>(m, "MatPtr")
       // No init, only created within C++.

compression/python/compression_test.py

@@ -102,8 +102,8 @@ def test_sbs_writer(self):
         configs.Type.kSFP,
         configs.PromptWrapping.GEMMA_IT,
     )
-    tokenizer_path = ""  # no tokenizer required for testing
-    writer.write(config, tokenizer_path)
+    tokenizer_blob = ""  # no tokenizer required for testing
+    writer.write(config, tokenizer_blob)
 
     print("Ignore next two warnings; test does not enable model deduction.")
     reader = compression.SbsReader(temp_file.full_path)

gemma/configs.h

@@ -52,6 +52,18 @@ enum class PromptWrapping {
   kSentinel  // must be last
 };
 
+// Which tokenizer engine the `.sbs` tokenizer blob targets.
+enum class TokenizerKind {
+  kSentencePiece = 0,  // blob is a SentencePiece proto
+  kHfBpe = 1,  // blob is a compact BPE format built from HF `tokenizer.json`
+  kSentinel    // must be last
+};
+
+static inline bool EnumValid(TokenizerKind kind) {
+  return static_cast<size_t>(kind) <
+         static_cast<size_t>(TokenizerKind::kSentinel);
+}
+
 // This is used in `ModelConfig.Specifier`, so the strings will not change,
 // though new ones may be added.
 static inline const char* WrappingSuffix(PromptWrapping wrapping) {
@@ -605,6 +617,8 @@ struct ModelConfig : public IFields {
     visitor(hc_eps);
     visitor(num_mtp_layers);
 
+    visitor(tokenizer_kind);
+
     // Append new fields here, then update `python/configs.cc`.
   }
 
@@ -760,6 +774,8 @@ struct ModelConfig : public IFields {
   // is an extra transformer layer used for speculative decoding; it is not
   // counted in `num_layers` and never runs as part of the main stack.
   uint32_t num_mtp_layers = 0;
+
+  TokenizerKind tokenizer_kind = TokenizerKind::kSentencePiece;
 };
 
 // Returns the sub-config for the ViT model of the PaliGemma model.

gemma/model_store.cc

@@ -34,6 +34,7 @@
 #include "io/blob_store.h"
 #include "io/fields.h"
 #include "io/io.h"  // Path
+#include "tokenizer/bpe_tokenizer.h"
 #include "util/basics.h"
 #include "util/threading_context.h"
 #include "hwy/base.h"
@@ -58,12 +59,8 @@ static void WarnIfExtra(const IFields::ReadResult& result, const char* name) {
   }
 }
 
-// Returns the serialized tokenizer (std::string is required for proto).
-// Reads it from a blob or from a separate file if pre-2025.
-static std::string ReadTokenizer(BlobReader& reader,
-                                 const Path& tokenizer_path) {
-  PROFILER_ZONE("Startup.ReadTokenizer");
-
+// Reads the raw `tokenizer` blob bytes, or empty if absent.
+static std::string ReadTokenizerBlob(BlobReader& reader) {
   std::string tokenizer;
   // Check prevents `CallWithSpan` from printing a warning.
   if (reader.Find(kTokenizerName)) {
@@ -76,27 +73,51 @@ static std::string ReadTokenizer(BlobReader& reader,
           "instead specify a tokenizer file via --tokenizer.");
     }
   }
+  return tokenizer;
+}
+
+// Constructs the tokenizer.
+static GemmaTokenizer ReadTokenizer(BlobReader& reader,
+                                    const Path& tokenizer_path,
+                                    const ModelConfig& config) {
+  PROFILER_ZONE("Startup.ReadTokenizer");
+
+  if (config.tokenizer_kind == TokenizerKind::kHfBpe) {
+    const std::string blob = ReadTokenizerBlob(reader);
+    if (!blob.empty() && blob != kMockTokenizer) {
+      return GemmaTokenizer(CreateBpeTokenizer(blob));
+    }
+    if (!tokenizer_path.Empty()) {
+      HWY_WARN("No separate tokenizer file is supported for HF BPE models. "
+               "Tests may continue but inference will fail. %s.",
+               tokenizer_path.path.c_str());
+      return GemmaTokenizer(kMockTokenizer);
+    }
+    HWY_WARN("BlobStore has no HuggingFace BPE tokenizer specified. Tests may "
+             "continue but inference will fail.");
+    return GemmaTokenizer(kMockTokenizer);
+  }
 
-  // Read actual tokenizer from blob.
-  if (!tokenizer.empty() && tokenizer != kMockTokenizer) {
+  // SentencePiece: the blob is a SentencePiece proto.
+  const std::string blob = ReadTokenizerBlob(reader);
+  if (!blob.empty() && blob != kMockTokenizer) {
     if (!tokenizer_path.Empty()) {
       HWY_WARN("--weights has tokenizer but overriding with %s.",
                tokenizer_path.path.c_str());
-      return ReadFileToString(tokenizer_path);
+      return GemmaTokenizer(ReadFileToString(tokenizer_path));
     }
-
-    return tokenizer;
+    return GemmaTokenizer(blob);
   }
 
   // No blob but user specified path to file: read it or abort.
   if (!tokenizer_path.Empty()) {
-    return ReadFileToString(tokenizer_path);
+    return GemmaTokenizer(ReadFileToString(tokenizer_path));
   }
 
   HWY_WARN(
       "BlobStore does not contain a tokenizer and no --tokenizer was "
       "specified. Tests may continue but inference will fail.");
-  return kMockTokenizer;
+  return GemmaTokenizer(kMockTokenizer);
 }
 
 using KeyVec = std::vector<std::string>;
@@ -393,7 +414,7 @@ void ModelStore::CreateMatPtrs(BlobReader& reader) {
 ModelStore::ModelStore(BlobReader& reader, const Path& tokenizer_path,
                        Tristate wrapping)
     : config_(ReadOrDeduceConfig(reader, wrapping)),
-      tokenizer_(ReadTokenizer(reader, tokenizer_path)) {
+      tokenizer_(ReadTokenizer(reader, tokenizer_path, config_)) {
   if (!ReadMatPtrs(reader)) {  // Pre-2025 format.
     CreateMatPtrs(reader);
     scales_ = ReadScales(reader, config_);

gemma/tokenizer.cc

@@ -19,89 +19,49 @@
 
 #include <memory>
 #include <string>
+#include <utility>
 #include <vector>
 
 #include "gemma/configs.h"  // PromptWrapping
-#include "hwy/base.h"         // HWY_ASSERT
-#include "hwy/profiler.h"
-// copybara:import_next_line:sentencepiece
-#include "src/sentencepiece_processor.h"
+#include "tokenizer/sentencepiece_tokenizer.h"
 
 namespace gcpp {
 
 // Set this to true to debug tokenizer tokens.
 constexpr bool kShowTokenization = false;
 
-class GemmaTokenizer::Impl {
- public:
-  Impl() = default;
-  // Loads the tokenizer from a serialized proto.
-  explicit Impl(const std::string& tokenizer_proto) {
-    if (tokenizer_proto == kMockTokenizer) return;
-    PROFILER_ZONE("Startup.tokenizer");
-    spp_ = std::make_unique<sentencepiece::SentencePieceProcessor>();
-    if (!spp_->LoadFromSerializedProto(tokenizer_proto).ok()) {
-      HWY_ABORT("Failed to load tokenizer from %zu byte serialized proto.",
-                tokenizer_proto.size());
-    }
-  }
-
-  std::string Serialize() const {
-    return spp_ ? spp_->serialized_model_proto() : kMockTokenizer;
-  }
-
-  bool Encode(const std::string& input,
-              std::vector<std::string>* pieces) const {
-    return spp_ && spp_->Encode(input, pieces).ok();
-  }
-
-  bool Encode(const std::string& input, std::vector<int>* ids) const {
-    if constexpr (kShowTokenization) {
-      bool is_ok = spp_ && spp_->Encode(input, ids).ok();
-      for (int i = 0; i < static_cast<int>(ids->size()); i++) {
-        fprintf(stderr, "%3d: %d\n", i, (*ids)[i]);
-      }
-      return is_ok;
-    } else {
-      return spp_ && spp_->Encode(input, ids).ok();
-    }
-  }
-
-  // Given a sequence of ids, decodes it into a detokenized output.
-  bool Decode(const std::vector<int>& ids, std::string* detokenized) const {
-    return spp_ && spp_->Decode(ids, detokenized).ok();
-  }
-
- private:
-  std::unique_ptr<sentencepiece::SentencePieceProcessor> spp_;
-};
-
-GemmaTokenizer::GemmaTokenizer(const std::string& tokenizer_proto)
-    : impl_(std::make_unique<Impl>(tokenizer_proto)) {
-  HWY_ASSERT(impl_);
+GemmaTokenizer::GemmaTokenizer(const std::string& tokenizer_proto) {
+  if (tokenizer_proto == kMockTokenizer) return;
+  impl_ = CreateSentencePieceTokenizer(tokenizer_proto);
 }
 
+GemmaTokenizer::GemmaTokenizer(std::unique_ptr<Tokenizer> impl)
+    : impl_(std::move(impl)) {}
+
 // Default suffices, but they must be defined after GemmaTokenizer::Impl.
 GemmaTokenizer::~GemmaTokenizer() = default;
 GemmaTokenizer::GemmaTokenizer(GemmaTokenizer&& other) = default;
 GemmaTokenizer& GemmaTokenizer::operator=(GemmaTokenizer&& other) = default;
 
-std::string GemmaTokenizer::Serialize() const { return impl_->Serialize(); }
-
-bool GemmaTokenizer::Encode(const std::string& input,
-                            std::vector<std::string>* pieces) const {
-  return impl_->Encode(input, pieces);
+std::string GemmaTokenizer::Serialize() const {
+  return impl_ ? impl_->Serialize() : kMockTokenizer;
 }
 
 bool GemmaTokenizer::Encode(const std::string& input,
                             std::vector<int>* ids) const {
-  return impl_->Encode(input, ids);
+  if (!impl_ || !impl_->Encode(input, *ids)) return false;
+  if constexpr (kShowTokenization) {
+    for (int i = 0; i < static_cast<int>(ids->size()); i++) {
+      fprintf(stderr, "%3d: %d\n", i, (*ids)[i]);
+    }
+  }
+  return true;
 }
 
 // Given a sequence of ids, decodes it into a detokenized output.
 bool GemmaTokenizer::Decode(const std::vector<int>& ids,
                             std::string* detokenized) const {
-  return impl_->Decode(ids, detokenized);
+  return impl_ && impl_->Decode(ids, *detokenized);
 }
 
 // Negligible CPU time in the ctor body.