diff --git a/gemma/gemma.cc b/gemma/gemma.cc
index 7aaf912..99c4d53 100644
--- a/gemma/gemma.cc
+++ b/gemma/gemma.cc
@@ -548,6 +548,13 @@ struct GemmaImpl : public GemmaInterface {
   hwy::AlignedUniquePtr<Activations<Config, 1>> state;
 };
 
+template <class TConfig>
+std::string TokenString(GemmaImpl<TConfig>& gemma, int token) {
+  std::string token_str;
+  gemma.Tokenizer()->Decode({token}, &token_str);
+  return "'" + std::regex_replace(token_str, std::regex("\n"), "\\n") + "'";
+}
+
 }  // namespace gcpp
 #endif  // GEMMA_ONCE
 
@@ -1121,13 +1128,6 @@ void GenerateImpl(GemmaImpl<TConfig>& gemma, size_t max_tokens,
   }
 }
 
-template <class TConfig>
-std::string TokenString(GemmaImpl<TConfig>& gemma, int token) {
-  std::string token_str;
-  gemma.Tokenizer()->Decode({token}, &token_str);
-  return "'" + std::regex_replace(token_str, std::regex("\n"), "\\n") + "'";
-}
-
 #define TOKEN(token_id) TokenString(gemma, token_id).c_str()
 
 template <class TConfig>