From f0e467c59f29ea3cf4797e0cb8ca459529350c32 Mon Sep 17 00:00:00 2001
From: "The gemma.cpp Authors" <no-reply@google.com>
Date: Thu, 3 Sep 2026 10:56:38 -0700
Subject: [PATCH] Internal Changes

PiperOrigin-RevId: 975815177
---
 deepseek/dsv4_tokenizer.cc | 32 +++++++++++++++++++++++++++-----
 deepseek/dsv4_tokenizer.h  |  5 +++--
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/deepseek/dsv4_tokenizer.cc b/deepseek/dsv4_tokenizer.cc
index e339d809..eff36463 100644
--- a/deepseek/dsv4_tokenizer.cc
+++ b/deepseek/dsv4_tokenizer.cc
@@ -18,6 +18,7 @@
 #include <stdio.h>
 
 #include <algorithm>
+#include <cstring>
 #include <fstream>
 #include <limits>
 #include <sstream>
@@ -25,7 +26,6 @@
 #include <string_view>
 #include <vector>
 
-
 #include "hwy/base.h"  // HWY_ABORT
 #include "nlohmann/json.hpp"
 
@@ -177,13 +177,14 @@ Dsv4Tokenizer::Dsv4Tokenizer(const std::string& tokenizer_json_path) {
   Init(contents);
 }
 
-Dsv4Tokenizer::Dsv4Tokenizer(std::string_view json_content, bool /*is_content*/) {
+Dsv4Tokenizer::Dsv4Tokenizer(std::string_view json_content,
+                             bool /*is_content*/) {
   Init(json_content);
 }
 
 void Dsv4Tokenizer::Init(std::string_view json_content) {
   json j = json::parse(json_content.begin(), json_content.end(), /*cb=*/nullptr,
-                      /*allow_exceptions=*/false);
+                       /*allow_exceptions=*/false);
   if (j.is_discarded()) {
     HWY_ABORT("Failed to parse tokenizer JSON");
   }
@@ -274,11 +275,15 @@ std::string Dsv4Tokenizer::WrapChat(const std::string& user_msg,
   static const char kAssistant[] =
       "<\xEF\xBD\x9C"
       "Assistant\xEF\xBD\x9C>";
-  std::string out(kBos);
+  const char* think_marker = thinking ? "<think>" : "</think>";
+  std::string out;
+  out.reserve((sizeof(kBos) - 1) + (sizeof(kUser) - 1) + user_msg.size() +
+              (sizeof(kAssistant) - 1) + std::strlen(think_marker));
+  out += kBos;
   out += kUser;
   out += user_msg;
   out += kAssistant;
-  out += thinking ? "<think>" : "</think>";
+  out += think_marker;
   return out;
 }
 
@@ -494,4 +499,21 @@ void Dsv4Tokenizer::AppendDecoded(int id, std::string& out) const {
   out += id_to_bytes_[id];
 }
 
+std::string Dsv4Tokenizer::Decode(const std::vector<int>& ids) const {
+  size_t total_size = 0;
+  for (int id : ids) {
+    if (id >= 0 && static_cast<size_t>(id) < id_to_bytes_.size() &&
+        !is_special_[id]) {
+      total_size += id_to_bytes_[id].size();
+    }
+  }
+
+  std::string out;
+  out.reserve(total_size);
+  for (int id : ids) {
+    AppendDecoded(id, out);
+  }
+  return out;
+}
+
 }  // namespace gcpp
diff --git a/deepseek/dsv4_tokenizer.h b/deepseek/dsv4_tokenizer.h
index 705e7b5e..2a08ed02 100644
--- a/deepseek/dsv4_tokenizer.h
+++ b/deepseek/dsv4_tokenizer.h
@@ -40,7 +40,6 @@ class Dsv4Tokenizer {
   // 'is_content' is used to distinguish from path constructor.
   Dsv4Tokenizer(std::string_view json_content, bool is_content);
 
-
   // Extracts added tokens (chat markers etc.), pre-tokenizes and BPE-encodes
   // everything in between. Equivalent to HF encode(add_special_tokens=false).
   std::vector<int> Encode(const std::string& text) const;
@@ -55,6 +54,9 @@ class Dsv4Tokenizer {
   // skip_special_tokens=true; other added tokens (e.g. </think>) are kept.
   void AppendDecoded(int id, std::string& out) const;
 
+  // Decodes a sequence of token IDs, skipping special tokens.
+  std::string Decode(const std::vector<int>& ids) const;
+
   size_t VocabSize() const { return id_to_bytes_.size(); }
 
  private:
@@ -65,7 +67,6 @@ class Dsv4Tokenizer {
 
   void Init(std::string_view json_content);
 
-
   // Splits `text` (a span with no added tokens) into pre-tokenization pieces
   // and BPE-encodes each, appending ids.
   void EncodeSegment(const char* bytes, size_t len,