246 lines
8.6 KiB
C++
246 lines
8.6 KiB
C++
#ifndef HUMANUS_TOKENIZER_BPE_H
|
||
#define HUMANUS_TOKENIZER_BPE_H
|
||
|
||
#include "base.h"
|
||
#include "../mcp/common/base64.hpp"
|
||
#include <unordered_map>
|
||
#include <queue>
|
||
#include <fstream>
|
||
#include <sstream>
|
||
#include <iostream>
|
||
#include <algorithm>
|
||
#include <memory>
|
||
#include <utility>
|
||
#include <functional>
|
||
#include <limits>
|
||
|
||
namespace humanus {
|
||
|
||
/**
|
||
* @brief BPE (Byte Pair Encoding) Tokenizer实现
|
||
*
|
||
* 这个类实现了基于BPE (Byte Pair Encoding)算法的分词器，
|
||
* 使用tiktoken格式的词汇表和合并规则文件。
|
||
* 使用优先队列进行高效的BPE合并操作。
|
||
*/
|
||
class BPETokenizer : public BaseTokenizer {
|
||
private:
|
||
// 辅助结构，用于哈希pair
|
||
struct PairHash {
|
||
template <class T1, class T2>
|
||
std::size_t operator() (const std::pair<T1, T2>& pair) const {
|
||
return std::hash<T1>()(pair.first) ^ std::hash<T2>()(pair.second);
|
||
}
|
||
};
|
||
|
||
// 字节对到对应token ID的映射
|
||
std::unordered_map<std::string, size_t> encoder;
|
||
// token ID到字节对的映射
|
||
std::unordered_map<size_t, std::string> decoder;
|
||
|
||
// 合并优先级映射，优先级越小越优先合并
|
||
std::unordered_map<std::pair<std::string, std::string>, size_t, PairHash> merge_ranks;
|
||
|
||
/**
|
||
* @brief 解码base64编码的字符串
|
||
* @param encoded base64编码的字符串
|
||
* @return 解码后的字符串
|
||
*/
|
||
std::string base64_decode(const std::string& encoded) const {
|
||
if (encoded.empty()) return "";
|
||
return base64::decode(encoded);
|
||
}
|
||
|
||
/**
|
||
* @brief 用于优先队列的比较函数
|
||
*
|
||
* 根据merge_ranks中的优先级比较两个token对。
|
||
* 优先级较低的值（即较小的rank值）具有更高的优先级。
|
||
*/
|
||
struct MergeComparator {
|
||
const std::unordered_map<std::pair<std::string, std::string>, size_t, PairHash>& ranks;
|
||
|
||
MergeComparator(const std::unordered_map<std::pair<std::string, std::string>, size_t, PairHash>& r)
|
||
: ranks(r) {}
|
||
|
||
bool operator()(const std::pair<std::pair<std::string, std::string>, size_t>& a,
|
||
const std::pair<std::pair<std::string, std::string>, size_t>& b) const {
|
||
// 首先按照merge_ranks比较，如果不存在则使用最大值
|
||
size_t rank_a = ranks.count(a.first) ? ranks.at(a.first) : std::numeric_limits<size_t>::max();
|
||
size_t rank_b = ranks.count(b.first) ? ranks.at(b.first) : std::numeric_limits<size_t>::max();
|
||
|
||
// 优先队列是最大堆，所以我们需要反向比较（较小的rank优先级更高）
|
||
return rank_a > rank_b;
|
||
}
|
||
};
|
||
|
||
public:
|
||
/**
|
||
* @brief 从tiktoken格式文件构造BPE tokenizer
|
||
* @param tokenizer_path tiktoken格式词汇表文件的路径
|
||
*
|
||
* 文件格式：每行包含一个base64编码的token和对应的token ID
|
||
* 例如: "IQ== 0"，其中"IQ=="是base64编码的token，0是对应的ID
|
||
*/
|
||
BPETokenizer(const std::string& tokenizer_path) {
|
||
std::ifstream file(tokenizer_path);
|
||
if (!file.is_open()) {
|
||
throw std::runtime_error("无法打开tokenizer文件: " + tokenizer_path);
|
||
}
|
||
|
||
std::string line;
|
||
while (std::getline(file, line)) {
|
||
std::istringstream iss(line);
|
||
std::string token_base64;
|
||
size_t rank;
|
||
|
||
if (iss >> token_base64 >> rank) {
|
||
std::string token = base64_decode(token_base64);
|
||
|
||
// 存储token和其ID的映射关系
|
||
encoder[token] = rank;
|
||
decoder[rank] = token;
|
||
}
|
||
}
|
||
|
||
// 构建merge_ranks
|
||
build_merge_ranks();
|
||
}
|
||
|
||
/**
|
||
* @brief 构建合并优先级映射
|
||
*
|
||
* 利用词汇表中的tokens推断可能的合并规则。
|
||
* 对于长度大于1的tokens，尝试所有可能的分割，如果分割后的两个部分也在词汇表中，
|
||
* 则假设这是一个有效的合并规则。
|
||
*/
|
||
void build_merge_ranks() {
|
||
// 对于tiktoken格式，我们可以利用编码中长度>1的token来构建合并规则
|
||
for (const auto& [token, id] : encoder) {
|
||
if (token.length() <= 1) continue;
|
||
|
||
// 尝试所有可能的分割点
|
||
for (size_t i = 1; i < token.length(); ++i) {
|
||
std::string first = token.substr(0, i);
|
||
std::string second = token.substr(i);
|
||
|
||
// 如果两个部分都在词汇表中，假设这是一个有效的合并规则
|
||
if (encoder.count(first) && encoder.count(second)) {
|
||
// 使用ID作为优先级 - 较小的ID表示更高的优先级
|
||
merge_ranks[{first, second}] = id;
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
/**
|
||
* @brief 设置合并优先级
|
||
* @param ranks 新的合并优先级映射
|
||
*/
|
||
void set_merge_ranks(const std::unordered_map<std::pair<std::string, std::string>, size_t, PairHash>& ranks) {
|
||
merge_ranks = ranks;
|
||
}
|
||
|
||
/**
|
||
* @brief 执行BPE编码
|
||
* @param text 要编码的文本
|
||
* @return 编码后的token IDs
|
||
*
|
||
* 该方法使用BPE算法对输入文本进行编码。
|
||
* 1. 首先将文本分解为单个字节
|
||
* 2. 使用优先队列根据merge_ranks中的优先级对相邻token进行合并
|
||
* 3. 将最终的tokens转换为对应的IDs
|
||
*/
|
||
std::vector<size_t> encode(const std::string& text) const override {
|
||
if (text.empty()) {
|
||
return {};
|
||
}
|
||
|
||
// 将文本分解为单个字符tokens
|
||
std::vector<std::string> tokens;
|
||
for (unsigned char c : text) {
|
||
tokens.push_back(std::string(1, c));
|
||
}
|
||
|
||
// 使用优先队列执行BPE合并
|
||
while (tokens.size() > 1) {
|
||
// 构建优先队列，用于选择最高优先级的合并对
|
||
using MergePair = std::pair<std::pair<std::string, std::string>, size_t>;
|
||
MergeComparator comparator(merge_ranks);
|
||
std::priority_queue<MergePair, std::vector<MergePair>, MergeComparator> merge_candidates(comparator);
|
||
|
||
// 查找所有可能的合并对
|
||
for (size_t i = 0; i < tokens.size() - 1; ++i) {
|
||
std::pair<std::string, std::string> pair = {tokens[i], tokens[i+1]};
|
||
if (merge_ranks.count(pair)) {
|
||
merge_candidates.push({pair, i});
|
||
}
|
||
}
|
||
|
||
// 如果没有可合并的对，退出循环
|
||
if (merge_candidates.empty()) {
|
||
break;
|
||
}
|
||
|
||
// 执行优先级最高的合并（优先队列中排在最前面的）
|
||
auto top_merge = merge_candidates.top();
|
||
auto pair = top_merge.first; // 要合并的token对
|
||
size_t pos = top_merge.second; // 要合并的位置
|
||
|
||
// 合并token
|
||
std::string merged_token = pair.first + pair.second;
|
||
tokens[pos] = merged_token;
|
||
tokens.erase(tokens.begin() + pos + 1);
|
||
}
|
||
|
||
// 将tokens转换为IDs
|
||
std::vector<size_t> ids;
|
||
ids.reserve(tokens.size());
|
||
for (const auto& token : tokens) {
|
||
auto it = encoder.find(token);
|
||
if (it != encoder.end()) {
|
||
ids.push_back(it->second);
|
||
}
|
||
// 未知token将被跳过
|
||
}
|
||
|
||
return ids;
|
||
}
|
||
|
||
/**
|
||
* @brief 解码BPE tokens
|
||
* @param tokens 要解码的token IDs
|
||
* @return 解码后的文本
|
||
*
|
||
* 该方法将编码后的token IDs转换回原始文本。
|
||
* 简单地将每个ID对应的token字符串连接起来。
|
||
*/
|
||
std::string decode(const std::vector<size_t>& tokens) const override {
|
||
std::string result;
|
||
result.reserve(tokens.size() * 2); // 预估大小，避免频繁重新分配
|
||
|
||
for (size_t id : tokens) {
|
||
auto it = decoder.find(id);
|
||
if (it != decoder.end()) {
|
||
result += it->second;
|
||
}
|
||
// 未知ID将被跳过
|
||
}
|
||
|
||
return result;
|
||
}
|
||
|
||
/**
|
||
* @brief 从tiktoken文件加载BPE词汇和合并规则
|
||
* @param file_path tiktoken格式文件的路径
|
||
* @return 共享指针指向创建的BPE tokenizer
|
||
*/
|
||
static std::shared_ptr<BPETokenizer> load_from_tiktoken(const std::string& file_path) {
|
||
return std::make_shared<BPETokenizer>(file_path);
|
||
}
|
||
};
|
||
|
||
} // namespace humanus
|
||
|
||
#endif // HUMANUS_TOKENIZER_BPE_H
|