tok.h

/* tok.h — Minimal BPE tokenizer for Picchio.
 *
 * Reads the HuggingFace tokenizer.json format:
 *   - "model.vocab": {"token": id, ...}
 *   - "model.merges": ["a b", "c d", ...]
 *
 * Implements:
 *   - Decode: token_id → UTF-8 string (direct lookup)
 *   - Encode: UTF-8 string → token_id[] (greedy BPE)
 *
 * The GPT-OSS tokenizer is based on o200k (200K+ tokens).
 * Vocabulary: ~201088 tokens.
 *
 * Simplified encode strategy:
 *   1. Split by whitespace/punctuation (rough pre-tokenization)
 *   2. Byte-fallback: each unrecognized byte → byte token (0x00..0xFF)
 *   3. Iterative BPE merge (longest match first)
 *
 * For v1 we use a "decode-only + byte-level encode" approach:
 *   - Decode is trivial (lookup table)
 *   - Encode uses longest-prefix match on the vocab (greedy, not optimal
 *     but working for generation, where only the prompt encode is needed)
 */

#ifndef PICCHIO_TOK_H
#define PICCHIO_TOK_H

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdint.h>

#define TOK_MAX_VOCAB 210000
#define TOK_MAX_TOKEN_LEN 256

/* ── Tokenizer structure ── */

typedef struct {
    /* Vocabulary: id → string (for decode) */
    char **vocab;         /* [vocab_size] pointers to strings */
    int *vocab_len;       /* [vocab_size] length of each token in bytes */
    int vocab_size;

    /* Hash table for encode: string → id */
    /* Simple open-addressing hash table */
    uint32_t *ht_hash;   /* hash values */
    int *ht_id;           /* token ids (-1 = empty) */
    int ht_cap;           /* capacity (power of 2) */

    /* Special tokens */
    int bos_id;
    int eos_id;
    int pad_id;
} Tokenizer;

/* ── FNV-1a hash ── */

static uint32_t tok_hash(const char *s, int len) {
    uint32_t h = 2166136261u;
    for (int i = 0; i < len; i++) {
        h ^= (uint8_t)s[i];
        h *= 16777619u;
    }
    return h;
}

/* ── Hash table: insert ── */

static void tok_ht_insert(Tokenizer *t, const char *s, int len, int id) {
    uint32_t h = tok_hash(s, len);
    int idx = h & (t->ht_cap - 1);
    while (t->ht_id[idx] != -1) {
        idx = (idx + 1) & (t->ht_cap - 1);
    }
    t->ht_hash[idx] = h;
    t->ht_id[idx] = id;
}

/* ── Hash table: lookup ── */

static int tok_ht_lookup(Tokenizer *t, const char *s, int len) {
    uint32_t h = tok_hash(s, len);
    int idx = h & (t->ht_cap - 1);
    while (t->ht_id[idx] != -1) {
        if (t->ht_hash[idx] == h) {
            int id = t->ht_id[idx];
            if (id < t->vocab_size && t->vocab_len[id] == len &&
                memcmp(t->vocab[id], s, len) == 0) {
                return id;
            }
        }
        idx = (idx + 1) & (t->ht_cap - 1);
    }
    return -1;
}

/* ── Decode a byte escape (e.g. "Ġ" → ' ', "Ċ" → '\n') ── */
/* The HF/tiktoken tokenizer uses a bytes-to-unicode mapping:
 * byte 0x20 (space) → 'Ġ' (U+0120)
 * byte 0x0A (newline) → 'Ċ' (U+010A)
 * etc. For decode we must invert this mapping. */

