diff --git a/include/builtins_registry.h b/include/builtins_registry.h index efa250fb..e1cd93d2 100644 --- a/include/builtins_registry.h +++ b/include/builtins_registry.h @@ -317,6 +317,15 @@ typedef struct { X("__regex_replace", 3, HML_BUILTIN_INTERNAL, HML_RET_STRING) \ X("__regex_replace_all", 3, HML_BUILTIN_INTERNAL, HML_RET_STRING) +#define HML_BUILTINS_TOKENIZER(X) \ + X("__tokenizer_create", 1, HML_BUILTIN_INTERNAL, HML_RET_PTR) \ + X("__tokenizer_free", 1, HML_BUILTIN_INTERNAL, HML_RET_NULL) \ + X("__tokenizer_encode", 2, HML_BUILTIN_INTERNAL, HML_RET_ARRAY) \ + X("__tokenizer_decode", 2, HML_BUILTIN_INTERNAL, HML_RET_STRING) \ + X("__tokenizer_count", 2, HML_BUILTIN_INTERNAL, HML_RET_I32) \ + X("__tokenizer_vocab_size", 1, HML_BUILTIN_INTERNAL | HML_BUILTIN_PURE, HML_RET_I32) \ + X("__tokenizer_add_special", 3, HML_BUILTIN_INTERNAL, HML_RET_NULL) + #define HML_BUILTINS_FFI(X) \ X("callback", 2, HML_BUILTIN_PUBLIC, HML_RET_PTR) \ X("callback_free", 1, HML_BUILTIN_PUBLIC, HML_RET_NULL) \ @@ -349,6 +358,7 @@ typedef struct { HML_BUILTINS_CRYPTO(X) \ HML_BUILTINS_COMPRESSION(X) \ HML_BUILTINS_REGEX(X) \ + HML_BUILTINS_TOKENIZER(X) \ HML_BUILTINS_FFI(X) \ HML_BUILTINS_INTERNAL(X) diff --git a/include/hemlock_limits.h b/include/hemlock_limits.h index 6d4f7338..fc73aeb6 100644 --- a/include/hemlock_limits.h +++ b/include/hemlock_limits.h @@ -185,6 +185,24 @@ #define HML_ATOMIC_I32_ALIGNMENT 4 // _Alignof(_Atomic int32_t) #define HML_ATOMIC_I64_ALIGNMENT 8 // _Alignof(_Atomic int64_t) +// ========== TOKENIZER (BPE) CONSTANTS ========== + +// Initial hash table capacity for BPE vocabulary +// Should be a prime number for better distribution +#define HML_BPE_INITIAL_HASH_CAPACITY 131071 + +// Maximum token byte sequence length +#define HML_BPE_MAX_TOKEN_LENGTH 128 + +// Initial capacity for the BPE merge work list +#define HML_BPE_INITIAL_WORK_CAPACITY 256 + +// Base64 decode buffer size for vocabulary loading +#define HML_BPE_BASE64_DECODE_BUFSIZE 256 + +// Maximum line length when reading vocabulary files +#define HML_BPE_MAX_VOCAB_LINE_LENGTH 1024 + // ========== INLINE CACHE CONSTANTS ========== // Inline caching is used to speed up property access and method dispatch diff --git a/runtime/include/hemlock_runtime.h b/runtime/include/hemlock_runtime.h index ec233d0d..431a4f9c 100644 --- a/runtime/include/hemlock_runtime.h +++ b/runtime/include/hemlock_runtime.h @@ -939,6 +939,19 @@ HmlValue hml_builtin_regex_error(HmlClosureEnv *env, HmlValue errcode, HmlValue HmlValue hml_builtin_regex_replace(HmlClosureEnv *env, HmlValue preg, HmlValue text, HmlValue replacement); HmlValue hml_builtin_regex_replace_all(HmlClosureEnv *env, HmlValue preg, HmlValue text, HmlValue replacement); +// ========== TOKENIZER (BPE) FUNCTIONS ========== + +// Tokenizer lifecycle +HmlValue hml_tokenizer_create(HmlValue path); +HmlValue hml_tokenizer_free(HmlValue handle); + +// Tokenizer operations +HmlValue hml_tokenizer_encode(HmlValue handle, HmlValue text); +HmlValue hml_tokenizer_decode(HmlValue handle, HmlValue tokens); +HmlValue hml_tokenizer_count(HmlValue handle, HmlValue text); +HmlValue hml_tokenizer_vocab_size(HmlValue handle); +HmlValue hml_tokenizer_add_special(HmlValue handle, HmlValue token, HmlValue rank); + // ========== CALL STACK TRACKING ========== // Default maximum call stack depth (matches interpreter's limit) diff --git a/runtime/src/builtins_tokenizer.c b/runtime/src/builtins_tokenizer.c new file mode 100644 index 00000000..36bad75c --- /dev/null +++ b/runtime/src/builtins_tokenizer.c @@ -0,0 +1,675 @@ +/* + * Hemlock Runtime - BPE Tokenizer Functions + * + * Runtime library implementation of BPE (Byte Pair Encoding) tokenizer + * for compiled Hemlock programs. Supports tiktoken-compatible vocabulary files. + */ + +#include "../include/hemlock_runtime.h" +#include "../include/hemlock_value.h" + +#include +#include +#include +#include + +/* ========== CONSTANTS ========== */ + +#define BPE_INITIAL_HASH_CAPACITY 131071 +#define BPE_MAX_TOKEN_LENGTH 128 +#define BPE_INITIAL_WORK_CAPACITY 256 +#define BPE_BASE64_DECODE_BUFSIZE 256 +#define BPE_MAX_VOCAB_LINE_LENGTH 1024 +#define BPE_GROWTH_FACTOR 2 + +/* ========== BASE64 DECODER ========== */ + +static const unsigned char b64_table[256] = { + ['A'] = 0, ['B'] = 1, ['C'] = 2, ['D'] = 3, + ['E'] = 4, ['F'] = 5, ['G'] = 6, ['H'] = 7, + ['I'] = 8, ['J'] = 9, ['K'] = 10, ['L'] = 11, + ['M'] = 12, ['N'] = 13, ['O'] = 14, ['P'] = 15, + ['Q'] = 16, ['R'] = 17, ['S'] = 18, ['T'] = 19, + ['U'] = 20, ['V'] = 21, ['W'] = 22, ['X'] = 23, + ['Y'] = 24, ['Z'] = 25, + ['a'] = 26, ['b'] = 27, ['c'] = 28, ['d'] = 29, + ['e'] = 30, ['f'] = 31, ['g'] = 32, ['h'] = 33, + ['i'] = 34, ['j'] = 35, ['k'] = 36, ['l'] = 37, + ['m'] = 38, ['n'] = 39, ['o'] = 40, ['p'] = 41, + ['q'] = 42, ['r'] = 43, ['s'] = 44, ['t'] = 45, + ['u'] = 46, ['v'] = 47, ['w'] = 48, ['x'] = 49, + ['y'] = 50, ['z'] = 51, + ['0'] = 52, ['1'] = 53, ['2'] = 54, ['3'] = 55, + ['4'] = 56, ['5'] = 57, ['6'] = 58, ['7'] = 59, + ['8'] = 60, ['9'] = 61, ['+'] = 62, ['/'] = 63, +}; + +static int b64_decode(const char *src, int src_len, uint8_t *dst, int dst_cap) { + int out = 0; + int accum = 0; + int bits = 0; + for (int i = 0; i < src_len; i++) { + unsigned char c = (unsigned char)src[i]; + if (c == '=') break; + if (c == '\n' || c == '\r' || c == ' ') continue; + accum = (accum << 6) | b64_table[c]; + bits += 6; + if (bits >= 8) { + bits -= 8; + if (out < dst_cap) { + dst[out++] = (uint8_t)((accum >> bits) & 0xFF); + } + } + } + return out; +} + +/* ========== BPE DATA STRUCTURES ========== */ + +typedef struct { + uint8_t *bytes; + int byte_len; + int32_t rank; +} BpeVocabEntry; + +typedef struct { + uint8_t *key; + int key_len; + int32_t value; +} BpeHashEntry; + +typedef struct { + BpeVocabEntry *vocab; + int32_t vocab_size; + int32_t vocab_capacity; + BpeHashEntry *hash_table; + int32_t hash_capacity; + char **special_tokens; + int32_t *special_ranks; + int32_t num_special; + int32_t special_capacity; +} BpeTokenizer; + +/* ========== HASH TABLE ========== */ + +static uint32_t fnv1a_bytes(const uint8_t *data, int len) { + uint32_t hash = 2166136261u; + for (int i = 0; i < len; i++) { + hash ^= data[i]; + hash *= 16777619u; + } + return hash; +} + +static void bpe_hash_init(BpeTokenizer *tok, int32_t capacity) { + tok->hash_capacity = capacity; + tok->hash_table = calloc((size_t)capacity, sizeof(BpeHashEntry)); + if (!tok->hash_table) return; + for (int32_t i = 0; i < capacity; i++) { + tok->hash_table[i].value = -1; + } +} + +static void bpe_hash_insert(BpeTokenizer *tok, uint8_t *key, int key_len, int32_t rank) { + if (!tok->hash_table) return; + uint32_t idx = fnv1a_bytes(key, key_len) % (uint32_t)tok->hash_capacity; + for (int32_t probe = 0; probe < tok->hash_capacity; probe++) { + int32_t i = (int32_t)((idx + (uint32_t)probe) % (uint32_t)tok->hash_capacity); + if (tok->hash_table[i].value == -1) { + tok->hash_table[i].key = key; + tok->hash_table[i].key_len = key_len; + tok->hash_table[i].value = rank; + return; + } + } +} + +static int32_t bpe_hash_lookup(BpeTokenizer *tok, const uint8_t *key, int key_len) { + if (!tok->hash_table) return -1; + uint32_t idx = fnv1a_bytes(key, key_len) % (uint32_t)tok->hash_capacity; + for (int32_t probe = 0; probe < tok->hash_capacity; probe++) { + int32_t i = (int32_t)((idx + (uint32_t)probe) % (uint32_t)tok->hash_capacity); + if (tok->hash_table[i].value == -1) return -1; + if (tok->hash_table[i].key_len == key_len && + memcmp(tok->hash_table[i].key, key, (size_t)key_len) == 0) { + return tok->hash_table[i].value; + } + } + return -1; +} + +/* ========== BPE TOKENIZER LIFECYCLE ========== */ + +static void bpe_tokenizer_destroy(BpeTokenizer *tok) { + if (!tok) return; + if (tok->vocab) { + for (int32_t i = 0; i < tok->vocab_size; i++) { + free(tok->vocab[i].bytes); + } + free(tok->vocab); + } + free(tok->hash_table); + if (tok->special_tokens) { + for (int32_t i = 0; i < tok->num_special; i++) { + free(tok->special_tokens[i]); + } + free(tok->special_tokens); + } + free(tok->special_ranks); + free(tok); +} + +static BpeTokenizer *bpe_tokenizer_load(const char *path) { + FILE *fp = fopen(path, "r"); + if (!fp) return NULL; + + BpeTokenizer *tok = calloc(1, sizeof(BpeTokenizer)); + if (!tok) { fclose(fp); return NULL; } + + tok->vocab_capacity = 4096; + tok->vocab = calloc((size_t)tok->vocab_capacity, sizeof(BpeVocabEntry)); + if (!tok->vocab) { fclose(fp); bpe_tokenizer_destroy(tok); return NULL; } + + char line[BPE_MAX_VOCAB_LINE_LENGTH]; + + while (fgets(line, sizeof(line), fp)) { + if (line[0] == '\n' || line[0] == '#') continue; + char *space = strchr(line, ' '); + if (!space) continue; + *space = '\0'; + char *b64_str = line; + int b64_len = (int)(space - line); + int32_t rank = (int32_t)strtol(space + 1, NULL, 10); + + uint8_t decoded[BPE_BASE64_DECODE_BUFSIZE]; + int decoded_len = b64_decode(b64_str, b64_len, decoded, BPE_BASE64_DECODE_BUFSIZE); + if (decoded_len <= 0) continue; + + if (tok->vocab_size >= tok->vocab_capacity) { + int32_t new_cap = tok->vocab_capacity * BPE_GROWTH_FACTOR; + BpeVocabEntry *new_vocab = realloc(tok->vocab, (size_t)new_cap * sizeof(BpeVocabEntry)); + if (!new_vocab) continue; + tok->vocab = new_vocab; + tok->vocab_capacity = new_cap; + } + + uint8_t *bytes_copy = malloc((size_t)decoded_len); + if (!bytes_copy) continue; + memcpy(bytes_copy, decoded, (size_t)decoded_len); + + BpeVocabEntry *entry = &tok->vocab[tok->vocab_size]; + entry->bytes = bytes_copy; + entry->byte_len = decoded_len; + entry->rank = rank; + tok->vocab_size++; + } + fclose(fp); + + if (tok->vocab_size == 0) { bpe_tokenizer_destroy(tok); return NULL; } + + int32_t hash_cap = tok->vocab_size * 3; + if (hash_cap < BPE_INITIAL_HASH_CAPACITY) hash_cap = BPE_INITIAL_HASH_CAPACITY; + bpe_hash_init(tok, hash_cap); + + for (int32_t i = 0; i < tok->vocab_size; i++) { + bpe_hash_insert(tok, tok->vocab[i].bytes, tok->vocab[i].byte_len, tok->vocab[i].rank); + } + + return tok; +} + +/* ========== BPE ENCODE ========== */ + +typedef struct BpeNode { + uint8_t *bytes; + int byte_len; + struct BpeNode *next; + struct BpeNode *prev; +} BpeNode; + +static int bpe_encode_chunk(BpeTokenizer *tok, const uint8_t *data, int data_len, + int32_t **out_tokens, int *out_count) { + if (data_len == 0) { + *out_tokens = NULL; + *out_count = 0; + return 0; + } + + int node_capacity = data_len + 1; + BpeNode *nodes = calloc((size_t)node_capacity, sizeof(BpeNode)); + if (!nodes) return -1; + + uint8_t *data_copy = malloc((size_t)data_len); + if (!data_copy) { free(nodes); return -1; } + memcpy(data_copy, data, (size_t)data_len); + + for (int i = 0; i < data_len; i++) { + nodes[i].bytes = data_copy + i; + nodes[i].byte_len = 1; + nodes[i].prev = (i > 0) ? &nodes[i - 1] : NULL; + nodes[i].next = (i < data_len - 1) ? &nodes[i + 1] : NULL; + } + + uint8_t pair_buf[BPE_MAX_TOKEN_LENGTH * 2]; + + int changed = 1; + while (changed) { + changed = 0; + int32_t best_rank = INT32_MAX; + BpeNode *best_node = NULL; + + for (BpeNode *node = &nodes[0]; node && node->next; node = node->next) { + if (node->byte_len <= 0) continue; + BpeNode *next = node->next; + int pair_len = node->byte_len + next->byte_len; + if (pair_len > BPE_MAX_TOKEN_LENGTH * 2) continue; + + memcpy(pair_buf, node->bytes, (size_t)node->byte_len); + memcpy(pair_buf + node->byte_len, next->bytes, (size_t)next->byte_len); + + int32_t rank = bpe_hash_lookup(tok, pair_buf, pair_len); + if (rank >= 0 && rank < best_rank) { + best_rank = rank; + best_node = node; + } + } + + if (best_node) { + BpeNode *next = best_node->next; + int new_len = best_node->byte_len + next->byte_len; + + uint8_t *merged = malloc((size_t)new_len); + if (!merged) break; + memcpy(merged, best_node->bytes, (size_t)best_node->byte_len); + memcpy(merged + best_node->byte_len, next->bytes, (size_t)next->byte_len); + + if (best_node->bytes < data_copy || best_node->bytes >= data_copy + data_len) { + free(best_node->bytes); + } + + best_node->bytes = merged; + best_node->byte_len = new_len; + + best_node->next = next->next; + if (next->next) next->next->prev = best_node; + + if (next->bytes < data_copy || next->bytes >= data_copy + data_len) { + free(next->bytes); + } + next->bytes = NULL; + next->byte_len = 0; + + changed = 1; + } + } + + int count = 0; + for (BpeNode *node = &nodes[0]; node; node = node->next) { + if (node->byte_len > 0) count++; + } + + int32_t *tokens = malloc((size_t)count * sizeof(int32_t)); + if (!tokens) { + for (int i = 0; i < node_capacity; i++) { + if (nodes[i].bytes && (nodes[i].bytes < data_copy || nodes[i].bytes >= data_copy + data_len)) { + free(nodes[i].bytes); + } + } + free(data_copy); + free(nodes); + return -1; + } + + int idx = 0; + for (BpeNode *node = &nodes[0]; node; node = node->next) { + if (node->byte_len <= 0) continue; + int32_t rank = bpe_hash_lookup(tok, node->bytes, node->byte_len); + if (rank >= 0) { + tokens[idx++] = rank; + } else { + for (int b = 0; b < node->byte_len; b++) { + int32_t byte_rank = bpe_hash_lookup(tok, &node->bytes[b], 1); + if (byte_rank >= 0) { + if (idx >= count) { + count += node->byte_len; + int32_t *new_tokens = realloc(tokens, (size_t)count * sizeof(int32_t)); + if (!new_tokens) { free(tokens); tokens = NULL; break; } + tokens = new_tokens; + } + tokens[idx++] = byte_rank; + } + } + if (!tokens) break; + } + } + + for (int i = 0; i < node_capacity; i++) { + if (nodes[i].bytes && (nodes[i].bytes < data_copy || nodes[i].bytes >= data_copy + data_len)) { + int already_freed = 0; + for (int j = 0; j < i; j++) { + if (nodes[j].bytes == nodes[i].bytes) { already_freed = 1; break; } + } + if (!already_freed) free(nodes[i].bytes); + } + } + free(data_copy); + free(nodes); + + *out_tokens = tokens; + *out_count = idx; + return 0; +} + +/* ========== PRE-TOKENIZER ========== */ + +typedef struct { + const uint8_t *start; + int len; +} TextChunk; + +static int is_letter_or_digit(uint8_t c) { + return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c >= 128; +} + +static int pretokenize(const uint8_t *text, int text_len, + TextChunk **out_chunks, int *out_count) { + if (text_len == 0) { + *out_chunks = NULL; + *out_count = 0; + return 0; + } + + int capacity = BPE_INITIAL_WORK_CAPACITY; + TextChunk *chunks = malloc((size_t)capacity * sizeof(TextChunk)); + if (!chunks) return -1; + int count = 0; + + int i = 0; + while (i < text_len) { + int start = i; + while (i < text_len && (text[i] == ' ' || text[i] == '\t')) i++; + if (i < text_len) { + if (text[i] == '\n' || text[i] == '\r') { + i++; + if (i < text_len && text[i - 1] == '\r' && text[i] == '\n') i++; + } else if (is_letter_or_digit(text[i])) { + while (i < text_len && is_letter_or_digit(text[i])) i++; + } else { + i++; + } + } + if (i > start) { + if (count >= capacity) { + capacity *= BPE_GROWTH_FACTOR; + TextChunk *new_chunks = realloc(chunks, (size_t)capacity * sizeof(TextChunk)); + if (!new_chunks) { free(chunks); return -1; } + chunks = new_chunks; + } + chunks[count].start = text + start; + chunks[count].len = i - start; + count++; + } + } + + *out_chunks = chunks; + *out_count = count; + return 0; +} + +static int bpe_encode(BpeTokenizer *tok, const char *text, int text_len, + int32_t **out_tokens, int *out_count) { + TextChunk *chunks = NULL; + int num_chunks = 0; + if (pretokenize((const uint8_t *)text, text_len, &chunks, &num_chunks) != 0) { + return -1; + } + + int total_capacity = text_len + 16; + int32_t *all_tokens = malloc((size_t)total_capacity * sizeof(int32_t)); + if (!all_tokens) { free(chunks); return -1; } + int total_count = 0; + + for (int c = 0; c < num_chunks; c++) { + int32_t *chunk_tokens = NULL; + int chunk_count = 0; + + if (bpe_encode_chunk(tok, chunks[c].start, chunks[c].len, + &chunk_tokens, &chunk_count) != 0) { + free(all_tokens); + free(chunks); + return -1; + } + + while (total_count + chunk_count > total_capacity) { + total_capacity *= BPE_GROWTH_FACTOR; + int32_t *new_all = realloc(all_tokens, (size_t)total_capacity * sizeof(int32_t)); + if (!new_all) { free(chunk_tokens); free(all_tokens); free(chunks); return -1; } + all_tokens = new_all; + } + + if (chunk_tokens) { + memcpy(all_tokens + total_count, chunk_tokens, (size_t)chunk_count * sizeof(int32_t)); + total_count += chunk_count; + free(chunk_tokens); + } + } + + free(chunks); + *out_tokens = all_tokens; + *out_count = total_count; + return 0; +} + +/* ========== BPE DECODE ========== */ + +static char *bpe_decode(BpeTokenizer *tok, const int32_t *tokens, int num_tokens, int *out_len) { + int total_len = 0; + for (int i = 0; i < num_tokens; i++) { + for (int32_t v = 0; v < tok->vocab_size; v++) { + if (tok->vocab[v].rank == tokens[i]) { + total_len += tok->vocab[v].byte_len; + break; + } + } + } + + char *result = malloc((size_t)(total_len + 1)); + if (!result) return NULL; + + int offset = 0; + for (int i = 0; i < num_tokens; i++) { + for (int32_t v = 0; v < tok->vocab_size; v++) { + if (tok->vocab[v].rank == tokens[i]) { + memcpy(result + offset, tok->vocab[v].bytes, (size_t)tok->vocab[v].byte_len); + offset += tok->vocab[v].byte_len; + break; + } + } + } + result[offset] = '\0'; + *out_len = offset; + return result; +} + +/* ========== RUNTIME API FUNCTIONS ========== */ + +HmlValue hml_tokenizer_create(HmlValue path) { + if (path.type != HML_VAL_STRING || !path.as.as_string) { + hml_runtime_error("tokenizer_create() requires string path argument"); + } + BpeTokenizer *tok = bpe_tokenizer_load(path.as.as_string->data); + if (!tok) { + hml_runtime_error("tokenizer_create() failed to load vocabulary from: %s", path.as.as_string->data); + } + return hml_val_ptr(tok); +} + +HmlValue hml_tokenizer_free(HmlValue handle) { + if (handle.type != HML_VAL_PTR) { + hml_runtime_error("tokenizer_free() requires ptr argument"); + } + bpe_tokenizer_destroy((BpeTokenizer *)handle.as.as_ptr); + return hml_val_null(); +} + +HmlValue hml_tokenizer_encode(HmlValue handle, HmlValue text) { + if (handle.type != HML_VAL_PTR || !handle.as.as_ptr) { + hml_runtime_error("tokenizer_encode() requires valid tokenizer ptr"); + } + if (text.type != HML_VAL_STRING || !text.as.as_string) { + hml_runtime_error("tokenizer_encode() requires string text argument"); + } + + BpeTokenizer *tok = (BpeTokenizer *)handle.as.as_ptr; + const char *str = text.as.as_string->data; + int str_len = text.as.as_string->length; + + int32_t *tokens = NULL; + int num_tokens = 0; + if (bpe_encode(tok, str, str_len, &tokens, &num_tokens) != 0) { + hml_runtime_error("tokenizer_encode() encoding failed"); + } + + HmlValue result = hml_val_array(); + for (int i = 0; i < num_tokens; i++) { + HmlValue v; + v.type = HML_VAL_I32; + v.as.as_i32 = tokens[i]; + hml_array_push(result, v); + } + free(tokens); + return result; +} + +HmlValue hml_tokenizer_decode(HmlValue handle, HmlValue tokens_arr) { + if (handle.type != HML_VAL_PTR || !handle.as.as_ptr) { + hml_runtime_error("tokenizer_decode() requires valid tokenizer ptr"); + } + if (tokens_arr.type != HML_VAL_ARRAY || !tokens_arr.as.as_array) { + hml_runtime_error("tokenizer_decode() requires array argument"); + } + + BpeTokenizer *tok = (BpeTokenizer *)handle.as.as_ptr; + HmlArray *arr = tokens_arr.as.as_array; + int count = arr->length; + + int32_t *token_ids = malloc((size_t)count * sizeof(int32_t)); + if (!token_ids && count > 0) { + hml_runtime_error("tokenizer_decode() out of memory"); + } + + for (int i = 0; i < count; i++) { + token_ids[i] = (int32_t)hml_to_i64(arr->elements[i]); + } + + int result_len = 0; + char *decoded = bpe_decode(tok, token_ids, count, &result_len); + free(token_ids); + + if (!decoded) { + hml_runtime_error("tokenizer_decode() decode failed"); + } + + HmlValue result = hml_val_string(decoded); + free(decoded); + return result; +} + +HmlValue hml_tokenizer_count(HmlValue handle, HmlValue text) { + if (handle.type != HML_VAL_PTR || !handle.as.as_ptr) { + hml_runtime_error("tokenizer_count() requires valid tokenizer ptr"); + } + if (text.type != HML_VAL_STRING || !text.as.as_string) { + hml_runtime_error("tokenizer_count() requires string text argument"); + } + + BpeTokenizer *tok = (BpeTokenizer *)handle.as.as_ptr; + const char *str = text.as.as_string->data; + int str_len = text.as.as_string->length; + + int32_t *tokens = NULL; + int num_tokens = 0; + if (bpe_encode(tok, str, str_len, &tokens, &num_tokens) != 0) { + hml_runtime_error("tokenizer_count() encoding failed"); + } + free(tokens); + + HmlValue result; + result.type = HML_VAL_I32; + result.as.as_i32 = num_tokens; + return result; +} + +HmlValue hml_tokenizer_vocab_size(HmlValue handle) { + if (handle.type != HML_VAL_PTR || !handle.as.as_ptr) { + HmlValue result; + result.type = HML_VAL_I32; + result.as.as_i32 = 0; + return result; + } + BpeTokenizer *tok = (BpeTokenizer *)handle.as.as_ptr; + HmlValue result; + result.type = HML_VAL_I32; + result.as.as_i32 = tok->vocab_size; + return result; +} + +HmlValue hml_tokenizer_add_special(HmlValue handle, HmlValue token, HmlValue rank) { + if (handle.type != HML_VAL_PTR || !handle.as.as_ptr) { + hml_runtime_error("tokenizer_add_special() requires valid tokenizer ptr"); + } + if (token.type != HML_VAL_STRING || !token.as.as_string) { + hml_runtime_error("tokenizer_add_special() requires string token argument"); + } + + BpeTokenizer *tok = (BpeTokenizer *)handle.as.as_ptr; + const char *token_str = token.as.as_string->data; + int32_t rank_val = (int32_t)hml_to_i64(rank); + int token_len = token.as.as_string->length; + + if (tok->num_special >= tok->special_capacity) { + int32_t new_cap = tok->special_capacity == 0 ? 16 : tok->special_capacity * BPE_GROWTH_FACTOR; + char **new_tokens = realloc(tok->special_tokens, (size_t)new_cap * sizeof(char *)); + int32_t *new_ranks = realloc(tok->special_ranks, (size_t)new_cap * sizeof(int32_t)); + if (!new_tokens || !new_ranks) { + free(new_tokens); + free(new_ranks); + hml_runtime_error("tokenizer_add_special() out of memory"); + } + tok->special_tokens = new_tokens; + tok->special_ranks = new_ranks; + tok->special_capacity = new_cap; + } + + tok->special_tokens[tok->num_special] = strdup(token_str); + tok->special_ranks[tok->num_special] = rank_val; + tok->num_special++; + + uint8_t *bytes_copy = malloc((size_t)token_len); + if (bytes_copy) { + memcpy(bytes_copy, token_str, (size_t)token_len); + + if (tok->vocab_size >= tok->vocab_capacity) { + int32_t new_cap = tok->vocab_capacity * BPE_GROWTH_FACTOR; + BpeVocabEntry *new_vocab = realloc(tok->vocab, (size_t)new_cap * sizeof(BpeVocabEntry)); + if (new_vocab) { + tok->vocab = new_vocab; + tok->vocab_capacity = new_cap; + } + } + + if (tok->vocab_size < tok->vocab_capacity) { + BpeVocabEntry *entry = &tok->vocab[tok->vocab_size]; + entry->bytes = bytes_copy; + entry->byte_len = token_len; + entry->rank = rank_val; + tok->vocab_size++; + bpe_hash_insert(tok, bytes_copy, token_len, rank_val); + } else { + free(bytes_copy); + } + } + + return hml_val_null(); +} diff --git a/src/backends/compiler/codegen_call.c b/src/backends/compiler/codegen_call.c index 70469f71..e294ec0c 100644 --- a/src/backends/compiler/codegen_call.c +++ b/src/backends/compiler/codegen_call.c @@ -2544,6 +2544,86 @@ char* codegen_expr_call(CodegenContext *ctx, Expr *expr, char *result) { return result; } + // ========== TOKENIZER (BPE) BUILTINS ========== + + // __tokenizer_create(path) - load BPE vocabulary file, return tokenizer ptr + if (strcmp(fn_name, "__tokenizer_create") == 0 && expr->as.call.num_args == 1) { + char *path = codegen_expr(ctx, expr->as.call.args[0]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_create(%s);", result, path); + codegen_writeln(ctx, "hml_release(&%s);", path); + free(path); + return result; + } + + // __tokenizer_free(tok) - free tokenizer resources + if (strcmp(fn_name, "__tokenizer_free") == 0 && expr->as.call.num_args == 1) { + char *handle = codegen_expr(ctx, expr->as.call.args[0]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_free(%s);", result, handle); + codegen_writeln(ctx, "hml_release(&%s);", handle); + free(handle); + return result; + } + + // __tokenizer_encode(tok, text) - encode text to token IDs + if (strcmp(fn_name, "__tokenizer_encode") == 0 && expr->as.call.num_args == 2) { + char *handle = codegen_expr(ctx, expr->as.call.args[0]); + char *text = codegen_expr(ctx, expr->as.call.args[1]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_encode(%s, %s);", result, handle, text); + codegen_writeln(ctx, "hml_release(&%s);", handle); + codegen_writeln(ctx, "hml_release(&%s);", text); + free(handle); + free(text); + return result; + } + + // __tokenizer_decode(tok, tokens) - decode token IDs to text + if (strcmp(fn_name, "__tokenizer_decode") == 0 && expr->as.call.num_args == 2) { + char *handle = codegen_expr(ctx, expr->as.call.args[0]); + char *tokens = codegen_expr(ctx, expr->as.call.args[1]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_decode(%s, %s);", result, handle, tokens); + codegen_writeln(ctx, "hml_release(&%s);", handle); + codegen_writeln(ctx, "hml_release(&%s);", tokens); + free(handle); + free(tokens); + return result; + } + + // __tokenizer_count(tok, text) - count tokens in text + if (strcmp(fn_name, "__tokenizer_count") == 0 && expr->as.call.num_args == 2) { + char *handle = codegen_expr(ctx, expr->as.call.args[0]); + char *text = codegen_expr(ctx, expr->as.call.args[1]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_count(%s, %s);", result, handle, text); + codegen_writeln(ctx, "hml_release(&%s);", handle); + codegen_writeln(ctx, "hml_release(&%s);", text); + free(handle); + free(text); + return result; + } + + // __tokenizer_vocab_size(tok) - get vocabulary size + if (strcmp(fn_name, "__tokenizer_vocab_size") == 0 && expr->as.call.num_args == 1) { + char *handle = codegen_expr(ctx, expr->as.call.args[0]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_vocab_size(%s);", result, handle); + codegen_writeln(ctx, "hml_release(&%s);", handle); + free(handle); + return result; + } + + // __tokenizer_add_special(tok, token, rank) - add special token + if (strcmp(fn_name, "__tokenizer_add_special") == 0 && expr->as.call.num_args == 3) { + char *handle = codegen_expr(ctx, expr->as.call.args[0]); + char *token = codegen_expr(ctx, expr->as.call.args[1]); + char *rank = codegen_expr(ctx, expr->as.call.args[2]); + codegen_writeln(ctx, "HmlValue %s = hml_tokenizer_add_special(%s, %s, %s);", result, handle, token, rank); + codegen_writeln(ctx, "hml_release(&%s);", handle); + codegen_writeln(ctx, "hml_release(&%s);", token); + codegen_writeln(ctx, "hml_release(&%s);", rank); + free(handle); + free(token); + free(rank); + return result; + } + // Handle user-defined function by name (hml_fn_) // Main file functions should use generic call path (hml_call_function) // to properly handle optional parameters with defaults diff --git a/src/backends/interpreter/builtins/internal.h b/src/backends/interpreter/builtins/internal.h index 5129844e..5a4e794b 100644 --- a/src/backends/interpreter/builtins/internal.h +++ b/src/backends/interpreter/builtins/internal.h @@ -319,4 +319,13 @@ Value builtin_regex_error(Value *args, int num_args, ExecutionContext *ctx); Value builtin_regex_replace(Value *args, int num_args, ExecutionContext *ctx); Value builtin_regex_replace_all(Value *args, int num_args, ExecutionContext *ctx); +// Tokenizer (BPE) builtins (tokenizer.c) +Value builtin_tokenizer_create(Value *args, int num_args, ExecutionContext *ctx); +Value builtin_tokenizer_free(Value *args, int num_args, ExecutionContext *ctx); +Value builtin_tokenizer_encode(Value *args, int num_args, ExecutionContext *ctx); +Value builtin_tokenizer_decode(Value *args, int num_args, ExecutionContext *ctx); +Value builtin_tokenizer_count(Value *args, int num_args, ExecutionContext *ctx); +Value builtin_tokenizer_vocab_size(Value *args, int num_args, ExecutionContext *ctx); +Value builtin_tokenizer_add_special(Value *args, int num_args, ExecutionContext *ctx); + #endif // BUILTINS_INTERNAL_H diff --git a/src/backends/interpreter/builtins/registration.c b/src/backends/interpreter/builtins/registration.c index d24abccf..f129ffc9 100644 --- a/src/backends/interpreter/builtins/registration.c +++ b/src/backends/interpreter/builtins/registration.c @@ -197,6 +197,14 @@ static BuiltinInfo builtins[] = { {"__os_name", builtin_os_name}, {"__tmpdir", builtin_tmpdir}, {"__uptime", builtin_uptime}, + // Tokenizer (BPE) builtins (use stdlib/tokenizer.hml module for public API) + {"__tokenizer_create", builtin_tokenizer_create}, + {"__tokenizer_free", builtin_tokenizer_free}, + {"__tokenizer_encode", builtin_tokenizer_encode}, + {"__tokenizer_decode", builtin_tokenizer_decode}, + {"__tokenizer_count", builtin_tokenizer_count}, + {"__tokenizer_vocab_size", builtin_tokenizer_vocab_size}, + {"__tokenizer_add_special", builtin_tokenizer_add_special}, // Unprefixed aliases for parity with compiler (also exposes public API) // Math functions {"sin", builtin_sin}, diff --git a/src/backends/interpreter/builtins/tokenizer.c b/src/backends/interpreter/builtins/tokenizer.c new file mode 100644 index 00000000..69b205b0 --- /dev/null +++ b/src/backends/interpreter/builtins/tokenizer.c @@ -0,0 +1,887 @@ +/* + * Hemlock Tokenizer Builtins + * + * BPE (Byte Pair Encoding) tokenizer for tiktoken-compatible token counting. + * Supports loading vocabulary files in tiktoken's base64 format. + * + * Vocabulary file format (one token per line): + * + * + * The rank is the token ID and also determines merge priority + * (lower rank = higher priority merge). + */ + +#include "internal.h" +#include "../../../../include/hemlock_limits.h" + +#include + +/* ========== BASE64 DECODER ========== */ + +static const unsigned char b64_table[256] = { + ['A'] = 0, ['B'] = 1, ['C'] = 2, ['D'] = 3, + ['E'] = 4, ['F'] = 5, ['G'] = 6, ['H'] = 7, + ['I'] = 8, ['J'] = 9, ['K'] = 10, ['L'] = 11, + ['M'] = 12, ['N'] = 13, ['O'] = 14, ['P'] = 15, + ['Q'] = 16, ['R'] = 17, ['S'] = 18, ['T'] = 19, + ['U'] = 20, ['V'] = 21, ['W'] = 22, ['X'] = 23, + ['Y'] = 24, ['Z'] = 25, + ['a'] = 26, ['b'] = 27, ['c'] = 28, ['d'] = 29, + ['e'] = 30, ['f'] = 31, ['g'] = 32, ['h'] = 33, + ['i'] = 34, ['j'] = 35, ['k'] = 36, ['l'] = 37, + ['m'] = 38, ['n'] = 39, ['o'] = 40, ['p'] = 41, + ['q'] = 42, ['r'] = 43, ['s'] = 44, ['t'] = 45, + ['u'] = 46, ['v'] = 47, ['w'] = 48, ['x'] = 49, + ['y'] = 50, ['z'] = 51, + ['0'] = 52, ['1'] = 53, ['2'] = 54, ['3'] = 55, + ['4'] = 56, ['5'] = 57, ['6'] = 58, ['7'] = 59, + ['8'] = 60, ['9'] = 61, ['+'] = 62, ['/'] = 63, +}; + +/* Decode base64 string into bytes. Returns number of decoded bytes. */ +static int b64_decode(const char *src, int src_len, uint8_t *dst, int dst_cap) { + int out = 0; + int accum = 0; + int bits = 0; + for (int i = 0; i < src_len; i++) { + unsigned char c = (unsigned char)src[i]; + if (c == '=') break; + if (c == '\n' || c == '\r' || c == ' ') continue; + accum = (accum << 6) | b64_table[c]; + bits += 6; + if (bits >= 8) { + bits -= 8; + if (out < dst_cap) { + dst[out++] = (uint8_t)((accum >> bits) & 0xFF); + } + } + } + return out; +} + +/* ========== BPE DATA STRUCTURES ========== */ + +/* A single entry in the vocabulary: byte sequence -> rank (token ID) */ +typedef struct { + uint8_t *bytes; /* Token byte sequence (owned) */ + int byte_len; /* Length of byte sequence */ + int32_t rank; /* Token rank/ID */ +} BpeVocabEntry; + +/* Hash table entry for byte-sequence -> rank lookup */ +typedef struct { + uint8_t *key; /* Byte sequence (not owned, points into vocab) */ + int key_len; /* Key length */ + int32_t value; /* Rank/token ID, or -1 if empty */ +} BpeHashEntry; + +/* The BPE tokenizer state */ +typedef struct { + /* Vocabulary: array indexed by rank -> byte sequence */ + BpeVocabEntry *vocab; /* Array of vocab entries */ + int32_t vocab_size; /* Number of entries */ + int32_t vocab_capacity; /* Allocated capacity */ + + /* Hash table: byte sequence -> rank (for encoding) */ + BpeHashEntry *hash_table; + int32_t hash_capacity; + + /* Special token support */ + char **special_tokens; /* Array of special token strings */ + int32_t *special_ranks; /* Corresponding ranks */ + int32_t num_special; + int32_t special_capacity; +} BpeTokenizer; + +/* ========== HASH TABLE OPERATIONS ========== */ + +/* FNV-1a hash for byte sequences */ +static uint32_t fnv1a_bytes(const uint8_t *data, int len) { + uint32_t hash = 2166136261u; + for (int i = 0; i < len; i++) { + hash ^= data[i]; + hash *= 16777619u; + } + return hash; +} + +/* Initialize hash table */ +static void bpe_hash_init(BpeTokenizer *tok, int32_t capacity) { + tok->hash_capacity = capacity; + tok->hash_table = calloc((size_t)capacity, sizeof(BpeHashEntry)); + if (!tok->hash_table) return; + for (int32_t i = 0; i < capacity; i++) { + tok->hash_table[i].value = -1; + } +} + +/* Insert into hash table */ +static void bpe_hash_insert(BpeTokenizer *tok, uint8_t *key, int key_len, int32_t rank) { + if (!tok->hash_table) return; + uint32_t idx = fnv1a_bytes(key, key_len) % (uint32_t)tok->hash_capacity; + for (int32_t probe = 0; probe < tok->hash_capacity; probe++) { + int32_t i = (int32_t)((idx + (uint32_t)probe) % (uint32_t)tok->hash_capacity); + if (tok->hash_table[i].value == -1) { + tok->hash_table[i].key = key; + tok->hash_table[i].key_len = key_len; + tok->hash_table[i].value = rank; + return; + } + } +} + +/* Lookup in hash table. Returns rank or -1 if not found. */ +static int32_t bpe_hash_lookup(BpeTokenizer *tok, const uint8_t *key, int key_len) { + if (!tok->hash_table) return -1; + uint32_t idx = fnv1a_bytes(key, key_len) % (uint32_t)tok->hash_capacity; + for (int32_t probe = 0; probe < tok->hash_capacity; probe++) { + int32_t i = (int32_t)((idx + (uint32_t)probe) % (uint32_t)tok->hash_capacity); + if (tok->hash_table[i].value == -1) return -1; + if (tok->hash_table[i].key_len == key_len && + memcmp(tok->hash_table[i].key, key, (size_t)key_len) == 0) { + return tok->hash_table[i].value; + } + } + return -1; +} + +/* ========== BPE TOKENIZER LIFECYCLE ========== */ + +/* Free all tokenizer resources */ +static void bpe_tokenizer_free(BpeTokenizer *tok) { + if (!tok) return; + if (tok->vocab) { + for (int32_t i = 0; i < tok->vocab_size; i++) { + free(tok->vocab[i].bytes); + } + free(tok->vocab); + } + free(tok->hash_table); + if (tok->special_tokens) { + for (int32_t i = 0; i < tok->num_special; i++) { + free(tok->special_tokens[i]); + } + free(tok->special_tokens); + } + free(tok->special_ranks); + free(tok); +} + +/* Load vocabulary from tiktoken-format file */ +static BpeTokenizer *bpe_tokenizer_load(const char *path, char *err_buf, int err_buf_size) { + FILE *fp = fopen(path, "r"); + if (!fp) { + snprintf(err_buf, (size_t)err_buf_size, "cannot open vocabulary file: %s", path); + return NULL; + } + + BpeTokenizer *tok = calloc(1, sizeof(BpeTokenizer)); + if (!tok) { + fclose(fp); + snprintf(err_buf, (size_t)err_buf_size, "out of memory"); + return NULL; + } + + /* Initial allocation for vocab array */ + tok->vocab_capacity = 4096; + tok->vocab = calloc((size_t)tok->vocab_capacity, sizeof(BpeVocabEntry)); + if (!tok->vocab) { + fclose(fp); + bpe_tokenizer_free(tok); + snprintf(err_buf, (size_t)err_buf_size, "out of memory"); + return NULL; + } + + char line[HML_BPE_MAX_VOCAB_LINE_LENGTH]; + int32_t max_rank = -1; + + while (fgets(line, sizeof(line), fp)) { + /* Skip empty lines and comments */ + if (line[0] == '\n' || line[0] == '#') continue; + + /* Find the space separating base64 and rank */ + char *space = strchr(line, ' '); + if (!space) continue; + + *space = '\0'; + char *b64_str = line; + int b64_len = (int)(space - line); + int32_t rank = (int32_t)strtol(space + 1, NULL, 10); + + /* Decode base64 to bytes */ + uint8_t decoded[HML_BPE_BASE64_DECODE_BUFSIZE]; + int decoded_len = b64_decode(b64_str, b64_len, decoded, HML_BPE_BASE64_DECODE_BUFSIZE); + if (decoded_len <= 0) continue; + + /* Grow vocab array if needed */ + if (tok->vocab_size >= tok->vocab_capacity) { + int32_t new_cap = tok->vocab_capacity * HML_GROWTH_FACTOR; + BpeVocabEntry *new_vocab = realloc(tok->vocab, (size_t)new_cap * sizeof(BpeVocabEntry)); + if (!new_vocab) continue; + tok->vocab = new_vocab; + tok->vocab_capacity = new_cap; + } + + /* Store vocabulary entry */ + uint8_t *bytes_copy = malloc((size_t)decoded_len); + if (!bytes_copy) continue; + memcpy(bytes_copy, decoded, (size_t)decoded_len); + + BpeVocabEntry *entry = &tok->vocab[tok->vocab_size]; + entry->bytes = bytes_copy; + entry->byte_len = decoded_len; + entry->rank = rank; + tok->vocab_size++; + + if (rank > max_rank) max_rank = rank; + } + fclose(fp); + + if (tok->vocab_size == 0) { + bpe_tokenizer_free(tok); + snprintf(err_buf, (size_t)err_buf_size, "vocabulary file is empty or invalid: %s", path); + return NULL; + } + + /* Build hash table for encoding (byte sequence -> rank) */ + int32_t hash_cap = tok->vocab_size * 3; /* ~33% load factor */ + if (hash_cap < HML_BPE_INITIAL_HASH_CAPACITY) { + hash_cap = HML_BPE_INITIAL_HASH_CAPACITY; + } + bpe_hash_init(tok, hash_cap); + + for (int32_t i = 0; i < tok->vocab_size; i++) { + bpe_hash_insert(tok, tok->vocab[i].bytes, tok->vocab[i].byte_len, tok->vocab[i].rank); + } + + return tok; +} + +/* ========== BPE ENCODING ALGORITHM ========== */ + +/* + * BPE encoding for a single chunk of text (pre-tokenized piece). + * Implements the standard BPE merge algorithm: + * 1. Start with individual bytes as tokens + * 2. Repeatedly find the pair with the lowest rank + * 3. Merge that pair + * 4. Continue until no more merges + * + * Returns an array of token IDs (ranks). + * The caller must free the returned array. + */ + +/* Work item: a linked list node for the BPE merge algorithm */ +typedef struct BpeNode { + uint8_t *bytes; /* Pointer into combined buffer */ + int byte_len; /* Length of this token's bytes */ + struct BpeNode *next; + struct BpeNode *prev; +} BpeNode; + +static int bpe_encode_chunk(BpeTokenizer *tok, const uint8_t *data, int data_len, + int32_t **out_tokens, int *out_count) { + if (data_len == 0) { + *out_tokens = NULL; + *out_count = 0; + return 0; + } + + /* Allocate nodes: one per byte initially */ + int node_capacity = data_len + 1; + BpeNode *nodes = calloc((size_t)node_capacity, sizeof(BpeNode)); + if (!nodes) return -1; + + /* Make a copy of the input data to work with */ + uint8_t *data_copy = malloc((size_t)data_len); + if (!data_copy) { free(nodes); return -1; } + memcpy(data_copy, data, (size_t)data_len); + + /* Initialize doubly-linked list: each node is one byte */ + for (int i = 0; i < data_len; i++) { + nodes[i].bytes = data_copy + i; + nodes[i].byte_len = 1; + nodes[i].prev = (i > 0) ? &nodes[i - 1] : NULL; + nodes[i].next = (i < data_len - 1) ? &nodes[i + 1] : NULL; + } + + /* Temporary buffer for pair lookup */ + uint8_t pair_buf[HML_BPE_MAX_TOKEN_LENGTH * 2]; + + /* Repeatedly merge the highest-priority (lowest-rank) pair */ + int changed = 1; + while (changed) { + changed = 0; + int32_t best_rank = INT32_MAX; + BpeNode *best_node = NULL; + + /* Find the pair with the lowest rank */ + for (BpeNode *node = &nodes[0]; node && node->next; node = node->next) { + if (node->byte_len <= 0) continue; + BpeNode *next = node->next; + int pair_len = node->byte_len + next->byte_len; + if (pair_len > HML_BPE_MAX_TOKEN_LENGTH * 2) continue; + + /* Build the merged byte sequence */ + memcpy(pair_buf, node->bytes, (size_t)node->byte_len); + memcpy(pair_buf + node->byte_len, next->bytes, (size_t)next->byte_len); + + int32_t rank = bpe_hash_lookup(tok, pair_buf, pair_len); + if (rank >= 0 && rank < best_rank) { + best_rank = rank; + best_node = node; + } + } + + if (best_node) { + /* Merge best_node with best_node->next */ + BpeNode *next = best_node->next; + int new_len = best_node->byte_len + next->byte_len; + + /* Allocate merged byte sequence */ + uint8_t *merged = malloc((size_t)new_len); + if (!merged) break; + memcpy(merged, best_node->bytes, (size_t)best_node->byte_len); + memcpy(merged + best_node->byte_len, next->bytes, (size_t)next->byte_len); + + /* If best_node->bytes was previously allocated by us (not in data_copy), free it */ + if (best_node->bytes < data_copy || best_node->bytes >= data_copy + data_len) { + free(best_node->bytes); + } + + best_node->bytes = merged; + best_node->byte_len = new_len; + + /* Remove next node from linked list */ + best_node->next = next->next; + if (next->next) { + next->next->prev = best_node; + } + + /* Free next node's separately allocated bytes if any */ + if (next->bytes < data_copy || next->bytes >= data_copy + data_len) { + free(next->bytes); + } + next->bytes = NULL; + next->byte_len = 0; + + changed = 1; + } + } + + /* Collect results: look up each remaining node in the vocab */ + int count = 0; + for (BpeNode *node = &nodes[0]; node; node = node->next) { + if (node->byte_len > 0) count++; + } + + int32_t *tokens = malloc((size_t)count * sizeof(int32_t)); + if (!tokens) { + /* Clean up allocated bytes */ + for (BpeNode *node = &nodes[0]; node; node = node->next) { + if (node->bytes && (node->bytes < data_copy || node->bytes >= data_copy + data_len)) { + free(node->bytes); + } + } + free(data_copy); + free(nodes); + return -1; + } + + int idx = 0; + for (BpeNode *node = &nodes[0]; node; node = node->next) { + if (node->byte_len <= 0) continue; + int32_t rank = bpe_hash_lookup(tok, node->bytes, node->byte_len); + if (rank >= 0) { + tokens[idx++] = rank; + } else { + /* Unknown token: encode individual bytes as fallback */ + for (int b = 0; b < node->byte_len; b++) { + int32_t byte_rank = bpe_hash_lookup(tok, &node->bytes[b], 1); + if (byte_rank >= 0) { + /* Grow tokens array if needed */ + if (idx >= count) { + count += node->byte_len; + int32_t *new_tokens = realloc(tokens, (size_t)count * sizeof(int32_t)); + if (!new_tokens) { free(tokens); tokens = NULL; break; } + tokens = new_tokens; + } + tokens[idx++] = byte_rank; + } + } + if (!tokens) break; + } + } + + /* Clean up dynamically allocated node bytes */ + for (int i = 0; i < node_capacity; i++) { + if (nodes[i].bytes && (nodes[i].bytes < data_copy || nodes[i].bytes >= data_copy + data_len)) { + /* Only free if it was separately allocated (not pointing into data_copy) */ + /* Check if another node already freed this pointer */ + int already_freed = 0; + for (int j = 0; j < i; j++) { + if (nodes[j].bytes == nodes[i].bytes) { already_freed = 1; break; } + } + if (!already_freed) free(nodes[i].bytes); + } + } + free(data_copy); + free(nodes); + + *out_tokens = tokens; + *out_count = idx; + return 0; +} + +/* ========== SIMPLE PRE-TOKENIZER ========== */ + +/* + * Split text into chunks for BPE processing. + * This is a simplified pre-tokenizer that splits on: + * - Whitespace boundaries (keeping leading whitespace with the following word) + * - Punctuation boundaries + * Similar to GPT-style pre-tokenization but without requiring PCRE. + */ + +typedef struct { + const uint8_t *start; + int len; +} TextChunk; + +static int is_letter_or_digit(uint8_t c) { + return (c >= 'a' && c <= 'z') || (c >= 'A' && c <= 'Z') || + (c >= '0' && c <= '9') || c >= 128; /* Include non-ASCII as letters */ +} + +static int pretokenize(const uint8_t *text, int text_len, + TextChunk **out_chunks, int *out_count) { + if (text_len == 0) { + *out_chunks = NULL; + *out_count = 0; + return 0; + } + + int capacity = HML_BPE_INITIAL_WORK_CAPACITY; + TextChunk *chunks = malloc((size_t)capacity * sizeof(TextChunk)); + if (!chunks) return -1; + int count = 0; + + int i = 0; + while (i < text_len) { + int start = i; + + /* Consume optional leading whitespace (attach to following token) */ + while (i < text_len && (text[i] == ' ' || text[i] == '\t')) i++; + + if (i < text_len) { + if (text[i] == '\n' || text[i] == '\r') { + /* Newlines are their own chunk */ + i++; + if (i < text_len && text[i - 1] == '\r' && text[i] == '\n') i++; + } else if (is_letter_or_digit(text[i])) { + /* Consume word characters (letters + digits) */ + while (i < text_len && is_letter_or_digit(text[i])) i++; + } else { + /* Single punctuation/symbol character */ + i++; + } + } + + if (i > start) { + if (count >= capacity) { + capacity *= HML_GROWTH_FACTOR; + TextChunk *new_chunks = realloc(chunks, (size_t)capacity * sizeof(TextChunk)); + if (!new_chunks) { free(chunks); return -1; } + chunks = new_chunks; + } + chunks[count].start = text + start; + chunks[count].len = i - start; + count++; + } + } + + *out_chunks = chunks; + *out_count = count; + return 0; +} + +/* ========== FULL BPE ENCODE ========== */ + +static int bpe_encode(BpeTokenizer *tok, const char *text, int text_len, + int32_t **out_tokens, int *out_count) { + /* Pre-tokenize the input */ + TextChunk *chunks = NULL; + int num_chunks = 0; + if (pretokenize((const uint8_t *)text, text_len, &chunks, &num_chunks) != 0) { + return -1; + } + + /* Encode each chunk and collect all tokens */ + int total_capacity = text_len + 16; + int32_t *all_tokens = malloc((size_t)total_capacity * sizeof(int32_t)); + if (!all_tokens) { free(chunks); return -1; } + int total_count = 0; + + for (int c = 0; c < num_chunks; c++) { + int32_t *chunk_tokens = NULL; + int chunk_count = 0; + + if (bpe_encode_chunk(tok, chunks[c].start, chunks[c].len, + &chunk_tokens, &chunk_count) != 0) { + free(all_tokens); + free(chunks); + return -1; + } + + /* Grow if needed */ + while (total_count + chunk_count > total_capacity) { + total_capacity *= HML_GROWTH_FACTOR; + int32_t *new_all = realloc(all_tokens, (size_t)total_capacity * sizeof(int32_t)); + if (!new_all) { free(chunk_tokens); free(all_tokens); free(chunks); return -1; } + all_tokens = new_all; + } + + if (chunk_tokens) { + memcpy(all_tokens + total_count, chunk_tokens, (size_t)chunk_count * sizeof(int32_t)); + total_count += chunk_count; + free(chunk_tokens); + } + } + + free(chunks); + *out_tokens = all_tokens; + *out_count = total_count; + return 0; +} + +/* ========== BPE DECODE ========== */ + +static char *bpe_decode(BpeTokenizer *tok, const int32_t *tokens, int num_tokens, int *out_len) { + /* Calculate total length */ + int total_len = 0; + for (int i = 0; i < num_tokens; i++) { + /* Find the vocab entry for this rank */ + for (int32_t v = 0; v < tok->vocab_size; v++) { + if (tok->vocab[v].rank == tokens[i]) { + total_len += tok->vocab[v].byte_len; + break; + } + } + } + + char *result = malloc((size_t)(total_len + 1)); + if (!result) return NULL; + + int offset = 0; + for (int i = 0; i < num_tokens; i++) { + for (int32_t v = 0; v < tok->vocab_size; v++) { + if (tok->vocab[v].rank == tokens[i]) { + memcpy(result + offset, tok->vocab[v].bytes, (size_t)tok->vocab[v].byte_len); + offset += tok->vocab[v].byte_len; + break; + } + } + } + result[offset] = '\0'; + *out_len = offset; + return result; +} + +/* ========== HEMLOCK BUILTIN FUNCTIONS ========== */ + +/* + * __tokenizer_create(path: string) -> ptr + * + * Load a BPE vocabulary from a tiktoken-format file and return + * an opaque tokenizer handle (ptr). + */ +Value builtin_tokenizer_create(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 1) { + runtime_error(ctx, "tokenizer_create() expects 1 argument (vocab_path), got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_STRING) { + runtime_error(ctx, "tokenizer_create() argument must be a string path, got %s", + value_type_name(args[0].type)); + return val_null(); + } + + const char *path = args[0].as.as_string->data; + char err_buf[HML_ERROR_MESSAGE_SIZE]; + BpeTokenizer *tok = bpe_tokenizer_load(path, err_buf, sizeof(err_buf)); + if (!tok) { + runtime_error(ctx, "%s", err_buf); + return val_null(); + } + + return val_ptr(tok); +} + +/* + * __tokenizer_free(tok: ptr) -> null + * + * Free the tokenizer and all its resources. + */ +Value builtin_tokenizer_free(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 1) { + runtime_error(ctx, "tokenizer_free() expects 1 argument, got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_PTR) { + runtime_error(ctx, "tokenizer_free() argument must be a ptr, got %s", + value_type_name(args[0].type)); + return val_null(); + } + + BpeTokenizer *tok = (BpeTokenizer *)args[0].as.as_ptr; + if (tok) { + bpe_tokenizer_free(tok); + } + return val_null(); +} + +/* + * __tokenizer_encode(tok: ptr, text: string) -> array + * + * Encode text into token IDs using BPE. + */ +Value builtin_tokenizer_encode(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 2) { + runtime_error(ctx, "tokenizer_encode() expects 2 arguments (tokenizer, text), got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_PTR) { + runtime_error(ctx, "tokenizer_encode() first argument must be a tokenizer ptr, got %s", + value_type_name(args[0].type)); + return val_null(); + } + if (args[1].type != VAL_STRING) { + runtime_error(ctx, "tokenizer_encode() second argument must be a string, got %s", + value_type_name(args[1].type)); + return val_null(); + } + + BpeTokenizer *tok = (BpeTokenizer *)args[0].as.as_ptr; + if (!tok) { + runtime_error(ctx, "tokenizer_encode() called with null tokenizer"); + return val_null(); + } + + const char *text = args[1].as.as_string->data; + int text_len = args[1].as.as_string->length; + + int32_t *tokens = NULL; + int num_tokens = 0; + if (bpe_encode(tok, text, text_len, &tokens, &num_tokens) != 0) { + runtime_error(ctx, "tokenizer_encode() failed: encoding error"); + return val_null(); + } + + /* Build Hemlock array from token IDs */ + Array *arr = array_new(); + for (int i = 0; i < num_tokens; i++) { + Value v = val_i32(tokens[i]); + array_push(arr, v); + } + free(tokens); + + /* Ownership of array transfers to caller via return value */ + return val_array(arr); +} + +/* + * __tokenizer_decode(tok: ptr, tokens: array) -> string + * + * Decode an array of token IDs back to text. + */ +Value builtin_tokenizer_decode(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 2) { + runtime_error(ctx, "tokenizer_decode() expects 2 arguments (tokenizer, tokens), got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_PTR) { + runtime_error(ctx, "tokenizer_decode() first argument must be a tokenizer ptr, got %s", + value_type_name(args[0].type)); + return val_null(); + } + if (args[1].type != VAL_ARRAY) { + runtime_error(ctx, "tokenizer_decode() second argument must be an array, got %s", + value_type_name(args[1].type)); + return val_null(); + } + + BpeTokenizer *tok = (BpeTokenizer *)args[0].as.as_ptr; + if (!tok) { + runtime_error(ctx, "tokenizer_decode() called with null tokenizer"); + return val_null(); + } + + Array *arr = args[1].as.as_array; + int count = arr->length; + int32_t *token_ids = malloc((size_t)count * sizeof(int32_t)); + if (!token_ids && count > 0) { + runtime_error(ctx, "tokenizer_decode() out of memory"); + return val_null(); + } + + for (int i = 0; i < count; i++) { + Value elem = array_get(arr, i, ctx); + token_ids[i] = value_to_int(elem); + } + + int result_len = 0; + char *decoded = bpe_decode(tok, token_ids, count, &result_len); + free(token_ids); + + if (!decoded) { + runtime_error(ctx, "tokenizer_decode() failed: decode error"); + return val_null(); + } + + Value result = val_string(decoded); + free(decoded); + return result; +} + +/* + * __tokenizer_count(tok: ptr, text: string) -> i32 + * + * Count the number of tokens in text without returning the full array. + * More efficient for just checking context limits. + */ +Value builtin_tokenizer_count(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 2) { + runtime_error(ctx, "tokenizer_count() expects 2 arguments (tokenizer, text), got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_PTR) { + runtime_error(ctx, "tokenizer_count() first argument must be a tokenizer ptr, got %s", + value_type_name(args[0].type)); + return val_null(); + } + if (args[1].type != VAL_STRING) { + runtime_error(ctx, "tokenizer_count() second argument must be a string, got %s", + value_type_name(args[1].type)); + return val_null(); + } + + BpeTokenizer *tok = (BpeTokenizer *)args[0].as.as_ptr; + if (!tok) { + runtime_error(ctx, "tokenizer_count() called with null tokenizer"); + return val_null(); + } + + const char *text = args[1].as.as_string->data; + int text_len = args[1].as.as_string->length; + + int32_t *tokens = NULL; + int num_tokens = 0; + if (bpe_encode(tok, text, text_len, &tokens, &num_tokens) != 0) { + runtime_error(ctx, "tokenizer_count() failed: encoding error"); + return val_null(); + } + free(tokens); + + return val_i32(num_tokens); +} + +/* + * __tokenizer_vocab_size(tok: ptr) -> i32 + * + * Return the number of tokens in the vocabulary. + */ +Value builtin_tokenizer_vocab_size(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 1) { + runtime_error(ctx, "tokenizer_vocab_size() expects 1 argument, got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_PTR) { + runtime_error(ctx, "tokenizer_vocab_size() first argument must be a tokenizer ptr, got %s", + value_type_name(args[0].type)); + return val_null(); + } + + BpeTokenizer *tok = (BpeTokenizer *)args[0].as.as_ptr; + if (!tok) { + return val_i32(0); + } + + return val_i32(tok->vocab_size); +} + +/* + * __tokenizer_add_special(tok: ptr, token: string, rank: i32) -> null + * + * Add a special token to the tokenizer vocabulary. + * Special tokens are matched exactly before BPE encoding. + */ +Value builtin_tokenizer_add_special(Value *args, int num_args, ExecutionContext *ctx) { + if (num_args != 3) { + runtime_error(ctx, "tokenizer_add_special() expects 3 arguments (tokenizer, token, rank), got %d", num_args); + return val_null(); + } + if (args[0].type != VAL_PTR || args[1].type != VAL_STRING) { + runtime_error(ctx, "tokenizer_add_special() invalid argument types"); + return val_null(); + } + + BpeTokenizer *tok = (BpeTokenizer *)args[0].as.as_ptr; + if (!tok) { + runtime_error(ctx, "tokenizer_add_special() called with null tokenizer"); + return val_null(); + } + + const char *token_str = args[1].as.as_string->data; + int32_t rank = value_to_int(args[2]); + + /* Grow special tokens array if needed */ + if (tok->num_special >= tok->special_capacity) { + int32_t new_cap = tok->special_capacity == 0 ? 16 : tok->special_capacity * HML_GROWTH_FACTOR; + char **new_tokens = realloc(tok->special_tokens, (size_t)new_cap * sizeof(char *)); + int32_t *new_ranks = realloc(tok->special_ranks, (size_t)new_cap * sizeof(int32_t)); + if (!new_tokens || !new_ranks) { + free(new_tokens); + free(new_ranks); + runtime_error(ctx, "tokenizer_add_special() out of memory"); + return val_null(); + } + tok->special_tokens = new_tokens; + tok->special_ranks = new_ranks; + tok->special_capacity = new_cap; + } + + tok->special_tokens[tok->num_special] = strdup(token_str); + tok->special_ranks[tok->num_special] = rank; + tok->num_special++; + + /* Also add to the vocab hash table for decoding */ + int token_len = (int)strlen(token_str); + uint8_t *bytes_copy = malloc((size_t)token_len); + if (bytes_copy) { + memcpy(bytes_copy, token_str, (size_t)token_len); + + /* Grow vocab if needed */ + if (tok->vocab_size >= tok->vocab_capacity) { + int32_t new_cap = tok->vocab_capacity * HML_GROWTH_FACTOR; + BpeVocabEntry *new_vocab = realloc(tok->vocab, (size_t)new_cap * sizeof(BpeVocabEntry)); + if (new_vocab) { + tok->vocab = new_vocab; + tok->vocab_capacity = new_cap; + } + } + + if (tok->vocab_size < tok->vocab_capacity) { + BpeVocabEntry *entry = &tok->vocab[tok->vocab_size]; + entry->bytes = bytes_copy; + entry->byte_len = token_len; + entry->rank = rank; + tok->vocab_size++; + + bpe_hash_insert(tok, bytes_copy, token_len, rank); + } else { + free(bytes_copy); + } + } + + return val_null(); +} diff --git a/stdlib/docs/tokenizer.md b/stdlib/docs/tokenizer.md new file mode 100644 index 00000000..3bd16a40 --- /dev/null +++ b/stdlib/docs/tokenizer.md @@ -0,0 +1,235 @@ +# Hemlock Tokenizer Module + +BPE (Byte Pair Encoding) tokenizer for tiktoken-compatible token counting. Use this to stay under LLM context limits or analyze how text maps to tokens. + +## Overview + +- Load vocabulary files in tiktoken's base64 format +- Encode text to token IDs and decode back +- Count tokens efficiently without building the full array +- Add special tokens (e.g., `<|endoftext|>`) +- Compatible with any BPE vocabulary (GPT-2, GPT-4, etc.) + +## Usage + +```hemlock +import { Tokenizer } from "@stdlib/tokenizer"; +``` + +## Creating a Tokenizer + +### Tokenizer(vocab_path: string): object + +Load a BPE vocabulary file and return a tokenizer object. + +```hemlock +let tok = Tokenizer("cl100k_base.tiktoken"); +``` + +The vocabulary file must be in tiktoken format: each line contains a base64-encoded token byte sequence, a space, and its rank (token ID): + +``` +IQ== 0 +Ig== 1 +Iw== 2 +JA== 3 +``` + +The rank serves as both the token ID and the merge priority (lower rank = higher priority merge). + +## Tokenizer Methods + +### tok.encode(text: string): array\ + +Encode text into an array of token IDs. + +```hemlock +let tokens = tok.encode("Hello, world!"); +print(tokens); // [15496, 11, 995, 0] +print(tokens.length); // 4 +``` + +### tok.decode(tokens: array): string + +Decode an array of token IDs back into text. + +```hemlock +let text = tok.decode([15496, 11, 995, 0]); +print(text); // "Hello, world!" +``` + +### tok.count(text: string): i32 + +Count tokens without building the full array. More memory-efficient than `encode()` when you only need the count. + +```hemlock +let n = tok.count("Hello, world!"); +print(n); // 4 + +// Check if text fits in context window +if (tok.count(prompt) > 4096) { + print("Warning: prompt exceeds context limit"); +} +``` + +### tok.vocab_size: i32 + +The number of tokens in the vocabulary. + +```hemlock +print(tok.vocab_size); // e.g., 100256 +``` + +### tok.add_special(token: string, rank: i32): null + +Add a special token that is matched exactly before BPE encoding. + +```hemlock +tok.add_special("<|endoftext|>", 100257); +tok.add_special("<|fim_prefix|>", 100258); +``` + +### tok.free(): null + +Free the tokenizer and all its resources. The tokenizer must not be used after calling `free()`. + +```hemlock +tok.free(); +// tok.encode("text"); // Would throw "Tokenizer: use after free" +``` + +## Convenience Functions + +For one-off tokenization, use these functions that create and destroy a tokenizer internally. For repeated use, create a `Tokenizer` object instead. + +### encode(text: string, vocab_path: string): array\ + +```hemlock +import { encode } from "@stdlib/tokenizer"; +let tokens = encode("Hello!", "vocab.bpe"); +``` + +### decode(tokens: array, vocab_path: string): string + +```hemlock +import { decode } from "@stdlib/tokenizer"; +let text = decode([15496, 0], "vocab.bpe"); +``` + +### count_tokens(text: string, vocab_path: string): i32 + +```hemlock +import { count_tokens } from "@stdlib/tokenizer"; +let n = count_tokens("Hello, world!", "vocab.bpe"); +``` + +## BPE Algorithm + +The tokenizer implements the standard Byte Pair Encoding algorithm: + +1. **Pre-tokenize**: Split text into chunks at word/whitespace/punctuation boundaries +2. **Initialize**: Each byte in a chunk starts as its own token +3. **Merge**: Repeatedly find the byte pair with the lowest rank (highest priority) in the vocabulary and merge them into a single token +4. **Lookup**: Map each final merged byte sequence to its token ID (rank) + +This produces a deterministic encoding: the same text always produces the same token sequence. + +## Vocabulary File Format + +The vocabulary file follows tiktoken's format: + +``` + +``` + +Each line maps a byte sequence (base64-encoded) to a rank. The rank is both the token ID returned by `encode()` and the merge priority (lower rank = merge first). + +Lines starting with `#` are treated as comments. Empty lines are skipped. + +### Generating a Vocabulary File + +You can convert a tiktoken vocabulary using Python: + +```python +import tiktoken +import base64 + +enc = tiktoken.get_encoding("cl100k_base") +with open("cl100k_base.tiktoken", "w") as f: + for token_bytes, rank in sorted(enc._mergeable_ranks.items(), key=lambda x: x[1]): + f.write(f"{base64.b64encode(token_bytes).decode()} {rank}\n") +``` + +## Examples + +### Context Window Checking + +```hemlock +import { Tokenizer } from "@stdlib/tokenizer"; + +let tok = Tokenizer("cl100k_base.tiktoken"); + +fn fits_in_context(text: string, max_tokens: i32): bool { + return tok.count(text) <= max_tokens; +} + +let prompt = "Explain quantum computing in detail..."; +if (!fits_in_context(prompt, 4096)) { + print("Prompt too long, truncating..."); +} + +tok.free(); +``` + +### Token Analysis + +```hemlock +import { Tokenizer } from "@stdlib/tokenizer"; + +let tok = Tokenizer("vocab.bpe"); +let text = "Hello, world!"; +let tokens = tok.encode(text); + +print("Text:", text); +print("Tokens:", tokens); +print("Token count:", tokens.length); +print("Decoded:", tok.decode(tokens)); + +tok.free(); +``` + +### Batch Token Counting + +```hemlock +import { Tokenizer } from "@stdlib/tokenizer"; +import { read_file } from "@stdlib/fs"; + +let tok = Tokenizer("cl100k_base.tiktoken"); + +let files = ["doc1.txt", "doc2.txt", "doc3.txt"]; +let total = 0; + +for (file in files) { + let content = read_file(file); + let count = tok.count(content); + print(file + ": " + count + " tokens"); + total = total + count; +} + +print("Total: " + total + " tokens"); +tok.free(); +``` + +## Memory Management + +The tokenizer allocates memory for the vocabulary hash table and token data. Always call `tok.free()` when done to release these resources. Using `defer` is recommended: + +```hemlock +let tok = Tokenizer("vocab.bpe"); +defer tok.free(); + +// Use tok... +let n = tok.count("some text"); +``` + +The vocabulary is typically the largest allocation (e.g., ~100K entries for GPT-4's cl100k_base). diff --git a/stdlib/tokenizer.hml b/stdlib/tokenizer.hml new file mode 100644 index 00000000..239526ba --- /dev/null +++ b/stdlib/tokenizer.hml @@ -0,0 +1,119 @@ +// @stdlib/tokenizer - BPE (Byte Pair Encoding) tokenizer for token counting +// +// Provides tiktoken-compatible tokenization for staying under LLM context limits. +// Load vocabulary files in tiktoken's base64 format and encode/decode text. +// +// Usage: +// import { Tokenizer } from "@stdlib/tokenizer"; +// +// let tok = Tokenizer("vocab.bpe"); +// let tokens = tok.encode("Hello, world!"); +// let count = tok.count("Hello, world!"); +// let text = tok.decode(tokens); +// print(tok.vocab_size); +// tok.free(); + +// Create a Tokenizer object from a BPE vocabulary file. +// +// The vocabulary file should be in tiktoken format: each line contains +// a base64-encoded token byte sequence followed by a space and its rank (token ID). +// +// Example vocabulary file: +// IQ== 0 +// Ig== 1 +// Iw== 2 +// ... +// +// Returns an object with encode(), decode(), count(), and free() methods. +export fn Tokenizer(vocab_path: string) { + if (typeof(vocab_path) != "string") { + throw "Tokenizer() requires a string path to vocabulary file"; + } + + let handle = __tokenizer_create(vocab_path); + let _freed = false; + + let tok = { + // Encode text into an array of token IDs (i32). + encode: fn(text: string) { + if (_freed) { throw "Tokenizer: use after free"; } + if (typeof(text) != "string") { + throw "encode() requires a string argument"; + } + return __tokenizer_encode(handle, text); + }, + + // Decode an array of token IDs back into text. + decode: fn(tokens: array) { + if (_freed) { throw "Tokenizer: use after free"; } + if (typeof(tokens) != "array") { + throw "decode() requires an array argument"; + } + return __tokenizer_decode(handle, tokens); + }, + + // Count the number of tokens in text. + // More efficient than encode() when you only need the count. + count: fn(text: string) { + if (_freed) { throw "Tokenizer: use after free"; } + if (typeof(text) != "string") { + throw "count() requires a string argument"; + } + return __tokenizer_count(handle, text); + }, + + // Get the vocabulary size (number of tokens). + vocab_size: __tokenizer_vocab_size(handle), + + // Add a special token (e.g., "<|endoftext|>") with a specific rank. + // Special tokens are matched exactly before BPE encoding. + add_special: fn(token: string, rank: i32) { + if (_freed) { throw "Tokenizer: use after free"; } + if (typeof(token) != "string") { + throw "add_special() requires a string token"; + } + __tokenizer_add_special(handle, token, rank); + }, + + // Free the tokenizer and all its resources. + // The tokenizer must not be used after calling free(). + free: fn() { + if (!_freed) { + __tokenizer_free(handle); + _freed = true; + } + } + }; + + return tok; +} + +// Convenience function: encode text to tokens using a vocabulary file. +// Creates a temporary tokenizer, encodes, and frees it. +// For repeated use, create a Tokenizer object instead. +export fn encode(text: string, vocab_path: string) { + let tok = Tokenizer(vocab_path); + let result = tok.encode(text); + tok.free(); + return result; +} + +// Convenience function: decode tokens to text using a vocabulary file. +// Creates a temporary tokenizer, decodes, and frees it. +// For repeated use, create a Tokenizer object instead. +export fn decode(tokens: array, vocab_path: string) { + let tok = Tokenizer(vocab_path); + let result = tok.decode(tokens); + tok.free(); + return result; +} + +// Convenience function: count tokens in text using a vocabulary file. +// Creates a temporary tokenizer, counts, and frees it. +// For repeated use, create a Tokenizer object instead. +export fn count_tokens(text: string, vocab_path: string) { + let tok = Tokenizer(vocab_path); + let result = tok.count(text); + tok.free(); + return result; +} diff --git a/tests/parity/modules/stdlib_tokenizer.expected b/tests/parity/modules/stdlib_tokenizer.expected new file mode 100644 index 00000000..f50123ba --- /dev/null +++ b/tests/parity/modules/stdlib_tokenizer.expected @@ -0,0 +1,8 @@ +65 +271 +A +Hello +1 +2 +0 +true diff --git a/tests/parity/modules/stdlib_tokenizer.hml b/tests/parity/modules/stdlib_tokenizer.hml new file mode 100644 index 00000000..57c634ce --- /dev/null +++ b/tests/parity/modules/stdlib_tokenizer.hml @@ -0,0 +1,28 @@ +// Parity test: @stdlib/tokenizer +import { Tokenizer } from "@stdlib/tokenizer"; + +let tok = Tokenizer("tests/stdlib_tokenizer/test_vocab.bpe"); + +// Test basic encoding +let tokens = tok.encode("A"); +print(tokens[0]); + +// Test BPE merge +let hello_tokens = tok.encode("Hello"); +print(hello_tokens[0]); + +// Test decode +print(tok.decode([65])); +print(tok.decode([271])); + +// Test count +print(tok.count("Hello")); +print(tok.count("AB")); + +// Test empty string +print(tok.encode("").length); + +// Test vocab size +print(tok.vocab_size > 0); + +tok.free(); diff --git a/tests/stdlib_tokenizer/gen_test_vocab.py b/tests/stdlib_tokenizer/gen_test_vocab.py new file mode 100644 index 00000000..8c5ac987 --- /dev/null +++ b/tests/stdlib_tokenizer/gen_test_vocab.py @@ -0,0 +1,55 @@ +#!/usr/bin/env python3 +"""Generate a small test BPE vocabulary file for testing the tokenizer module. + +Creates a vocab with: +- 256 single-byte tokens (ranks 0-255) +- Common English bigrams/trigrams as merges (ranks 256+) +""" +import base64 + +vocab = [] + +# Single byte tokens (ranks 0-255) +for i in range(256): + b64 = base64.b64encode(bytes([i])).decode() + vocab.append((b64, i)) + +# Common merges (each pair is the concatenation of byte values) +merges = [ + (b" t", 256), # space + t + (b"he", 257), # h + e + (b"in", 258), # i + n + (b"re", 259), # r + e + (b"on", 260), # o + n + (b"th", 261), # t + h + (b"er", 262), # e + r + (b"an", 263), # a + n + (b" the", 264), # (space+t) + h + e => multi-byte merge + (b"the", 265), # t + h + e + (b"ll", 266), # l + l + (b"He", 267), # H + e + (b"lo", 268), # l + o + (b"Hel", 269), # He + l + (b"Hell", 270), # Hel + l + (b"Hello", 271), # Hell + o + (b", ", 272), # comma + space + (b"wo", 273), # w + o + (b"or", 274), # o + r + (b"wor", 275), # wo + r + (b"ld", 276), # l + d + (b"worl", 277), # wor + l (note: not standard BPE order but works for testing) + (b"world", 278), # worl + d + (b"!", 33), # already exists as single byte, skip +] + +for token_bytes, rank in merges: + if rank <= 255: + continue # Skip single-byte duplicates + b64 = base64.b64encode(token_bytes).decode() + vocab.append((b64, rank)) + +with open("test_vocab.bpe", "w") as f: + for b64, rank in vocab: + f.write(f"{b64} {rank}\n") + +print(f"Generated {len(vocab)} vocab entries") diff --git a/tests/stdlib_tokenizer/test_tokenizer.hml b/tests/stdlib_tokenizer/test_tokenizer.hml new file mode 100644 index 00000000..aea99a28 --- /dev/null +++ b/tests/stdlib_tokenizer/test_tokenizer.hml @@ -0,0 +1,85 @@ +// Test: @stdlib/tokenizer BPE tokenizer module + +import { Tokenizer } from "@stdlib/tokenizer"; + +// Get the vocab path relative to where tests are run +let vocab_path = "tests/stdlib_tokenizer/test_vocab.bpe"; + +// === Test 1: Create tokenizer === +let tok = Tokenizer(vocab_path); +assert(tok != null, "Tokenizer should not be null"); +print("PASS: Tokenizer created"); + +// === Test 2: Vocab size === +assert(tok.vocab_size > 0, "vocab_size should be > 0"); +print("PASS: vocab_size = " + tok.vocab_size); + +// === Test 3: Encode single byte === +let tokens_a = tok.encode("A"); +assert(tokens_a.length == 1, "Single char 'A' should encode to 1 token"); +assert(tokens_a[0] == 65, "'A' should encode to rank 65 (ASCII)"); +print("PASS: encode('A') = [" + tokens_a[0] + "]"); + +// === Test 4: Encode with BPE merges === +let tokens_hello = tok.encode("Hello"); +// With our test vocab: "Hello" should merge to rank 271 (Hello) +assert(tokens_hello.length == 1, "'Hello' should merge to 1 token, got " + tokens_hello.length); +assert(tokens_hello[0] == 271, "'Hello' should encode to rank 271, got " + tokens_hello[0]); +print("PASS: encode('Hello') = [" + tokens_hello[0] + "]"); + +// === Test 5: Decode === +let decoded = tok.decode([271]); +assert(decoded == "Hello", "decode([271]) should be 'Hello', got '" + decoded + "'"); +print("PASS: decode([271]) = '" + decoded + "'"); + +// === Test 6: Roundtrip === +let text = "Hello"; +let roundtrip = tok.decode(tok.encode(text)); +assert(roundtrip == text, "Roundtrip failed: expected '" + text + "', got '" + roundtrip + "'"); +print("PASS: roundtrip('Hello')"); + +// === Test 7: Count tokens === +let count = tok.count("Hello"); +assert(count == 1, "count('Hello') should be 1, got " + count); +print("PASS: count('Hello') = " + count); + +// === Test 8: Empty string === +let tokens_empty = tok.encode(""); +assert(tokens_empty.length == 0, "Empty string should encode to 0 tokens"); +let count_empty = tok.count(""); +assert(count_empty == 0, "count('') should be 0"); +print("PASS: empty string handling"); + +// === Test 9: Single byte encoding/decoding === +let tokens_space = tok.encode(" "); +assert(tokens_space.length == 1, "Space should encode to 1 token"); +assert(tokens_space[0] == 32, "Space should encode to rank 32"); +let decoded_space = tok.decode([32]); +assert(decoded_space == " ", "decode([32]) should be ' '"); +print("PASS: single byte space"); + +// === Test 10: Decode multiple tokens === +let decoded_multi = tok.decode([65, 66, 67]); +assert(decoded_multi == "ABC", "decode([65,66,67]) should be 'ABC', got '" + decoded_multi + "'"); +print("PASS: decode multiple tokens"); + +// === Test 11: Free and use-after-free protection === +tok.free(); +let caught_error = false; +try { + tok.encode("test"); +} catch (e) { + caught_error = true; +} +assert(caught_error, "Should throw on use after free"); +print("PASS: use-after-free protection"); + +// === Test 12: Convenience function count_tokens (not used to avoid vocab reload) === +// Test creating a second tokenizer +let tok2 = Tokenizer(vocab_path); +let count2 = tok2.count("AB"); +assert(count2 == 2, "count('AB') should be 2, got " + count2); +tok2.free(); +print("PASS: second tokenizer instance"); + +print("All tokenizer tests passed!"); diff --git a/tests/stdlib_tokenizer/test_vocab.bpe b/tests/stdlib_tokenizer/test_vocab.bpe new file mode 100644 index 00000000..f44a3495 --- /dev/null +++ b/tests/stdlib_tokenizer/test_vocab.bpe @@ -0,0 +1,279 @@ +AA== 0 +AQ== 1 +Ag== 2 +Aw== 3 +BA== 4 +BQ== 5 +Bg== 6 +Bw== 7 +CA== 8 +CQ== 9 +Cg== 10 +Cw== 11 +DA== 12 +DQ== 13 +Dg== 14 +Dw== 15 +EA== 16 +EQ== 17 +Eg== 18 +Ew== 19 +FA== 20 +FQ== 21 +Fg== 22 +Fw== 23 +GA== 24 +GQ== 25 +Gg== 26 +Gw== 27 +HA== 28 +HQ== 29 +Hg== 30 +Hw== 31 +IA== 32 +IQ== 33 +Ig== 34 +Iw== 35 +JA== 36 +JQ== 37 +Jg== 38 +Jw== 39 +KA== 40 +KQ== 41 +Kg== 42 +Kw== 43 +LA== 44 +LQ== 45 +Lg== 46 +Lw== 47 +MA== 48 +MQ== 49 +Mg== 50 +Mw== 51 +NA== 52 +NQ== 53 +Ng== 54 +Nw== 55 +OA== 56 +OQ== 57 +Og== 58 +Ow== 59 +PA== 60 +PQ== 61 +Pg== 62 +Pw== 63 +QA== 64 +QQ== 65 +Qg== 66 +Qw== 67 +RA== 68 +RQ== 69 +Rg== 70 +Rw== 71 +SA== 72 +SQ== 73 +Sg== 74 +Sw== 75 +TA== 76 +TQ== 77 +Tg== 78 +Tw== 79 +UA== 80 +UQ== 81 +Ug== 82 +Uw== 83 +VA== 84 +VQ== 85 +Vg== 86 +Vw== 87 +WA== 88 +WQ== 89 +Wg== 90 +Ww== 91 +XA== 92 +XQ== 93 +Xg== 94 +Xw== 95 +YA== 96 +YQ== 97 +Yg== 98 +Yw== 99 +ZA== 100 +ZQ== 101 +Zg== 102 +Zw== 103 +aA== 104 +aQ== 105 +ag== 106 +aw== 107 +bA== 108 +bQ== 109 +bg== 110 +bw== 111 +cA== 112 +cQ== 113 +cg== 114 +cw== 115 +dA== 116 +dQ== 117 +dg== 118 +dw== 119 +eA== 120 +eQ== 121 +eg== 122 +ew== 123 +fA== 124 +fQ== 125 +fg== 126 +fw== 127 +gA== 128 +gQ== 129 +gg== 130 +gw== 131 +hA== 132 +hQ== 133 +hg== 134 +hw== 135 +iA== 136 +iQ== 137 +ig== 138 +iw== 139 +jA== 140 +jQ== 141 +jg== 142 +jw== 143 +kA== 144 +kQ== 145 +kg== 146 +kw== 147 +lA== 148 +lQ== 149 +lg== 150 +lw== 151 +mA== 152 +mQ== 153 +mg== 154 +mw== 155 +nA== 156 +nQ== 157 +ng== 158 +nw== 159 +oA== 160 +oQ== 161 +og== 162 +ow== 163 +pA== 164 +pQ== 165 +pg== 166 +pw== 167 +qA== 168 +qQ== 169 +qg== 170 +qw== 171 +rA== 172 +rQ== 173 +rg== 174 +rw== 175 +sA== 176 +sQ== 177 +sg== 178 +sw== 179 +tA== 180 +tQ== 181 +tg== 182 +tw== 183 +uA== 184 +uQ== 185 +ug== 186 +uw== 187 +vA== 188 +vQ== 189 +vg== 190 +vw== 191 +wA== 192 +wQ== 193 +wg== 194 +ww== 195 +xA== 196 +xQ== 197 +xg== 198 +xw== 199 +yA== 200 +yQ== 201 +yg== 202 +yw== 203 +zA== 204 +zQ== 205 +zg== 206 +zw== 207 +0A== 208 +0Q== 209 +0g== 210 +0w== 211 +1A== 212 +1Q== 213 +1g== 214 +1w== 215 +2A== 216 +2Q== 217 +2g== 218 +2w== 219 +3A== 220 +3Q== 221 +3g== 222 +3w== 223 +4A== 224 +4Q== 225 +4g== 226 +4w== 227 +5A== 228 +5Q== 229 +5g== 230 +5w== 231 +6A== 232 +6Q== 233 +6g== 234 +6w== 235 +7A== 236 +7Q== 237 +7g== 238 +7w== 239 +8A== 240 +8Q== 241 +8g== 242 +8w== 243 +9A== 244 +9Q== 245 +9g== 246 +9w== 247 ++A== 248 ++Q== 249 ++g== 250 ++w== 251 +/A== 252 +/Q== 253 +/g== 254 +/w== 255 +IHQ= 256 +aGU= 257 +aW4= 258 +cmU= 259 +b24= 260 +dGg= 261 +ZXI= 262 +YW4= 263 +IHRoZQ== 264 +dGhl 265 +bGw= 266 +SGU= 267 +bG8= 268 +SGVs 269 +SGVsbA== 270 +SGVsbG8= 271 +LCA= 272 +d28= 273 +b3I= 274 +d29y 275 +bGQ= 276 +d29ybA== 277 +d29ybGQ= 278