spb/forge Public MIT

1// Author: Simon-Pierre Boucher — contact@spboucher.ai2//3// Tokenizer tests: decode(encode(x)) == x for byte-exact round-tripping,4// and — when tools/ has produced token ids for the same text — that the C++5// greedy lowest-id encoder agrees with the vectorized Python encoder used to6// build the .bin files. The test writes its own tiny .model so it needs no7// data prep; the Python-agreement half runs only if FORGE_TOK_MODEL and8// FORGE_TOK_IDS are set (wired up by tests/tokenizer_agreement.py).9#include "tokenizer/bpe.h"1011#include <cstdio>12#include <cstdlib>13#include <fstream>14#include <sstream>15#include <string>16#include <vector>1718namespace {1920int g_failures = 0;2122void expect(bool cond, const char* what) {23    if (cond) {24        std::printf("  ok: %s\n", what);25    } else {26        std::printf("  FAIL: %s\n", what);27        ++g_failures;28    }29}3031// Minimal hand-built vocab: merges chosen so encoding "aaabdaaabac" exercises32// repeated-pair overlap handling (minbpe's worked example).33void write_test_model(const std::string& path) {34    std::ofstream out(path);35    out << "forgebpe v1\n";36    out << 259 << "\n";37    out << "256 97 97\n";  // 'aa'38    out << "257 256 97\n"; // 'aaa'39    out << "258 257 98\n"; // 'aaab'40}4142} // namespace4344int main() {45    const std::string model_path = "test_tok.model";46    write_test_model(model_path);4748    forge::tok::BPETokenizer tok;49    tok.load(model_path);50    expect(tok.vocab_size() == 259, "vocab size");5152    // "aaabdaaabac": 'aaab' merges twice, leaving d/a/c as raw bytes53    const std::string text = "aaabdaaabac";54    std::vector<int32_t> ids = tok.encode(text);55    expect(ids == std::vector<int32_t>({258, 100, 258, 97, 99}), "merge order");56    expect(tok.decode(ids) == text, "round-trip (merged)");5758    // Overlap: "aaaa" applies merge 256 ('aa') to BOTH non-overlapping pairs59    // in one round, so 'aaa' (id 257) never forms — greedy is lowest-id-first,60    // not longest-match. "aaa" does reach 257 (256 then 256+97).61    expect(tok.decode(tok.encode("aaaa")) == "aaaa", "round-trip (overlap)");62    expect(tok.encode("aaaa") == std::vector<int32_t>({256, 256}), "overlap merge");63    expect(tok.encode("aaa") == std::vector<int32_t>({257}), "chained merge");6465    // Bytes with no merges, and full 0-255 range including UTF-8 and NUL66    std::string bytes;67    for (int i = 1; i < 256; ++i) bytes += char(i);68    expect(tok.decode(tok.encode(bytes)) == bytes, "round-trip (all byte values)");69    const std::string utf8 = "héllo wörld — ünïcode ✓";70    expect(tok.decode(tok.encode(utf8)) == utf8, "round-trip (utf-8)");71    expect(tok.encode("").empty(), "empty input");7273    std::remove(model_path.c_str());7475    // Optional: agreement with the Python encoder on real text.76    const char* py_model = std::getenv("FORGE_TOK_MODEL");77    const char* py_ids = std::getenv("FORGE_TOK_IDS");78    const char* py_text = std::getenv("FORGE_TOK_TEXT");79    if (py_model && py_ids && py_text) {80        forge::tok::BPETokenizer real;81        real.load(py_model);82        std::ifstream tf(py_text, std::ios::binary);83        std::stringstream ts;84        ts << tf.rdbuf();85        std::vector<int32_t> mine = real.encode(ts.str());8687        std::ifstream idf(py_ids);88        std::vector<int32_t> theirs;89        int32_t v;90        while (idf >> v) theirs.push_back(v);9192        char label[128];93        std::snprintf(label, sizeof(label),94                      "python encoder agreement (%zu vs %zu tokens)", mine.size(),95                      theirs.size());96        expect(mine == theirs, label);97        expect(real.decode(mine) == ts.str(), "round-trip (real vocab, real text)");98    } else {99        std::printf("  skip: python agreement (set FORGE_TOK_MODEL/IDS/TEXT)\n");100    }101102    if (g_failures) {103        std::printf("\n%d tokenizer test(s) FAILED\n", g_failures);104        return 1;105    }106    std::printf("\nall tokenizer tests passed\n");107    return 0;108}109