Files
2026-07-13 13:33:03 +08:00

318 lines
12 KiB
C++

//
// test_tokenizer.cpp
//
// Tokenizer unit test for PipelineTokenizer (.mtok format).
// Loads test_cases.jsonl from each model directory, encodes input,
// and compares with expected ids_raw / decoded_full.
//
// Usage: test_tokenizer [models_path] [model_filter]
//
#include <fstream>
#include <iostream>
#include <sstream>
#include <string>
#include <vector>
#include <cstring>
#include <dirent.h>
#include <sys/stat.h>
#include "tokenizer/tokenizer.hpp"
#include "ujson.hpp"
using namespace MNN::Transformer;
// ANSI colors
static const char* RED = "\033[31m";
static const char* GREEN = "\033[32m";
static const char* YELLOW = "\033[33m";
static const char* RESET = "\033[0m";
static std::vector<std::string> list_dirs(const std::string& path) {
std::vector<std::string> dirs;
DIR* d = opendir(path.c_str());
if (!d) return dirs;
struct dirent* entry;
while ((entry = readdir(d)) != nullptr) {
if (entry->d_name[0] == '.') continue;
std::string full = path + "/" + entry->d_name;
struct stat st;
if (stat(full.c_str(), &st) == 0 && S_ISDIR(st.st_mode)) {
dirs.push_back(entry->d_name);
}
}
closedir(d);
std::sort(dirs.begin(), dirs.end());
return dirs;
}
static std::string ids_to_str(const std::vector<int>& ids) {
std::string s = "[";
for (size_t i = 0; i < ids.size(); i++) {
if (i) s += ",";
s += std::to_string(ids[i]);
}
s += "]";
return s;
}
struct TestResult {
std::string model;
int total;
int passed;
int failed;
int skipped;
};
static TestResult test_model(const std::string& models_path, const std::string& model_name) {
TestResult result = {model_name, 0, 0, 0, 0};
std::string model_dir = models_path + "/" + model_name;
// Find tokenizer file (.mtok preferred, then .txt)
std::string tok_path = model_dir + "/tokenizer.mtok";
{
struct stat st;
if (stat(tok_path.c_str(), &st) != 0) {
tok_path = model_dir + "/tokenizer.txt";
if (stat(tok_path.c_str(), &st) != 0) {
std::cerr << YELLOW << " SKIP: no tokenizer file found" << RESET << std::endl;
return result;
}
}
}
// Load tokenizer
Tokenizer* tok = Tokenizer::createTokenizer(tok_path);
if (!tok) {
std::cerr << YELLOW << " SKIP: failed to load tokenizer" << RESET << std::endl;
return result;
}
// Read test_cases.jsonl
std::string jsonl_path = model_dir + "/test_cases.jsonl";
std::ifstream jsonl_file(jsonl_path);
if (!jsonl_file.is_open()) {
std::cerr << YELLOW << " SKIP: no test_cases.jsonl" << RESET << std::endl;
delete tok;
return result;
}
std::string line;
int line_no = 0;
while (std::getline(jsonl_file, line)) {
line_no++;
if (line.empty()) continue;
ujson::json j = ujson::json::parse(line);
if (j.is_null()) {
result.skipped++;
continue;
}
std::string type = j.value("type", "");
if (type == "chat") {
// Chat template test
if (tok->chat_template().empty()) {
result.skipped++;
continue;
}
std::string name = j.value("name", "");
bool add_gen = j.value("add_generation_prompt", true);
std::string expected_text = j.value("formatted_text", "");
// Parse messages — use extra_json for messages with complex fields
ChatMessages messages;
if (j.contains("messages") && j["messages"].is_array()) {
size_t n = j["messages"].size();
for (size_t i = 0; i < n; i++) {
auto msg = j["messages"][i];
std::string role = msg.value("role", "");
std::string content = msg.value("content", "");
bool has_extra = msg.contains("tool_calls") ||
msg.contains("reasoning_content");
if (has_extra) {
// "json" role signals second is full JSON message object
messages.push_back({"json", msg.dump()});
} else {
messages.push_back({role, content});
}
}
}
result.total++;
std::string actual_text = tok->apply_chat_template(messages, add_gen);
// Normalize dynamic dates: strftime_now generates current date which varies per run.
// Replace "Today Date: <actual_date>" with the expected date for stable comparison.
{
const std::string prefix = "Today Date: ";
auto epos = expected_text.find(prefix);
auto apos = actual_text.find(prefix);
if (epos != std::string::npos && apos != std::string::npos) {
auto eend = expected_text.find('\n', epos);
auto aend = actual_text.find('\n', apos);
if (eend != std::string::npos && aend != std::string::npos) {
actual_text.replace(apos + prefix.size(), aend - apos - prefix.size(),
expected_text.substr(epos + prefix.size(), eend - epos - prefix.size()));
}
}
}
if (actual_text == expected_text) {
result.passed++;
} else {
result.failed++;
std::cerr << RED << " FAIL chat \"" << name << "\" line " << line_no << RESET << std::endl;
// Show first difference
size_t diff_pos = 0;
while (diff_pos < actual_text.size() && diff_pos < expected_text.size() &&
actual_text[diff_pos] == expected_text[diff_pos]) diff_pos++;
std::cerr << " diff at pos " << diff_pos << std::endl;
size_t ctx_start = diff_pos > 20 ? diff_pos - 20 : 0;
size_t ctx_end_e = std::min(diff_pos + 40, expected_text.size());
size_t ctx_end_a = std::min(diff_pos + 40, actual_text.size());
std::cerr << " expected: \"" << expected_text.substr(ctx_start, ctx_end_e - ctx_start) << "\"" << std::endl;
std::cerr << " actual: \"" << actual_text.substr(ctx_start, ctx_end_a - ctx_start) << "\"" << std::endl;
}
continue;
}
if (type != "basic") {
result.skipped++;
continue;
}
result.total++;
std::string input = j.value("input", "");
// Get expected ids_raw
std::vector<int> expected_ids;
if (j.contains("ids_raw") && j["ids_raw"].is_array()) {
size_t n = j["ids_raw"].size();
for (size_t i = 0; i < n; i++) {
expected_ids.push_back(j["ids_raw"][i].get<int>());
}
}
// Get expected decoded_full
std::string expected_decoded = j.value("decoded_full", "");
// Encode (without special tokens — encode() adds prefix_tokens automatically,
// so we need to strip them for comparison with ids_raw)
std::vector<int> actual_ids = tok->encode(input);
// Strip prefix tokens if present
// The prefix tokens are added by Tokenizer::encode() wrapper.
// For models with BOS/prefix tokens, actual_ids will have them but ids_raw won't.
// We detect and strip by checking if actual_ids starts with tokens not in ids_raw.
// Simpler approach: compare with ids_full first, then ids_raw.
bool encode_pass = false;
if (actual_ids == expected_ids) {
encode_pass = true;
} else {
// Try comparing with ids_full (which may include BOS/EOS)
std::vector<int> expected_full;
if (j.contains("ids_full") && j["ids_full"].is_array()) {
size_t n = j["ids_full"].size();
for (size_t i = 0; i < n; i++) {
expected_full.push_back(j["ids_full"][i].get<int>());
}
}
if (!expected_full.empty() && actual_ids == expected_full) {
encode_pass = true;
}
}
// Decode: use full-sequence decode (uses decoder chain for proper spacing)
std::string actual_decoded = tok->decode(expected_ids);
// Full-sequence decode may have an extra leading space (from ▁/Metaspace)
// that HuggingFace strips. Allow this difference.
bool decode_pass = (actual_decoded == expected_decoded);
if (!decode_pass && !actual_decoded.empty() && actual_decoded[0] == ' ' &&
actual_decoded.substr(1) == expected_decoded) {
decode_pass = true;
}
if (encode_pass && decode_pass) {
result.passed++;
} else {
result.failed++;
std::cerr << RED << " FAIL line " << line_no << ": \"" << input << "\"" << RESET << std::endl;
if (!encode_pass) {
std::cerr << " encode expected: " << ids_to_str(expected_ids) << std::endl;
std::cerr << " encode actual: " << ids_to_str(actual_ids) << std::endl;
}
if (!decode_pass) {
std::cerr << " decode expected: \"" << expected_decoded << "\"" << std::endl;
std::cerr << " decode actual: \"" << actual_decoded << "\"" << std::endl;
}
}
}
delete tok;
return result;
}
int main(int argc, char* argv[]) {
std::string models_path = "./models";
std::string filter = "";
if (argc > 1) models_path = argv[1];
if (argc > 2) filter = argv[2];
auto models = list_dirs(models_path);
if (models.empty()) {
std::cerr << "No model directories found in: " << models_path << std::endl;
return 1;
}
std::vector<TestResult> results;
int total_models = 0, total_passed = 0, total_failed = 0, total_skipped = 0;
std::vector<std::string> failed_models;
for (const auto& model : models) {
if (!filter.empty() && model.find(filter) == std::string::npos) continue;
std::cout << "Testing: " << model << " ..." << std::endl;
auto result = test_model(models_path, model);
results.push_back(result);
if (result.total > 0) {
total_models++;
total_passed += result.passed;
total_failed += result.failed;
total_skipped += result.skipped;
const char* color = (result.failed == 0) ? GREEN : RED;
std::cout << color << " " << result.passed << "/" << result.total
<< " passed" << RESET;
if (result.failed > 0) {
std::cout << RED << ", " << result.failed << " failed" << RESET;
failed_models.push_back(model);
}
if (result.skipped > 0) {
std::cout << YELLOW << ", " << result.skipped << " skipped" << RESET;
}
std::cout << std::endl;
}
}
// Summary
std::cout << "\n========================================" << std::endl;
std::cout << "Summary: " << total_models << " models tested" << std::endl;
std::cout << GREEN << " Passed: " << total_passed << RESET << std::endl;
if (total_failed > 0) {
std::cout << RED << " Failed: " << total_failed << RESET << std::endl;
}
if (total_skipped > 0) {
std::cout << YELLOW << " Skipped: " << total_skipped << RESET << std::endl;
}
if (!failed_models.empty()) {
std::cout << RED << "\nFailed models:" << RESET << std::endl;
for (const auto& m : failed_models) {
std::cout << RED << " - " << m << RESET << std::endl;
}
}
std::cout << "========================================" << std::endl;
return total_failed > 0 ? 1 : 0;
}