wehub/mlc-ai--mlc-llm
1
0
Fork 0
Code Issues Pull Requests Actions Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Lint / lint (push) Has been cancelled
Details
Build Docs / Deploy Docs (push) Has been cancelled
Details
Windows CI / Windows (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:23:58 +08:00
commit 770d92cb1f
694 changed files with 114634 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
cpp/tokenizers/streamer.cc
+292
View File
@@ -0,0 +1,292 @@
/*!
* Copyright (c) 2023-2025 by Contributors
* \file streamer.cc
*/
#include "streamer.h"
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <algorithm>
#include <string>
#include "tokenizers.h"
namespace mlc {
namespace llm {
TVM_FFI_STATIC_INIT_BLOCK() {
TextStreamerObj::RegisterReflection();
StopStrHandlerObj::RegisterReflection();
}
/****************** TextStreamer ******************/
TextStreamerObj::TextStreamerObj(Tokenizer tokenizer) : tokenizer_(std::move(tokenizer)) {}
TextStreamer::TextStreamer(Tokenizer tokenizer) {
data_ = tvm::ffi::make_object<TextStreamerObj>(std::move(tokenizer));
}
std::string TextStreamerObj::Put(const std::vector<int32_t>& delta_tokens) {
TVM_FFI_ICHECK(!finished_) << "`put` is not expected to be invoked after finish.";
if (delta_tokens.empty()) {
return "";
}
std::string ret;
// We process delta tokens one by one.
for (int32_t delta_token : delta_tokens) {
// push to pending tokens.
pending_tokens_.push_back(delta_token);
// all_tokens = prefix_tokens_ + pending_tokens_
std::vector<int32_t> all_tokens;
all_tokens.reserve(prefix_tokens_.size() + pending_tokens_.size());
all_tokens.insert(all_tokens.end(), prefix_tokens_.begin(), prefix_tokens_.end());
all_tokens.insert(all_tokens.end(), pending_tokens_.begin(), pending_tokens_.end());
// Decode prefix_tokens_ and all_tokens.
std::string prefix_str = prefix_tokens_.empty() ? "" : tokenizer_->Decode(prefix_tokens_);
std::string full_str = tokenizer_->Decode(all_tokens);
std::string validated_str;
std::vector<int32_t> new_pending_tokens;
if (full_str.compare(0, prefix_str.length(), prefix_str) == 0) {
// Case 1. prefix_str is a prefix of `full_str`.
// validated_str = full_str[len(prefix_str):]
validated_str = full_str.substr(prefix_str.length());
// Pop UTF-8 replacement character from the back of pending tokens.
// - The UTF-8 replacement character take 3 chars.
// - A valid UTF-8 has 4 chars at most.
// So there will be at most 3 tokens popped.
while (!pending_tokens_.empty() && //
static_cast<int>(new_pending_tokens.size()) < 3 && //
validated_str.length() >= 3 && //
validated_str.compare(validated_str.length() - 3, /*n=*/3, kReplacementCharacter) ==
0) {
new_pending_tokens.push_back(pending_tokens_.back());
pending_tokens_.pop_back();
all_tokens.pop_back();
validated_str = tokenizer_->Decode(all_tokens).substr(prefix_str.length());
}
} else {
// Case 2. prefix_str is not a prefix of `full_str`.
// Pop pending tokens from the back.
// - Pop until prefix_str is indeed a prefix of full_str.
// - A valid UTF-8 has 4 chars at most.
// So there will be at most 3 tokens popped.
// - If there are no more than 3 pending tokens, skip popping.
// This is because it is impossible to make full_str contain
// prefix_str without popping all the pending tokens.
if (static_cast<int>(pending_tokens_.size()) < 3) {
continue;
}
bool get_valid_full_str = false;
while (!pending_tokens_.empty() && static_cast<int>(new_pending_tokens.size()) < 3) {
new_pending_tokens.push_back(pending_tokens_.back());
pending_tokens_.pop_back();
all_tokens.pop_back();
full_str = tokenizer_->Decode(all_tokens);
if (full_str.compare(0, prefix_str.length(), prefix_str) == 0) {
get_valid_full_str = true;
break;
}
}
if (get_valid_full_str) {
// We find a full_str which starts from prefix_str.
// So we return the sliced full string without the prefix.
validated_str = full_str.substr(prefix_str.length());
} else {
// We cannot find a full_str which starts from prefix_str by
// popping 3 tokens.
// In this case, the remaining pending tokens are invalid UTF-8
// characters already, so we return the decoded pending tokens.
validated_str = tokenizer_->Decode(pending_tokens_);
}
}
if (!pending_tokens_.empty()) {
// Set the new prefix.
prefix_tokens_ = pending_tokens_;
}
std::reverse(new_pending_tokens.begin(), new_pending_tokens.end());
pending_tokens_ = new_pending_tokens;
ret += validated_str;
}
return ret;
}
std::string TextStreamerObj::Finish() {
// all_tokens = prefix_tokens_ + pending_tokens_
std::vector<int32_t> all_tokens;
all_tokens.reserve(prefix_tokens_.size() + pending_tokens_.size());
all_tokens.insert(all_tokens.end(), prefix_tokens_.begin(), prefix_tokens_.end());
all_tokens.insert(all_tokens.end(), pending_tokens_.begin(), pending_tokens_.end());
// Decode prefix_tokens_ and all_tokens.
std::string prefix_str = prefix_tokens_.empty() ? "" : tokenizer_->Decode(prefix_tokens_);
std::string full_str = all_tokens.empty() ? "" : tokenizer_->Decode(all_tokens);
finished_ = true;
if (full_str.compare(0, prefix_str.length(), prefix_str) == 0) {
// Case 1. prefix_str is a prefix of `full_str`.
return full_str.substr(prefix_str.length());
} else {
// Case 2. prefix_str is not a prefix of `full_str`.
// In this case, the remaining pending tokens are invalid UTF-8
// characters already, so we return the decoded pending tokens.
return tokenizer_->Decode(pending_tokens_);
}
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("mlc.tokenizers.TextStreamer",
[](Tokenizer tokenizer) { return TextStreamer(std::move(tokenizer)); })
.def("mlc.tokenizers.TextStreamerPut",
[](TextStreamer text_streamer, const Shape& delta_tokens) {
return text_streamer->Put(
{delta_tokens->data, delta_tokens->data + delta_tokens->size});
})
.def_method("mlc.tokenizers.TextStreamerFinish", &TextStreamerObj::Finish);
}
/****************** StopStrHandler ******************/
/*! \brief Create the KMP partial match table for the input string. */
inline std::vector<int> CreatePartialMatchTable(const String& str) {
int length = str.length();
std::vector<int> partial_match_table = {-1};
partial_match_table.reserve(length);
for (int i = 1; i < length; ++i) {
int ptr = partial_match_table[i - 1];
while (ptr != -1 && str.at(ptr) != str.at(i - 1)) {
ptr = partial_match_table[ptr];
}
partial_match_table.push_back(ptr + 1);
}
return partial_match_table;
}
StopStrHandlerObj::StopStrHandlerObj(Array<String> stop_strs,
const std::vector<std::string>& token_table)
: stop_strs_(std::move(stop_strs)), token_table_(token_table) {
int num_stop_strs = stop_strs_.size();
cur_match_lengths_.resize(num_stop_strs, 0);
// Create the KMP partial match table for each stop string.
partial_match_tables_.reserve(num_stop_strs);
for (const String& stop_str : stop_strs_) {
TVM_FFI_ICHECK(!stop_str.empty()) << "Stop string cannot be empty.";
partial_match_tables_.push_back(CreatePartialMatchTable(stop_str));
}
}
void StopStrHandlerObj::Put(int32_t token_id, std::vector<int64_t>* return_token_ids) {
TVM_FFI_ICHECK_NOTNULL(return_token_ids);
// Return the input token id if there is no stop string.
if (stop_strs_.empty()) {
return_token_ids->push_back(token_id);
return;
}
TVM_FFI_ICHECK(!stop_triggered_) << "Cannot put new token when already stopped.";
TVM_FFI_ICHECK_LT(token_id, static_cast<int>(token_table_.size()));
const std::string& token = token_table_[token_id];
pending_token_ids_.push_back(token_id);
pending_token_lengths_.push_back(token.length());
for (char ch : token) {
// The earliest starting point of stop string.
int stop_starting_pos = std::numeric_limits<int>::max();
// The cutoff length that can be safely return.
int cutoff_length = std::numeric_limits<int>::max();
// The maximum matched length.
int max_match_length = 0;
for (int str_id = 0; str_id < static_cast<int>(stop_strs_.size()); ++str_id) {
// - Run one step of KMP algorithm.
const std::vector<int>& partial_match_table = partial_match_tables_[str_id];
int& cur_match_length = cur_match_lengths_[str_id];
while (cur_match_length != -1 && ch != stop_strs_[str_id].at(cur_match_length)) {
cur_match_length = partial_match_table[cur_match_length];
}
++cur_match_length;
// Case 1. The stop string is matched.
if (cur_match_length == stop_strs_[str_id].length()) {
stop_triggered_ = true;
stop_starting_pos =
std::min(stop_starting_pos,
pending_string_len_ + 1 - static_cast<int>(stop_strs_[str_id].length()));
continue;
}
// Case 2. The stop string is not matched.
// - Get the cutoff length that can be safely return.
TVM_FFI_ICHECK_GE(pending_string_len_ + 1, cur_match_length);
cutoff_length = std::min(cutoff_length, pending_string_len_ + 1 - cur_match_length);
// - Get the updated pending string length.
max_match_length = std::max(max_match_length, cur_match_length);
}
// Collect the token ids that can be safely cut off and returned.
if (stop_triggered_) {
cutoff_length = stop_starting_pos;
}
TVM_FFI_ICHECK_NE(cutoff_length, std::numeric_limits<int>::max());
TVM_FFI_ICHECK_GE(cutoff_length, 0);
int cum_length = 0;
while (!pending_token_ids_.empty() &&
cum_length + pending_token_lengths_.front() <= cutoff_length) {
cum_length += pending_token_lengths_.front();
return_token_ids->push_back(pending_token_ids_.front());
pending_token_ids_.erase(pending_token_ids_.begin());
pending_token_lengths_.erase(pending_token_lengths_.begin());
}
if (stop_triggered_) {
return;
}
TVM_FFI_ICHECK_LE(cum_length, cutoff_length);
// `cum_length` is the prefix length what we actually cut off.
pending_string_len_ = (cutoff_length - cum_length) + max_match_length;
}
}
StopStrHandler::StopStrHandler(Array<String> stop_strs,
const std::vector<std::string>& token_table) {
data_ = tvm::ffi::make_object<StopStrHandlerObj>(std::move(stop_strs), token_table);
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("mlc.tokenizers.StopStrHandler",
[](Array<String> stop_strs, const Tokenizer& tokenizer) {
return StopStrHandler(std::move(stop_strs), tokenizer->PostProcessedTokenTable());
})
.def("mlc.tokenizers.StopStrHandlerPut",
[](StopStrHandler handler, int token_id) {
std::vector<int64_t> delta_tokens;
handler->Put(token_id, &delta_tokens);
return Shape(std::move(delta_tokens));
})
.def("mlc.tokenizers.StopStringHandlerFinish",
[](StopStrHandler handler) {
std::vector<int64_t> remaining_token_ids;
handler->Finish(&remaining_token_ids);
return Shape(std::move(remaining_token_ids));
})
.def_method("mlc.tokenizers.StopStrHandlerStopTriggered", &StopStrHandlerObj::StopTriggered);
}
} // namespace llm
} // namespace mlc
cpp/tokenizers/streamer.h
+160
View File
@@ -0,0 +1,160 @@
/*!
* Copyright (c) 2023-2025 by Contributors
* \file streamer.h
* \brief Header of streamers in MLC LLM.
*/
#ifndef MLC_LLM_STREAMER_H_
#define MLC_LLM_STREAMER_H_
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include "tokenizers.h"
namespace mlc {
namespace llm {
using namespace tvm::runtime;
using tvm::ffi::Object;
using tvm::ffi::ObjectRef;
/****************** TextStreamer ******************/
/*!
* \brief The class that streams back validated utf-8 text strings
* that generated by tokenizer.
*/
class TextStreamerObj : public Object {
public:
explicit TextStreamerObj(Tokenizer tokenizer);
/*!
* \brief Put new delta tokens into the streamer, and get the UTF-8-valid
* delta string. The text streamer may hold some of the input delta tokens
* which cannot decode into valid UTF-8 strings. The returned string
* is always guaranteed to be UTF-8 valid.
* \param delta_tokens The new tokens to put into the streamer.
* \return The decoded delta string after putting the input new tokens.
*/
std::string Put(const std::vector<int32_t>& delta_tokens);
/*! \brief Return the string decoded by remaining tokens. */
std::string Finish();
// REPLACEMENT CHARACTER (U+FFFD) in UTF-8.
static constexpr const char* kReplacementCharacter = "\xef\xbf\xbd";
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TextStreamerObj>();
}
static constexpr const bool _type_has_method_sequal_reduce = false;
static constexpr const bool _type_has_method_shash_reduce = false;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("mlc.TextStreamer", TextStreamerObj, Object);
private:
Tokenizer tokenizer_;
std::vector<int32_t> prefix_tokens_;
std::vector<int32_t> pending_tokens_;
bool finished_ = false;
};
/*!
* \brief Managed reference to TextStreamerObj
* \sa TextStreamerObj
*/
class TextStreamer : public ObjectRef {
public:
/*! \brief Construct a text streamer with tokenizer. */
explicit TextStreamer(Tokenizer tokenizer);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TextStreamer, ObjectRef, TextStreamerObj);
};
/****************** StopStrHandler ******************/
/*!
* \brief The stop string handler in MLC LLM, which takes input delta tokens
* one at a time, and return the output delta token before stopping due to
* stop strings.
*/
class StopStrHandlerObj : public Object {
public:
explicit StopStrHandlerObj(Array<String> stop_strs, const std::vector<std::string>& token_table);
/*!
* \brief Add new input delta token to the handler, push the output
* delta tokens before stopping into the given vector.
* The stop string handler may hold some of the input delta token
* which may be part of a stop string.
* The returned tokens are always guaranteed not to be part of stop string.
*/
void Put(int32_t token_id, std::vector<int64_t>* return_token_ids);
/*!
* \brief Stop string handling has finished, append the remaining
* cached token ids into the given vector.
*/
void Finish(std::vector<int64_t>* return_token_ids) const {
return_token_ids->insert(return_token_ids->end(), pending_token_ids_.begin(),
pending_token_ids_.end());
};
/*! \brief Check if the generation has stopped due to stop string. */
bool StopTriggered() const { return stop_triggered_; }
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<StopStrHandlerObj>();
}
static constexpr const bool _type_has_method_sequal_reduce = false;
static constexpr const bool _type_has_method_shash_reduce = false;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.StopStrHandler", StopStrHandlerObj, Object);
private:
/*! \brief The stop strings. */
Array<String> stop_strs_;
/*! \brief The partial match table for each stop string in the KMP algorithm. */
std::vector<std::vector<int>> partial_match_tables_;
/*! \brief The tokenizer token table for token id lookup. */
const std::vector<std::string>& token_table_;
/************ Global states across all stop strings. ************/
/*! \brief The globally pending string length. */
int pending_string_len_ = 0;
/*! \brief The globally pending token ids. */
std::vector<int32_t> pending_token_ids_;
/*! \brief The token string length of each pending token id. */
std::vector<int> pending_token_lengths_;
/*! \brief A boolean flag indicating if stop has been triggered. */
bool stop_triggered_ = false;
/************ Per-stop-string states. ************/
/*! \brief The current match position of the pending string to each stop string. */
std::vector<int> cur_match_lengths_;
};
/*!
* \brief Managed reference to StopStrHandlerObj
* \sa StopStrHandlerObj
*/
class StopStrHandler : public ObjectRef {
public:
explicit StopStrHandler(Array<String> stop_strs, const std::vector<std::string>& token_table);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(StopStrHandler, ObjectRef, StopStrHandlerObj);
};
} // namespace llm
} // namespace mlc
#endif // MLC_LLM_STREAMER_H_
cpp/tokenizers/tokenizers.cc
+535
View File
@@ -0,0 +1,535 @@
/*!
* Copyright (c) 2023-2025 by Contributors
* \file tokenizer.cc
*/
#include "tokenizers.h"
#include <tokenizers_cpp.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/extra/json.h>
#include <tvm/ffi/function.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/runtime/logging.h>
#include <array>
#include <filesystem>
#include <fstream>
#include <string>
#include <string_view>
#include "./../support/encoding.h"
#include "./../support/load_bytes_from_file.h"
namespace mlc {
namespace llm {
TVM_FFI_STATIC_INIT_BLOCK() {
TokenizerInfoNode::RegisterReflection();
TokenizerObj::RegisterReflection();
}
#ifndef COMPILE_MLC_WASM_RUNTIME
String TokenizerInfoNode::AsJSONString() const {
tvm::ffi::json::Object obj;
obj.Set("token_postproc_method", token_postproc_method);
obj.Set("prepend_space_in_encode", prepend_space_in_encode);
obj.Set("strip_space_in_decode", strip_space_in_decode);
return tvm::ffi::json::Stringify(obj);
}
TokenizerInfo TokenizerInfo::FromJSONString(String json_string) {
tvm::ffi::String err;
auto v = tvm::ffi::json::Parse(json_string, &err);
TVM_FFI_ICHECK(err.empty()) << "Failed to parse JSON: " << err;
TVM_FFI_ICHECK(v.try_cast<tvm::ffi::json::Object>().has_value()) << "JSON must be an object.";
const auto& obj = v.cast<tvm::ffi::json::Object>();
ObjectPtr<TokenizerInfoNode> n = tvm::ffi::make_object<TokenizerInfoNode>();
if (obj.count("token_postproc_method")) {
TVM_FFI_ICHECK(obj.at("token_postproc_method").try_cast<tvm::ffi::String>().has_value());
n->token_postproc_method = obj.at("token_postproc_method").cast<tvm::ffi::String>();
}
if (obj.count("prepend_space_in_encode")) {
TVM_FFI_ICHECK(obj.at("prepend_space_in_encode").try_cast<bool>().has_value());
n->prepend_space_in_encode = obj.at("prepend_space_in_encode").cast<bool>();
}
if (obj.count("strip_space_in_decode")) {
TVM_FFI_ICHECK(obj.at("strip_space_in_decode").try_cast<bool>().has_value());
n->strip_space_in_decode = obj.at("strip_space_in_decode").cast<bool>();
}
return TokenizerInfo(n);
}
Tokenizer::Tokenizer(std::unique_ptr<tokenizers::Tokenizer> tokenizer, TokenizerInfo info) {
ObjectPtr<TokenizerObj> n = tvm::ffi::make_object<TokenizerObj>();
n->tokenizer = std::move(tokenizer);
n->info_ = std::move(info);
data_ = std::move(n);
}
std::vector<int32_t> TokenizerObj::Encode(const std::string& text) const {
return tokenizer->Encode(text);
}
std::vector<int32_t> TokenizerObj::EncodeNoPrependSpace(const std::string& text) const {
// TODO(yixin): now this only supports tokenizers with tokenizer.json
// other tokenizers should be supported.
static const constexpr char* kPaddingPrefix = "\x01";
if (!info_->prepend_space_in_encode) {
return tokenizer->Encode(text);
}
auto result = tokenizer->Encode(kPaddingPrefix + text);
// remove the first two tokens: "▁" and "<0x01>"
result.erase(result.begin(), result.begin() + 2);
return result;
}
std::vector<std::vector<int32_t>> TokenizerObj::EncodeBatch(const Array<String>& texts) const {
std::vector<std::string> texts_vec;
for (const String& text : texts) {
texts_vec.push_back(text);
}
return tokenizer->EncodeBatch(texts_vec);
}
std::string TokenizerObj::Decode(const std::vector<int32_t>& token_ids) const {
return tokenizer->Decode(token_ids);
}
const DynamicBitset& TokenizerObj::GetPrefixTokenMask() {
if (prefix_token_mask_.Size() != 0) {
return prefix_token_mask_;
}
int vocab_size = GetVocabSize();
prefix_token_mask_ = DynamicBitset(vocab_size);
// Sort all tokens
const auto& token_table = PostProcessedTokenTable();
std::vector<std::pair<std::string, int>> sorted_tokens;
for (int32_t token_id = 0; token_id < vocab_size; ++token_id) {
sorted_tokens.emplace_back(token_table[token_id], token_id);
}
std::sort(sorted_tokens.begin(), sorted_tokens.end());
// Check every token if it is a prefix of another token
for (int idx = 0; idx < vocab_size - 1; ++idx) {
auto cur_token = sorted_tokens[idx].first;
auto nxt_token = sorted_tokens[idx + 1].first;
if (cur_token.length() <= nxt_token.length() &&
std::string_view(nxt_token).substr(0, cur_token.length()) == cur_token) {
prefix_token_mask_.Set(sorted_tokens[idx].second);
}
}
return prefix_token_mask_;
}
size_t TokenizerObj::GetVocabSize() const { return tokenizer->GetVocabSize(); }
std::string TokenizerObj::IdToToken(int32_t token_id) const {
return tokenizer->IdToToken(token_id);
}
int32_t TokenizerObj::TokenToId(const std::string& token) const {
return tokenizer->TokenToId(token);
}
Tokenizer Tokenizer::FromPath(const String& _path, std::optional<TokenizerInfo> info) {
TokenizerInfo info_value = info.value_or(DetectTokenizerInfo(_path));
std::filesystem::path path{std::string(_path)};
std::filesystem::path sentencepiece;
std::filesystem::path huggingface;
std::filesystem::path rwkvworld;
TVM_FFI_ICHECK(std::filesystem::exists(path)) << "Cannot find tokenizer via path: " << _path;
if (std::filesystem::is_directory(path)) {
sentencepiece = path / "tokenizer.model";
huggingface = path / "tokenizer.json";
rwkvworld = path / "tokenizer_model";
} else {
sentencepiece = path.parent_path() / "tokenizer.model";
huggingface = path.parent_path() / "tokenizer.json";
rwkvworld = path.parent_path() / "tokenizer_model";
}
if (std::filesystem::exists(huggingface)) {
// Check HuggingFace
return Tokenizer(tokenizers::Tokenizer::FromBlobJSON(LoadBytesFromFile(huggingface.string())),
info_value);
}
if (std::filesystem::exists(sentencepiece)) {
// Check SentencePiece
LOG(WARNING)
<< "Using `tokenizer.model` since we cannot locate `tokenizer.json`.\n"
<< "It is recommended to use `tokenizer.json` to ensure all token mappings are included, "
<< "since currently, files like `added_tokens.json`, `tokenizer_config.json` are ignored.\n"
<< "Consider converting `tokenizer.model` to `tokenizer.json` by compiling the model "
<< "with MLC again, or see if MLC's huggingface provides this file.";
return Tokenizer(
tokenizers::Tokenizer::FromBlobSentencePiece(LoadBytesFromFile(sentencepiece.string())),
info_value);
}
{
// Check ByteLevelBPE
std::filesystem::path merges_path = path / "merges.txt";
std::filesystem::path vocab_path = path / "vocab.json";
std::filesystem::path added_tokens_path = path / "added_tokens.json";
if (std::filesystem::exists(merges_path) && std::filesystem::exists(vocab_path) &&
std::filesystem::exists(added_tokens_path)) {
std::string vocab = LoadBytesFromFile(vocab_path.string());
std::string merges = LoadBytesFromFile(merges_path.string());
std::string added_tokens = LoadBytesFromFile(added_tokens_path.string());
return Tokenizer(tokenizers::Tokenizer::FromBlobByteLevelBPE(vocab, merges, added_tokens),
info_value);
}
}
if (std::filesystem::exists(rwkvworld)) {
// Check RWKV
return Tokenizer(tokenizers::Tokenizer::FromBlobRWKVWorld(rwkvworld.string()), info_value);
}
LOG(FATAL) << "Cannot find any tokenizer under: " << _path;
}
TokenizerInfo Tokenizer::DetectTokenizerInfo(const String& path_str) {
std::filesystem::path path{std::string(path_str)};
TVM_FFI_ICHECK(std::filesystem::exists(path)) << "Cannot find tokenizer via path: " << path_str;
if (!std::filesystem::is_directory(path)) {
path = path.parent_path();
}
path = path / "tokenizer.json";
if (!std::filesystem::exists(path)) {
LOG(WARNING) << "Tokenizer info is not detected as tokenizer.json is not found. The default "
<< "tokenizer info will be used.";
return TokenizerInfo(tvm::ffi::make_object<TokenizerInfoNode>());
}
std::string tokenizer_json = LoadBytesFromFile(path.string());
tvm::ffi::String err;
auto v = tvm::ffi::json::Parse(tokenizer_json, &err);
TVM_FFI_ICHECK(err.empty()) << "Failed to parse JSON: " << err;
TVM_FFI_ICHECK(v.try_cast<tvm::ffi::json::Object>().has_value()) << "JSON must be an object.";
const auto& obj = v.cast<tvm::ffi::json::Object>();
ObjectPtr<TokenizerInfoNode> n = tvm::ffi::make_object<TokenizerInfoNode>();
// Step 1. Detect token_postproc_method: byte_fallback or byte_level
// Detect {"type": "ByteLevel"} or {"type": "ByteFallback"} in "decoder" field of the tokenizer
if (!obj.count("decoder") || !obj.at("decoder").try_cast<tvm::ffi::json::Object>().has_value()) {
LOG(WARNING) << "Decoder field is not found in tokenizer.json. Use ByteFallback as default.";
n->token_postproc_method = "byte_fallback";
} else {
auto decoder_obj = obj.at("decoder").cast<tvm::ffi::json::Object>();
TVM_FFI_ICHECK(decoder_obj.count("type") &&
decoder_obj.at("type").try_cast<tvm::ffi::String>().has_value());
auto type = decoder_obj.at("type").cast<tvm::ffi::String>();
auto f_detect_decoder_type = [](ObjectPtr<TokenizerInfoNode> n,
const tvm::ffi::json::Value& decoder_json) {
TVM_FFI_ICHECK(decoder_json.try_cast<tvm::ffi::json::Object>().has_value());
TVM_FFI_ICHECK(decoder_json.cast<tvm::ffi::json::Object>().count("type") &&
decoder_json.cast<tvm::ffi::json::Object>()
.at("type")
.try_cast<tvm::ffi::String>()
.has_value());
auto type = decoder_json.cast<tvm::ffi::json::Object>().at("type").cast<tvm::ffi::String>();
if (type == "ByteLevel") {
n->token_postproc_method = "byte_level";
return true;
} else if (type == "ByteFallback") {
n->token_postproc_method = "byte_fallback";
return true;
}
return false;
};
bool found = false;
// For sequence, examine every decoder
if (type == "Sequence") {
TVM_FFI_ICHECK(decoder_obj.count("decoders") &&
decoder_obj.at("decoders").try_cast<tvm::ffi::json::Array>().has_value());
for (const tvm::ffi::json::Value& decoder :
decoder_obj.at("decoders").cast<tvm::ffi::json::Array>()) {
if (f_detect_decoder_type(n, decoder)) {
found = true;
}
}
} else {
if (f_detect_decoder_type(n, obj.at("decoder"))) {
found = true;
}
}
if (!found) {
LOG(WARNING) << "Neither ByteLevel nor ByteFallback decoder is detected in tokenizer.json. "
<< "Use ByteFallback as default.";
n->token_postproc_method = "byte_fallback";
}
}
// Step 2. Detect prepend_space_in_encode
// Find {"type": "Prepend", "prepend": "▁"} in "normalizer" field of the tokenizer
if (obj.count("normalizer") &&
obj.at("normalizer").try_cast<tvm::ffi::json::Object>().has_value()) {
const tvm::ffi::json::Value& normalizer_json = obj.at("normalizer");
auto f_handle_normalizer = [](ObjectPtr<TokenizerInfoNode> n,
const tvm::ffi::json::Value& normalizer_json) {
TVM_FFI_ICHECK(normalizer_json.try_cast<tvm::ffi::json::Object>().has_value());
auto obj = normalizer_json.cast<tvm::ffi::json::Object>();
TVM_FFI_ICHECK(obj.count("type") && obj.at("type").try_cast<tvm::ffi::String>().has_value());
if (obj.at("type").cast<tvm::ffi::String>() == "Prepend" && obj.count("prepend") &&
obj.at("prepend").try_cast<tvm::ffi::String>().has_value() &&
obj.at("prepend").cast<tvm::ffi::String>() == "\xe2\x96\x81") {
n->prepend_space_in_encode = true;
return true;
}
return false;
};
auto type = normalizer_json.cast<tvm::ffi::json::Object>().at("type").cast<tvm::ffi::String>();
if (type == "Sequence") {
TVM_FFI_ICHECK(normalizer_json.cast<tvm::ffi::json::Object>().count("normalizers") &&
normalizer_json.cast<tvm::ffi::json::Object>()
.at("normalizers")
.try_cast<tvm::ffi::json::Array>()
.has_value());
for (const tvm::ffi::json::Value& normalizer : normalizer_json.cast<tvm::ffi::json::Object>()
.at("normalizers")
.cast<tvm::ffi::json::Array>()) {
if (f_handle_normalizer(n, normalizer)) {
break;
}
}
} else {
f_handle_normalizer(n, normalizer_json);
}
}
// Step 3. Detect strip_space_in_decode
// Find {"type": "Strip", "content": " ", "start": 1, "stop": 0} in "decoder" field of the
// tokenizer
if (obj.count("decoder") && obj.at("decoder").try_cast<tvm::ffi::json::Object>().has_value()) {
const tvm::ffi::json::Value& decoders_json = obj.at("decoder");
auto f_handle_decoder = [](ObjectPtr<TokenizerInfoNode> n,
const tvm::ffi::json::Value& decoder_json) {
TVM_FFI_ICHECK(decoder_json.try_cast<tvm::ffi::json::Object>().has_value());
auto obj = decoder_json.cast<tvm::ffi::json::Object>();
TVM_FFI_ICHECK(obj.count("type") && obj.at("type").try_cast<tvm::ffi::String>().has_value());
if (obj.at("type").cast<tvm::ffi::String>() == "Strip" && obj.count("content") &&
obj.at("content").try_cast<tvm::ffi::String>().has_value() &&
obj.at("content").cast<tvm::ffi::String>() == " " && obj.count("start") &&
obj.at("start").try_cast<int64_t>().has_value() && obj.at("start").cast<int64_t>() == 1 &&
obj.count("stop") && obj.at("stop").try_cast<int64_t>().has_value() &&
obj.at("stop").cast<int64_t>() == 0) {
n->strip_space_in_decode = true;
return true;
}
return false;
};
auto type = decoders_json.cast<tvm::ffi::json::Object>().at("type").cast<tvm::ffi::String>();
if (type == "Sequence") {
TVM_FFI_ICHECK(decoders_json.cast<tvm::ffi::json::Object>().count("decoders") &&
decoders_json.cast<tvm::ffi::json::Object>()
.at("decoders")
.try_cast<tvm::ffi::json::Array>()
.has_value());
for (const tvm::ffi::json::Value& decoder : decoders_json.cast<tvm::ffi::json::Object>()
.at("decoders")
.cast<tvm::ffi::json::Array>()) {
if (f_handle_decoder(n, decoder)) {
break;
}
}
} else {
f_handle_decoder(n, decoders_json);
}
}
return TokenizerInfo(n);
}
#endif
/*! \brief ByteFallback decoder: transform tokens like <0x1B> to hex char byte 1B */
inline std::string ByteFallbackDecoder(const std::string& token) {
if (token.length() == 6 && token.substr(0, 3) == "<0x" && token.back() == '>') {
int byte = 0;
for (int i = 0; i < 2; ++i) {
byte *= 16;
byte +=
token[3 + i] >= '0' && token[3 + i] <= '9' ? token[3 + i] - '0' : token[3 + i] - 'A' + 10;
}
TVM_FFI_ICHECK(byte >= 0 && byte < 256);
return std::string(/*n=*/1, static_cast<char>(byte));
}
return token;
}
/*! \brief SpaceReplacer decoder: transform "\u2581" back to space */
inline std::string SpaceReplacerDecoder(const std::string& token) {
// \u2581 is the unicode for "lower one eighth block"
// UTF8 encoding for \u2581 is 0xE2 0x96 0x81
std::string result;
for (size_t i = 0; i < token.size(); ++i) {
if (i + 2 < token.size() && token[i] == char(0xE2) && token[i + 1] == char(0x96) &&
token[i + 2] == char(0x81)) {
result += ' ';
i += 2;
} else {
result += token[i];
}
}
return result;
}
/*! \brief ByteLevel decoder: inverses the bytes-to-unicode transformation in the encoding
* process as in
* https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59
*/
inline std::string ByteLevelDecoder(const std::string& token) {
// clang-format off
// The inverse map of bytes_to_unicode. -1 means there is no mapping to this unicode.
static const std::array<int, 324> char_to_byte_map = {
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45,
46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68,
69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84, 85, 86, 87, 88, 89, 90, 91,
92, 93, 94, 95, 96, 97, 98, 99, 100, 101, 102, 103, 104, 105, 106, 107, 108, 109, 110, 111,
112, 113, 114, 115, 116, 117, 118, 119, 120, 121, 122, 123, 124, 125, 126, -1, -1, -1, -1,
-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1,
-1, -1, -1, -1, -1, -1, -1, 161, 162, 163, 164, 165, 166, 167, 168, 169, 170, 171, 172, -1,
174, 175, 176, 177, 178, 179, 180, 181, 182, 183, 184, 185, 186, 187, 188, 189, 190, 191,
192, 193, 194, 195, 196, 197, 198, 199, 200, 201, 202, 203, 204, 205, 206, 207, 208, 209,
210, 211, 212, 213, 214, 215, 216, 217, 218, 219, 220, 221, 222, 223, 224, 225, 226, 227,
228, 229, 230, 231, 232, 233, 234, 235, 236, 237, 238, 239, 240, 241, 242, 243, 244, 245,
246, 247, 248, 249, 250, 251, 252, 253, 254, 255, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12,
13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 127, 128,
129, 130, 131, 132, 133, 134, 135, 136, 137, 138, 139, 140, 141, 142, 143, 144, 145, 146,
147, 148, 149, 150, 151, 152, 153, 154, 155, 156, 157, 158, 159, 160, 173
};
// clang-format on
auto unicode_codepoints = ParseUTF8(token.c_str(), UTF8ErrorPolicy::kReturnInvalid);
if (unicode_codepoints.size() == 1 && unicode_codepoints[0] == kInvalidUTF8) {
return token;
}
std::string decoded;
for (auto unicode_codepoint : unicode_codepoints) {
TVM_FFI_ICHECK(unicode_codepoint >= 0);
if (unicode_codepoint >= static_cast<int>(char_to_byte_map.size()) ||
char_to_byte_map[unicode_codepoint] == -1) {
// If there is no mapping, return the original token
return token;
}
decoded += static_cast<char>(char_to_byte_map[unicode_codepoint]);
}
return decoded;
}
/*!
* \brief Post-process a raw token to the actual token with the given post-processing method.
*/
inline std::string PostProcessToken(const std::string& token,
const std::string& token_postproc_method) {
if (token_postproc_method == "byte_fallback") {
return SpaceReplacerDecoder(ByteFallbackDecoder(token));
} else if (token_postproc_method == "byte_level") {
return ByteLevelDecoder(token);
} else {
LOG(FATAL) << "Unknown post-processing method: " << token_postproc_method;
}
}
std::vector<std::string> Tokenizer::PostProcessTokenTable(
const std::vector<std::string>& token_table, const std::string& token_postproc_method) {
std::vector<std::string> post_processed_token_table;
post_processed_token_table.reserve(token_table.size());
for (const std::string& token : token_table) {
post_processed_token_table.push_back(PostProcessToken(token, token_postproc_method));
}
return post_processed_token_table;
}
#ifndef COMPILE_MLC_WASM_RUNTIME
const std::vector<std::string>& TokenizerObj::PostProcessedTokenTable() {
if (!post_processed_token_table_.empty()) {
return post_processed_token_table_;
}
std::vector<std::string> raw_token_table;
int vocab_size = tokenizer->GetVocabSize();
raw_token_table.reserve(vocab_size);
for (int32_t token_id = 0; token_id < vocab_size; ++token_id) {
raw_token_table.push_back(tokenizer->IdToToken(token_id));
}
post_processed_token_table_ =
Tokenizer::PostProcessTokenTable(raw_token_table, info_->token_postproc_method);
return post_processed_token_table_;
}
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def("mlc.tokenizers.Tokenizer", [](const String& path) { return Tokenizer::FromPath(path); })
.def("mlc.tokenizers.TokenizerEncode",
[](const Tokenizer& tokenizer, const String& text) {
std::vector<int32_t> token_ids = tokenizer->Encode(text);
return Shape{token_ids.begin(), token_ids.end()};
})
.def("mlc.tokenizers.TokenizerEncodeBatch",
[](const Tokenizer& tokenizer, const Array<String>& texts) {
std::vector<std::vector<int32_t>> results = tokenizer->EncodeBatch(texts);
Array<Shape> ret;
ret.reserve(results.size());
for (const auto& result : results) {
ret.push_back(Shape{result.begin(), result.end()});
}
return ret;
})
.def("mlc.tokenizers.TokenizerDecode",
[](const Tokenizer& tokenizer, const Shape& token_ids) {
return tokenizer->Decode({token_ids->data, token_ids->data + token_ids->size});
})
.def("mlc.tokenizers.DetectTokenizerInfo",
[](const String& path) { return Tokenizer::DetectTokenizerInfo(path)->AsJSONString(); });
}
#endif
TVM_FFI_STATIC_INIT_BLOCK() {
namespace refl = tvm::ffi::reflection;
refl::GlobalDef()
.def_packed("mlc.tokenizers.PostProcessTokenTable",
[](tvm::ffi::PackedArgs args, tvm::ffi::Any* rv) {
Array<String> token_table_arr = args[0].cast<Array<String>>();
std::string token_postproc_method = args[args.size() - 1].cast<String>();
std::vector<std::string> token_table;
for (int i = 0; i < token_table_arr.size(); ++i) {
token_table.push_back(token_table_arr[i]);
}
std::vector<std::string> processed_token_table =
Tokenizer::PostProcessTokenTable(token_table, token_postproc_method);
// Convert std::vector<std::string> to Array<String>
Array<String> processed_token_table_tvm;
for (int i = 0; i < processed_token_table.size(); ++i) {
processed_token_table_tvm.push_back(processed_token_table[i]);
}
*rv = processed_token_table_tvm;
})
.def("mlc.tokenizers.PostProcessToken",
[](const String& token, const String& token_postproc_method) {
return PostProcessToken(token, token_postproc_method);
});
}
} // namespace llm
} // namespace mlc
cpp/tokenizers/tokenizers.h
+171
View File
@@ -0,0 +1,171 @@
/*!
* Copyright (c) 2023-2025 by Contributors
* \file tokenizers.h
* \brief Header of tokenizer related functions.
*/
#ifndef MLC_LLM_TOKENIZER_H_
#define MLC_LLM_TOKENIZER_H_
#include <tokenizers_cpp.h>
#include <tvm/ffi/container/array.h>
#include <tvm/ffi/container/shape.h>
#include <tvm/ffi/object.h>
#include <tvm/ffi/reflection/registry.h>
#include <tvm/ffi/string.h>
#include <optional>
#include <unordered_map>
#include "../base.h"
#include "../support/dynamic_bitset.h"
namespace mlc {
namespace llm {
using namespace tvm::runtime;
using tvm::ffi::Array;
using tvm::ffi::Object;
using tvm::ffi::ObjectPtr;
using tvm::ffi::ObjectRef;
using tvm::ffi::Shape;
using tvm::ffi::String;
/*! \brief Useful information of the tokenizer during generation. */
class TokenizerInfoNode : public Object {
public:
/*! \brief The method to post-process the tokens to their original strings.
* Possible values (each refers to a kind of tokenizer):
* - "byte_fallback": The same as the byte-fallback BPE tokenizer, including LLaMA-2,
* Mixtral-7b, etc. E.g. "▁of" -> " of", "<0x1B>" -> "\x1B".
* This method:
* 1) Transform tokens like <0x1B> to hex char byte 1B. (so-called byte-fallback)
* 2) Replace \\u2581 "▁" with space.
* - "byte_level": The same as the byte-level BPE tokenizer, including LLaMA-3, GPT-2,
* Phi-2, etc. E.g. "Ġin" -> " in", "ě" -> "\x1B"
* This method inverses the bytes-to-unicode transformation in the encoding process in
* https://github.com/huggingface/transformers/blob/87be06ca77166e6a6215eee5a990ab9f07238a18/src/transformers/models/gpt2/tokenization_gpt2.py#L38-L59
*/
String token_postproc_method = "byte_fallback";
/*! \brief Whether to prepend a space during encoding. */
bool prepend_space_in_encode = false;
/*! \brief Whether to strip the first space during decoding. */
bool strip_space_in_decode = false;
String AsJSONString() const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TokenizerInfoNode>();
}
static constexpr const bool _type_has_method_sequal_reduce = false;
static constexpr const bool _type_has_method_shash_reduce = false;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO("mlc.serve.TokenizerInfo", TokenizerInfoNode, Object);
};
class TokenizerInfo : public ObjectRef {
public:
/*! \brief Create a TokenizerInfo object from a dumped string. */
static TokenizerInfo FromJSONString(String json_string);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(TokenizerInfo, ObjectRef, TokenizerInfoNode);
};
/*! \brief A wrapper object class for tokenizer. */
class TokenizerObj : public Object {
public:
/*! \brief The underlying tokenizer. */
std::unique_ptr<tokenizers::Tokenizer> tokenizer;
/*! \brief Encode text into ids. */
std::vector<int32_t> Encode(const std::string& text) const;
/*! \brief Encode text into ids. Some tokenizers may prepend a space in encoding, this method
* guarantees the space is not prepended. */
std::vector<int32_t> EncodeNoPrependSpace(const std::string& text) const;
/*! \brief Encode texts into ids. */
std::vector<std::vector<int32_t>> EncodeBatch(const Array<String>& texts) const;
/*! \brief Decode token ids into text. */
std::string Decode(const std::vector<int32_t>& token_ids) const;
/*! \brief Return the post-processed token table of the tokenizer. Special tokens are included. */
const std::vector<std::string>& PostProcessedTokenTable();
/*! \brief Get the prefix token mask as a bitset. The tokens which is a prefix of another token
* are set to true, and others are set to false in the bitset. */
const DynamicBitset& GetPrefixTokenMask();
/*!
* \brief Returns the vocabulary size. Special tokens are considered. This may be smaller than the
* `vocab_size` in config.json (length of logits), see https://github.com/QwenLM/Qwen2/issues/147
* and https://huggingface.co/microsoft/Phi-3-mini-4k-instruct/discussions/47.
*/
size_t GetVocabSize() const;
/*!
* \brief Convert the given id to its corresponding token if it exists. If not, return an
* empty string.
*/
std::string IdToToken(int32_t token_id) const;
/*!
* \brief Convert the given token to its corresponding id if it exists. If not, return -1.
*/
int32_t TokenToId(const std::string& token) const;
static void RegisterReflection() {
namespace refl = tvm::ffi::reflection;
refl::ObjectDef<TokenizerObj>();
}
friend class Tokenizer;
static constexpr const bool _type_has_method_sequal_reduce = false;
static constexpr const bool _type_has_method_shash_reduce = false;
static constexpr const bool _type_mutable = true;
TVM_FFI_DECLARE_OBJECT_INFO_FINAL("mlc.Tokenizer", TokenizerObj, Object);
private:
/*! \brief Useful information of the tokenizer during generation. */
TokenizerInfo info_;
/*! \brief The cached token table. */
std::vector<std::string> post_processed_token_table_;
/*! \brief The cached prefix token mask. */
DynamicBitset prefix_token_mask_;
};
class Tokenizer : public ObjectRef {
public:
/*!
* \brief Create a tokenizer from a directory path on disk.
* \param path The path to the tokenizer or the tokenizer directory.
* \param info The tokenizer info. If not provided, the info will be detected automatically.
*/
MLC_LLM_DLL static Tokenizer FromPath(const String& path,
std::optional<TokenizerInfo> info = std::nullopt);
/*! \brief Detect the tokenizer info from the given path of the tokenizer. */
MLC_LLM_DLL static TokenizerInfo DetectTokenizerInfo(const String& path);
/*!
* \brief Post-process the token table to their original strings.
* \param token_table The raw token table.
* \param postproc_method The postprocessing method to use.
* \returns The postprocessed token table containing the original strings.
*/
static std::vector<std::string> PostProcessTokenTable(const std::vector<std::string>& token_table,
const std::string& token_postproc_method);
TVM_FFI_DEFINE_OBJECT_REF_METHODS_NULLABLE(Tokenizer, ObjectRef, TokenizerObj);
private:
explicit Tokenizer(std::unique_ptr<tokenizers::Tokenizer> tokenizer, TokenizerInfo info);
};
} // namespace llm
} // namespace mlc
#endif // MLC_LLM_TOKENIZER_H_