static int tok_decode_byte_char(const char *s, int *advance) {
    /* UTF-8 decode */
    uint8_t c = (uint8_t)s[0];
    if (c < 0x80) {
        *advance = 1;
        return c;  /* direct ASCII (bytes 33-126 map to themselves) */
    }
    if ((c & 0xE0) == 0xC0 && s[1]) {
        /* 2-byte UTF-8 */
        int cp = ((c & 0x1F) << 6) | (s[1] & 0x3F);
        *advance = 2;
        /* Map Unicode codepoint → original byte */
        /* tiktoken's bytes_to_unicode mapping:
         * byte 0..32 → U+0100..U+0120 (256+byte)
         * byte 127..160 → U+0121..U+014E
         * etc. Simplified approach: */
        if (cp >= 0x100 && cp <= 0x1FF) {
            /* This is the approximate inverse mapping */
            /* bytes 0-32, 127-160, 173 are mapped to U+0100+ */
            static const int byte_map[] = {
                /* U+0100..U+010F → byte 0..15 */
                0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15,
                /* U+0110..U+011F → byte 16..31 */
                16,17,18,19,20,21,22,23,24,25,26,27,28,29,30,31,
                /* U+0120 → byte 32 (space) */
                32,
                /* U+0121..U+017E → byte 127..160, 173 */
            };
            if (cp - 0x100 < (int)(sizeof(byte_map)/sizeof(byte_map[0])))
                return byte_map[cp - 0x100];
            /* Fallback for the high bytes (127-255) */
            return cp - 0x100 + 0;
        }
        return -1;  /* Not a byte escape */
    }
    if ((c & 0xF0) == 0xE0 && s[1] && s[2]) {
        *advance = 3;
        return -1;
    }
    *advance = 1;
    return -1;
}

/* ── Tokenizer parsing: Picchio binary format (picchio_vocab.bin) ── */
/* Much faster than parsing the 28 MB JSON.
 * Generated by export_vocab.py. */

static int tok_load_binary(Tokenizer *t, const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) return -1;

    /* Header */
    char magic[4];
    uint32_t total, n_added, eos_id, pad_id;
    if (fread(magic, 1, 4, f) != 4 || memcmp(magic, "PVOC", 4) != 0) {
        fclose(f); return -1;
    }
    fread(&total, 4, 1, f);
    fread(&n_added, 4, 1, f);
    fread(&eos_id, 4, 1, f);
    fread(&pad_id, 4, 1, f);

    t->vocab_size = (int)total;
    t->eos_id = (int)eos_id;
    t->pad_id = (int)pad_id;
    t->bos_id = -1;

    /* Allocate */
    t->vocab = (char **)calloc(total, sizeof(char *));
    t->vocab_len = (int *)calloc(total, sizeof(int));

    /* Hash table */
    t->ht_cap = 1;
    while (t->ht_cap < total * 4) t->ht_cap <<= 1;
    t->ht_hash = (uint32_t *)calloc(t->ht_cap, sizeof(uint32_t));
    t->ht_id = (int *)malloc(t->ht_cap * sizeof(int));
    for (int i = 0; i < t->ht_cap; i++) t->ht_id[i] = -1;

    /* Read tokens */
    for (uint32_t i = 0; i < total; i++) {
        uint16_t len;
        if (fread(&len, 2, 1, f) != 1) break;
        char *s = (char *)malloc(len + 1);
        if (len > 0) fread(s, 1, len, f);
        s[len] = '\0';
        t->vocab[i] = s;
        t->vocab_len[i] = (int)len;
        if (len > 0)
            tok_ht_insert(t, s, (int)len, (int)i);
    }

    fclose(f);
    fprintf(stderr, "  tok: %d tokens loaded (binary, eos=%d)\n",
            t->vocab_size, t->eos_id);
    return 0;
}

static int tok_load(Tokenizer *t, const char *path) {
    memset(t, 0, sizeof(*t));

    /* Try the binary format first (fast) */
    {
        /* Look for picchio_vocab.bin in the same directory */
        char bin_path[512];
        /* Extract the directory from the path */
        strncpy(bin_path, path, sizeof(bin_path) - 1);
        char *last_sep = strrchr(bin_path, '/');
        if (!last_sep) last_sep = strrchr(bin_path, '\\');
        if (last_sep) {
            strcpy(last_sep + 1, "picchio_vocab.bin");
        } else {
            strcpy(bin_path, "picchio_vocab.bin");
        }

        if (tok_load_binary(t, bin_path) == 0) return 0;
    }

    /* Fallback: parse tokenizer.json (slow, limited support) */
    FILE *f = fopen(path, "rb");
    if (!f) {
        fprintf(stderr, "tok: cannot open %s\n", path);
        fprintf(stderr, "     Generate the binary vocab: python export_vocab.py\n");
        return -1;
    }
    fclose(f);
    fprintf(stderr, "tok: tokenizer.json found but the binary format is required.\n");
    fprintf(stderr, "     Run: python export_vocab.py %s\n", path);
    return -1;
}

