feat(tokenizers): Add decoding support to HFTokenizer

pytorch#1251 Branch: TokenizersCpp-1251 Signed-off-by: Gabe Goodhart <ghart@us.ibm.com>
gabe-l-hart · Nov 15, 2024 · b8c7941 · b8c7941
1 parent 7084831
commit b8c7941
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 7 deletions.
diff --git a/tokenizer/hf_tokenizers.cpp b/tokenizer/hf_tokenizers.cpp
@@ -114,6 +114,13 @@ void HFTokenizer::load(const std::string& path) {
     exit(EXIT_FAILURE);
   }
 
+  // Set up the decoder (optional)
+  try {
+    _decoder = TokenDecoderConfig().parse_json(parsed_json.at("decoder")).create();
+  } catch (const json::out_of_range& e) {
+    // No decoder specified
+  }
+
   // TODO: Do we need to parse the merges?
 
   // If a tokenizer config file is found, parse it to look up the eos/bos tokens
@@ -241,3 +248,14 @@ void HFTokenizer::_encode(
     ret.insert(ret.end(), tokens.begin(), tokens.end());
   }
 }
+
+void HFTokenizer::_decode(
+  re2::StringPiece input,
+  std::string& ret
+) const {
+  if (_decoder) {
+    ret += _decoder->decode(input);
+  } else {
+    ret += input;
+  }
+}
diff --git a/tokenizer/tiktoken.cpp b/tokenizer/tiktoken.cpp
@@ -329,6 +329,12 @@ void Tiktoken::_encode(
   }
 }
 
+void Tiktoken::_decode(
+  re2::StringPiece input,
+  std::string& ret) const {
+  ret += input;
+}
+
 // -------------------------private method end-------------------------------
 // -------------------------public method start-------------------------------
 
@@ -386,8 +392,9 @@ std::string BPETokenizerBase::decode(uint64_t prev, uint64_t cur) const {
       exit(EXIT_FAILURE);
     }
   }
-  ret += token_bytes;
+  _decode(token_bytes, ret);
 
   return ret;
 }
+
 // -------------------------public method end-------------------------------
diff --git a/tokenizer/tokenizer.h b/tokenizer/tokenizer.h
@@ -25,6 +25,7 @@
 
 #include "pre_tokenizer.h"
 #include "sentencepiece_processor.h"
+#include "token_decoder.h"
 
 class Tokenizer {
  public:
@@ -125,10 +126,15 @@ class BPETokenizerBase : public Tokenizer {
   Decoder special_token_decoder_;
 
  private:
+
   virtual void _encode(
-      re2::StringPiece& input,
-      std::vector<uint64_t>& ret,
-      uint64_t& last_piece_token_len) const = 0;
+    re2::StringPiece& input,
+    std::vector<uint64_t>& ret,
+    uint64_t& last_piece_token_len) const = 0;
+
+  virtual void _decode(
+    re2::StringPiece input,
+    std::string& ret) const = 0;
 };
 
 class Tiktoken : public BPETokenizerBase {
@@ -160,9 +166,13 @@ class Tiktoken : public BPETokenizerBase {
   }
 
   void _encode(
-      re2::StringPiece& input,
-      std::vector<uint64_t>& ret,
-      uint64_t& last_piece_token_len) const override;
+    re2::StringPiece& input,
+    std::vector<uint64_t>& ret,
+    uint64_t& last_piece_token_len) const override;
+
+  void _decode(
+    re2::StringPiece input,
+    std::string& ret) const override;
 
   // Removed negative lookahead \s+(?!\S) since it's not supported by RE2.
   const std::string _pattern =
@@ -194,10 +204,16 @@ class HFTokenizer : public BPETokenizerBase {
   void load(const std::string& tokenizer_path) override;
 
  private:
+
   void _encode(
     re2::StringPiece& input,
     std::vector<uint64_t>& ret,
     uint64_t& last_piece_token_len) const override;
 
+  void _decode(
+    re2::StringPiece input,
+    std::string& ret) const override;
+
   PreTokenizer::Ptr _pretokenizer;
+  TokenDecoder::Ptr _decoder;
 };