383 lines
9.8 KiB
C
383 lines
9.8 KiB
C
/* SPDX-License-Identifier: AGPL-3.0-or-later */
|
|
/* Copyright (C) 2026 David (Corey) Herrera */
|
|
#include "../../lib/test.h"
|
|
#include "../tokenizer.h"
|
|
|
|
#include <string.h>
|
|
|
|
/* ============================================================= */
|
|
/* Creation Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_create_tokenizer(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ASSERT_NOT_NULL(tok);
|
|
ASSERT_INT_EQ(0, ft_vocab_size(tok));
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_add_tokens(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
|
|
int id1 = ft_add_token(tok, "hello");
|
|
int id2 = ft_add_token(tok, "world");
|
|
|
|
ASSERT_INT_EQ(0, id1);
|
|
ASSERT_INT_EQ(1, id2);
|
|
ASSERT_INT_EQ(2, ft_vocab_size(tok));
|
|
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_duplicate_token_returns_same_id(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
|
|
int id1 = ft_add_token(tok, "hello");
|
|
int id2 = ft_add_token(tok, "hello");
|
|
|
|
ASSERT_INT_EQ(id1, id2);
|
|
ASSERT_INT_EQ(1, ft_vocab_size(tok));
|
|
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Token Lookup Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_id_to_token(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "hello");
|
|
ft_add_token(tok, "world");
|
|
|
|
ASSERT_STR_EQ("hello", ft_id_to_token(tok, 0));
|
|
ASSERT_STR_EQ("world", ft_id_to_token(tok, 1));
|
|
ASSERT_NULL(ft_id_to_token(tok, 999));
|
|
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_token_to_id(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "hello");
|
|
ft_add_token(tok, "world");
|
|
|
|
ASSERT_INT_EQ(0, ft_token_to_id(tok, "hello"));
|
|
ASSERT_INT_EQ(1, ft_token_to_id(tok, "world"));
|
|
ASSERT_INT_EQ(-1, ft_token_to_id(tok, "unknown"));
|
|
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Special Token Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_special_tokens(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
|
|
/* Initially all -1 */
|
|
ASSERT_INT_EQ(-1, ft_unk_id(tok));
|
|
ASSERT_INT_EQ(-1, ft_bos_id(tok));
|
|
ASSERT_INT_EQ(-1, ft_eos_id(tok));
|
|
ASSERT_INT_EQ(-1, ft_pad_id(tok));
|
|
|
|
/* Set special tokens */
|
|
ft_set_unk_token(tok, "<unk>");
|
|
ft_set_bos_token(tok, "<s>");
|
|
ft_set_eos_token(tok, "</s>");
|
|
ft_set_pad_token(tok, "<pad>");
|
|
|
|
ASSERT_TRUE(ft_unk_id(tok) >= 0);
|
|
ASSERT_TRUE(ft_bos_id(tok) >= 0);
|
|
ASSERT_TRUE(ft_eos_id(tok) >= 0);
|
|
ASSERT_TRUE(ft_pad_id(tok) >= 0);
|
|
|
|
/* Verify tokens are in vocab */
|
|
ASSERT_STR_EQ("<unk>", ft_id_to_token(tok, ft_unk_id(tok)));
|
|
ASSERT_STR_EQ("<s>", ft_id_to_token(tok, ft_bos_id(tok)));
|
|
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_unk_for_unknown_tokens(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "hello");
|
|
ft_set_unk_token(tok, "<unk>");
|
|
|
|
/* Known token returns its ID */
|
|
ASSERT_INT_EQ(0, ft_token_to_id(tok, "hello"));
|
|
|
|
/* Unknown token returns UNK ID */
|
|
ASSERT_INT_EQ(ft_unk_id(tok), ft_token_to_id(tok, "unknown"));
|
|
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Encoding Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_encode_simple(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "h");
|
|
ft_add_token(tok, "e");
|
|
ft_add_token(tok, "l");
|
|
ft_add_token(tok, "o");
|
|
|
|
int length;
|
|
int* ids = ft_encode(tok, "hello", &length);
|
|
|
|
ASSERT_NOT_NULL(ids);
|
|
ASSERT_INT_EQ(5, length); /* h e l l o */
|
|
ASSERT_INT_EQ(0, ids[0]); /* h */
|
|
ASSERT_INT_EQ(1, ids[1]); /* e */
|
|
ASSERT_INT_EQ(2, ids[2]); /* l */
|
|
ASSERT_INT_EQ(2, ids[3]); /* l */
|
|
ASSERT_INT_EQ(3, ids[4]); /* o */
|
|
|
|
free(ids);
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_encode_empty(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
|
|
int length;
|
|
int* ids = ft_encode(tok, "", &length);
|
|
|
|
ASSERT_NOT_NULL(ids);
|
|
ASSERT_INT_EQ(0, length);
|
|
|
|
free(ids);
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_encode_with_bpe_merge(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "h");
|
|
ft_add_token(tok, "e");
|
|
ft_add_token(tok, "l");
|
|
ft_add_token(tok, "o");
|
|
ft_add_token(tok, "he"); /* Merged token */
|
|
ft_add_token(tok, "ll"); /* Merged token */
|
|
|
|
/* Add BPE merges */
|
|
ft_add_merge(tok, "h", "e");
|
|
ft_add_merge(tok, "l", "l");
|
|
|
|
int length;
|
|
int* ids = ft_encode(tok, "hello", &length);
|
|
|
|
ASSERT_NOT_NULL(ids);
|
|
ASSERT_INT_EQ(3, length); /* "he" "ll" "o" */
|
|
ASSERT_INT_EQ(4, ids[0]); /* he */
|
|
ASSERT_INT_EQ(5, ids[1]); /* ll */
|
|
ASSERT_INT_EQ(3, ids[2]); /* o */
|
|
|
|
free(ids);
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_encode_with_special_tokens(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "h");
|
|
ft_add_token(tok, "i");
|
|
ft_set_bos_token(tok, "<s>");
|
|
ft_set_eos_token(tok, "</s>");
|
|
|
|
int length;
|
|
int* ids = ft_encode_with_special(tok, "hi", &length, true, true);
|
|
|
|
ASSERT_NOT_NULL(ids);
|
|
ASSERT_INT_EQ(4, length); /* <s> h i </s> */
|
|
ASSERT_INT_EQ(ft_bos_id(tok), ids[0]);
|
|
ASSERT_INT_EQ(0, ids[1]); /* h */
|
|
ASSERT_INT_EQ(1, ids[2]); /* i */
|
|
ASSERT_INT_EQ(ft_eos_id(tok), ids[3]);
|
|
|
|
free(ids);
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Decoding Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_decode_simple(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "hel");
|
|
ft_add_token(tok, "lo");
|
|
|
|
int ids[] = {0, 1};
|
|
char* text = ft_decode(tok, ids, 2);
|
|
|
|
ASSERT_NOT_NULL(text);
|
|
ASSERT_STR_EQ("hello", text);
|
|
|
|
free(text);
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_decode_empty(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
|
|
char* text = ft_decode(tok, NULL, 0);
|
|
|
|
ASSERT_NOT_NULL(text);
|
|
ASSERT_STR_EQ("", text);
|
|
|
|
free(text);
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_decode_skip_special(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "hello");
|
|
ft_set_bos_token(tok, "<s>");
|
|
ft_set_eos_token(tok, "</s>");
|
|
|
|
int ids[] = {ft_bos_id(tok), 0, ft_eos_id(tok)};
|
|
char* text = ft_decode_skip_special(tok, ids, 3);
|
|
|
|
ASSERT_NOT_NULL(text);
|
|
ASSERT_STR_EQ("hello", text);
|
|
|
|
free(text);
|
|
ft_free(tok);
|
|
}
|
|
|
|
static void test_roundtrip(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "h");
|
|
ft_add_token(tok, "e");
|
|
ft_add_token(tok, "l");
|
|
ft_add_token(tok, "o");
|
|
ft_add_token(tok, " ");
|
|
|
|
const char* original = "hello";
|
|
|
|
int length;
|
|
int* ids = ft_encode(tok, original, &length);
|
|
char* decoded = ft_decode(tok, ids, length);
|
|
|
|
ASSERT_STR_EQ(original, decoded);
|
|
|
|
free(ids);
|
|
free(decoded);
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Batch Encoding Tests */
|
|
/* ============================================================= */
|
|
|
|
static void test_encode_batch(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
ft_add_token(tok, "a");
|
|
ft_add_token(tok, "b");
|
|
ft_add_token(tok, "c");
|
|
|
|
const char* texts[] = {"ab", "abc"};
|
|
int* lengths;
|
|
int** ids = ft_encode_batch(tok, texts, 2, &lengths);
|
|
|
|
ASSERT_NOT_NULL(ids);
|
|
ASSERT_NOT_NULL(lengths);
|
|
|
|
ASSERT_INT_EQ(2, lengths[0]);
|
|
ASSERT_INT_EQ(3, lengths[1]);
|
|
|
|
ASSERT_INT_EQ(0, ids[0][0]); /* a */
|
|
ASSERT_INT_EQ(1, ids[0][1]); /* b */
|
|
|
|
ASSERT_INT_EQ(0, ids[1][0]); /* a */
|
|
ASSERT_INT_EQ(1, ids[1][1]); /* b */
|
|
ASSERT_INT_EQ(2, ids[1][2]); /* c */
|
|
|
|
free(ids[0]);
|
|
free(ids[1]);
|
|
free(ids);
|
|
free(lengths);
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Edge Cases */
|
|
/* ============================================================= */
|
|
|
|
static void test_null_safety(void) {
|
|
ASSERT_INT_EQ(0, ft_vocab_size(NULL));
|
|
ASSERT_NULL(ft_id_to_token(NULL, 0));
|
|
ASSERT_INT_EQ(-1, ft_token_to_id(NULL, "test"));
|
|
ASSERT_INT_EQ(-1, ft_unk_id(NULL));
|
|
|
|
ft_free(NULL); /* Should not crash */
|
|
ASSERT_TRUE(1);
|
|
}
|
|
|
|
static void test_unicode_chars(void) {
|
|
ft_tokenizer* tok = ft_new();
|
|
|
|
/* Add some unicode characters */
|
|
ft_add_token(tok, "é");
|
|
ft_add_token(tok, "ñ");
|
|
ft_add_token(tok, "日");
|
|
|
|
ASSERT_INT_EQ(3, ft_vocab_size(tok));
|
|
|
|
int length;
|
|
int* ids = ft_encode(tok, "éñ", &length);
|
|
|
|
ASSERT_INT_EQ(2, length);
|
|
ASSERT_INT_EQ(0, ids[0]); /* é */
|
|
ASSERT_INT_EQ(1, ids[1]); /* ñ */
|
|
|
|
free(ids);
|
|
ft_free(tok);
|
|
}
|
|
|
|
/* ============================================================= */
|
|
/* Main */
|
|
/* ============================================================= */
|
|
|
|
int main(void) {
|
|
printf("\n Running tokenizer tests...\n\n");
|
|
|
|
/* Creation tests */
|
|
RUN_TEST(create_tokenizer);
|
|
RUN_TEST(add_tokens);
|
|
RUN_TEST(duplicate_token_returns_same_id);
|
|
|
|
/* Lookup tests */
|
|
RUN_TEST(id_to_token);
|
|
RUN_TEST(token_to_id);
|
|
|
|
/* Special token tests */
|
|
RUN_TEST(special_tokens);
|
|
RUN_TEST(unk_for_unknown_tokens);
|
|
|
|
/* Encoding tests */
|
|
RUN_TEST(encode_simple);
|
|
RUN_TEST(encode_empty);
|
|
RUN_TEST(encode_with_bpe_merge);
|
|
RUN_TEST(encode_with_special_tokens);
|
|
|
|
/* Decoding tests */
|
|
RUN_TEST(decode_simple);
|
|
RUN_TEST(decode_empty);
|
|
RUN_TEST(decode_skip_special);
|
|
RUN_TEST(roundtrip);
|
|
|
|
/* Batch tests */
|
|
RUN_TEST(encode_batch);
|
|
|
|
/* Edge cases */
|
|
RUN_TEST(null_safety);
|
|
RUN_TEST(unicode_chars);
|
|
|
|
TEST_REPORT();
|
|
|
|
return __ft_failed > 0 ? 1 : 0;
|
|
}
|