fix rwkv world tokenzier (mlc-ai#19)

* support rwkv world tokenizer

* refine

* rename

* refine

* switch msgpack version

* refine

* refine

* fix rwkv world tokenizer bug

* fix comment

Author: BBuf

include/rwkv_world_tokenizer.h

@@ -7,18 +7,44 @@
 #include <unordered_map>
 #include <string>
 #include <vector>
+#include <exception>
+#include <sstream>
+#include <stdexcept>
+
+#define STRINGIFY(...) STRINGIFY_(__VA_ARGS__)
+#define STRINGIFY_(...) #__VA_ARGS__
+#define RV_CHECK(...)                                                          \
+  for (bool _rv_check_status = (__VA_ARGS__); !_rv_check_status;)              \
+  throw FRException() << ("Check \"" STRINGIFY(__VA_ARGS__) "\" failed at " +  \
+                          std::to_string(__LINE__) +                           \
+                          " in " __FILE__ "\n  > Error msg: ")
+struct FRException : public std::runtime_error {
+  FRException() : std::runtime_error("") {}
+  const char *what() const noexcept override { return msg.c_str(); }
+  template <typename T> FRException &operator<<(const T &s) {
+    std::stringstream ss;
+    ss << s;
+    msg += ss.str();
+    return *this;
+  }
+  std::string msg;
+};
 
 namespace tokenizers {
-class RWKVWorldToolTokenizer {
+struct TrieTree;
+  
+class RWKVWorldToolTokenizer{
 public:
   RWKVWorldToolTokenizer(const std::string &path);
-  std::vector<int> encode(std::string_view str) const;
+  std::vector<int> encode(const std::string &str) const;
   std::string decode(const std::vector<int> &ids) const;
   std::string decode(int id) const;
 
 private:
   std::unordered_map<std::string, int> _word2idx;
   std::unordered_map<int, std::string> _idx2word;
+  std::unique_ptr<TrieTree> _tree;
 };
+
 } // namespace tokenizers
 

src/rwkv_world_tokenizer.cc

@@ -13,6 +13,57 @@
 
 namespace tokenizers {
 
+struct TrieTree {
+  std::unordered_map<int, std::unique_ptr<TrieTree>> children;
+  std::string word;
+  std::optional<int> token_id;
+  
+  TrieTree(const std::unordered_map<std::string, int>& word2id) {
+    for (auto& pair : word2id) {
+      add_word(pair.first, pair.second);
+    }
+  }
+
+  std::pair<std::string, int> find_longest_prefix(const std::string& str) const {
+    std::string prefix;
+    int token_id = -1;
+    const TrieTree* node = this;
+    for (int i = 0; i < str.size(); ++i) {
+      auto it = node->children.find(str[i]);
+      if (it == node->children.end()) {
+        break;
+      }
+      node = it->second.get();
+      RV_CHECK(node != nullptr);
+      if (node->token_id.has_value()) {
+        prefix = node->word;
+        token_id = node->token_id.value();
+      }
+    }
+    RV_CHECK(!prefix.empty());
+    RV_CHECK(token_id != -1);
+    return {prefix, token_id};
+  }
+
+  private:
+  TrieTree() = default;
+  void add_word(const std::string& word, int token_id) {
+    return _add_word(word, token_id, 0);
+  }
+  void _add_word(const std::string& word, int token_id, int idx) {
+    if (idx == word.size()) {
+      this->word = word;
+      this->token_id = token_id;
+      return;
+    }
+    auto& child = children[word[idx]];
+    if (!child) {
+      child = std::unique_ptr<TrieTree>(new TrieTree());
+    }
+    child->_add_word(word, token_id, idx + 1);
+  }
+};
+
 RWKVWorldToolTokenizer::RWKVWorldToolTokenizer(const std::string &path) {
   std::ifstream infile;
   infile.open(path, std::ios::binary | std::ios::in);
@@ -29,28 +80,17 @@ RWKVWorldToolTokenizer::RWKVWorldToolTokenizer(const std::string &path) {
   for (auto &pair : _idx2word) {
     _word2idx[pair.second] = pair.first;
   }
+  _tree = std::make_unique<TrieTree>(_word2idx);
 }
 
-std::vector<int> RWKVWorldToolTokenizer::encode(std::string_view str) const {
+std::vector<int> RWKVWorldToolTokenizer::encode(const std::string &str) const {
   std::vector<int> ids;
   int str_idx = 0;
-  int word_len = 1;
-  int id = 0;
+
   while (str_idx < str.size()) {
-    if (str_idx + word_len > str.size()) {
-      ids.push_back(id);
-      break;
-    }
-    auto substr = str.substr(str_idx, word_len);
-    auto it = _word2idx.find(std::string(substr));
-    if (it == _word2idx.end()) {
-      ids.push_back(id);
-      str_idx += (word_len - 1);
-      word_len = 1;
-    } else {
-      id = it->second;
-      word_len++;
-    }
+    auto [prefix, token_id] = _tree->find_longest_prefix(str.substr(str_idx));
+    ids.push_back(token_id);
+    str_idx += prefix.size();
   }
   return ids;
 }