spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2#include "tokenizer/bpe.h"34#include <cstdio>5#include <cstdlib>6#include <fstream>7#include <limits>8#include <sstream>910namespace forge::tok {1112namespace {13[[noreturn]] void die(const std::string& msg) {14    std::fprintf(stderr, "forge/tokenizer: %s\n", msg.c_str());15    std::abort();16}17} // namespace1819void BPETokenizer::load(const std::string& model_path) {20    std::ifstream in(model_path);21    if (!in) die("cannot open " + model_path);22    std::string header;23    std::getline(in, header);24    if (header != "forgebpe v1") die("bad model header in " + model_path);25    int64_t vocab_size = 0;26    in >> vocab_size;2728    vocab_.resize(size_t(vocab_size));29    for (int i = 0; i < 256 && i < vocab_size; ++i) vocab_[size_t(i)] = std::string(1, char(i));3031    int32_t id, left, right;32    while (in >> id >> left >> right) {33        merges_[{left, right}] = id;34        vocab_[size_t(id)] = vocab_[size_t(left)] + vocab_[size_t(right)];35    }36    if (int64_t(merges_.size()) != vocab_size - 256)37        die("merge count does not match vocab size in " + model_path);38}3940std::vector<int32_t> BPETokenizer::encode(const std::string& text) const {41    std::vector<int32_t> ids;42    ids.reserve(text.size());43    for (unsigned char c : text) ids.push_back(int32_t(c));4445    // greedy: repeatedly apply the LOWEST-id (earliest-learned) merge present46    while (ids.size() >= 2) {47        int32_t best_id = std::numeric_limits<int32_t>::max();48        std::pair<int32_t, int32_t> best_pair{-1, -1};49        for (size_t i = 0; i + 1 < ids.size(); ++i) {50            auto it = merges_.find({ids[i], ids[i + 1]});51            if (it != merges_.end() && it->second < best_id) {52                best_id = it->second;53                best_pair = it->first;54            }55        }56        if (best_pair.first < 0) break;57        std::vector<int32_t> next;58        next.reserve(ids.size());59        for (size_t i = 0; i < ids.size();) {60            if (i + 1 < ids.size() && ids[i] == best_pair.first &&61                ids[i + 1] == best_pair.second) {62                next.push_back(best_id);63                i += 2;64            } else {65                next.push_back(ids[i]);66                i += 1;67            }68        }69        ids = std::move(next);70    }71    return ids;72}7374std::string BPETokenizer::decode(const std::vector<int32_t>& ids) const {75    std::string out;76    for (int32_t id : ids) {77        if (id >= 0 && size_t(id) < vocab_.size()) out += vocab_[size_t(id)];78    }79    return out;80}8182} // namespace forge::tok83