/* ── Decode: token_id → UTF-8 string ── */
/* Returns a pointer to the token's string (do not copy, it lives in the vocab) */

static const char *tok_decode(Tokenizer *t, int id) {
    if (id < 0 || id >= t->vocab_size || !t->vocab[id])
        return "";
    return t->vocab[id];
}

/* Decode with the inverse bytes-to-unicode conversion.
 * Writes into dst (max dstlen), returns bytes written. */
static int tok_decode_raw(Tokenizer *t, int id, char *dst, int dstlen) {
    const char *s = tok_decode(t, id);
    int slen = t->vocab_len[id];
    int out = 0;

    for (int i = 0; i < slen && out < dstlen - 1; ) {
        int adv = 0;
        int byte = tok_decode_byte_char(s + i, &adv);
        if (byte >= 0) {
            dst[out++] = (char)byte;
            i += adv;
        } else {
            /* Copy the UTF-8 char as-is */
            uint8_t c = (uint8_t)s[i];
            if (c < 0x80) { dst[out++] = s[i]; i++; }
            else if ((c & 0xE0) == 0xC0) {
                if (out + 2 <= dstlen - 1) { dst[out++] = s[i]; dst[out++] = s[i+1]; }
                i += 2;
            }
            else if ((c & 0xF0) == 0xE0) {
                if (out + 3 <= dstlen - 1) { dst[out++] = s[i]; dst[out++] = s[i+1]; dst[out++] = s[i+2]; }
                i += 3;
            }
            else if ((c & 0xF8) == 0xF0) {
                if (out + 4 <= dstlen - 1) { dst[out++] = s[i]; dst[out++] = s[i+1]; dst[out++] = s[i+2]; dst[out++] = s[i+3]; }
                i += 4;
            }
            else { dst[out++] = s[i]; i++; }
        }
    }
    dst[out] = '\0';
    return out;
}

/* ── Encode: UTF-8 string → token_id[] (greedy longest match) ── */
/* Not optimal BPE, but works for prompt encoding.
 * Returns the number of tokens written to out_ids (max max_tokens). */

static int tok_encode(Tokenizer *t, const char *text, int *out_ids, int max_tokens) {
    int n = 0;
    int len = (int)strlen(text);
    int pos = 0;

    while (pos < len && n < max_tokens) {
        /* Find the longest token that matches starting at pos */
        int best_len = 0;
        int best_id = -1;

        /* Try from the maximum length (TOK_MAX_TOKEN_LEN) downward */
        int max_try = len - pos;
        if (max_try > TOK_MAX_TOKEN_LEN) max_try = TOK_MAX_TOKEN_LEN;

        for (int try_len = max_try; try_len >= 1; try_len--) {
            int id = tok_ht_lookup(t, text + pos, try_len);
            if (id >= 0) {
                best_len = try_len;
                best_id = id;
                break;
            }
        }

        if (best_id >= 0) {
            out_ids[n++] = best_id;
            pos += best_len;
        } else {
            /* Byte fallback: single byte as a token */
            /* Byte tokens in tiktoken/o200k are the first 256 IDs
             * or have a specific mapping. For now, skip. */
            pos++;
        }
    }

    return n;
}

/* ── Free ── */

static void tok_free(Tokenizer *t) {
    if (t->vocab) {
        for (int i = 0; i < t->vocab_size; i++)
            free(t->vocab[i]);
        free(t->vocab);
    }
    free(t->vocab_len);
    free(t->ht_hash);
    free(t->ht_id);
    memset(t, 0, sizeof(*t));
}

#endif /* PICCHIO_TOK_H */