init

.clang-format

@@ -0,0 +1,16 @@
+Language: Cpp
+BasedOnStyle: WebKit
+IndentWidth: 4
+ColumnLimit: 80
+IndentCaseLabels: true
+InsertNewlineAtEOF: true
+AllowShortFunctionsOnASingleLine: None
+
+BinPackArguments: false
+AllowAllArgumentsOnNextLine: true
+
+BinPackParameters: false
+AllowAllParametersOfDeclarationOnNextLine: true
+
+PointerAlignment: Right
+

.gitignore

@@ -0,0 +1,2 @@
+dataset_*
+

compile_flags.txt

@@ -0,0 +1,6 @@
+-xc
+-std=c17
+-Wall
+-Wextra
+-pedantic-errors
+

simple-ngram/.gitignore

@@ -0,0 +1,7 @@
+perf.svg
+tokens.txt
+dataset_*
+
+estimate_word_count.txt
+make_tokens
+train_model

simple-ngram/common.c

@@ -0,0 +1,413 @@
+#include "common.h"
+#include <assert.h>
+#include <ctype.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+static bool ch_is_punctuation(char ch)
+{
+    return strchr(",.!?'", ch);
+}
+
+static bool ch_is_letters(char ch)
+{
+    return isalpha(ch);
+}
+
+size_t next_word(char line[], size_t line_size, FILE *file, int *ch)
+{
+    size_t len = 0;
+
+    if (!*ch)
+        *ch = fgetc(file);
+
+    if (*ch != EOF && !ch_is_punctuation(*ch) && !ch_is_letters(*ch)) {
+        assert(len < line_size - 1);
+        line[len] = ' ';
+        len += 1;
+
+        int newlines = 0;
+
+        while (*ch != EOF && !ch_is_punctuation(*ch) && !ch_is_letters(*ch)) {
+            newlines += *ch == '\n';
+            *ch = fgetc(file);
+        }
+
+        if (newlines >= 2) {
+            line[len - 1] = '\n';
+            return len;
+        }
+    }
+
+    bool matched = false;
+
+    while (*ch != EOF && ch_is_punctuation(*ch)) {
+        matched = true;
+
+        assert(len < line_size - 1);
+        line[len] = *ch;
+        len += 1;
+
+        *ch = fgetc(file);
+    }
+    if (matched)
+        return len;
+
+    while (*ch != EOF && ch_is_letters(*ch)) {
+        matched = true;
+
+        assert(len < line_size - 1);
+        line[len] = tolower(*ch);
+        len += 1;
+
+        *ch = fgetc(file);
+    }
+
+    return len;
+}
+
+void array_push(
+    void **data, size_t *capacity, size_t *size, void *elem, size_t elem_size)
+{
+    if (!*capacity) {
+        *capacity = 4;
+    }
+    if (!*data) {
+        *data = malloc(*capacity * elem_size);
+        *size = 0;
+    }
+    if (*size >= *capacity) {
+        *capacity *= 2;
+        *data = realloc(*data, *capacity * elem_size);
+    }
+    memcpy((unsigned char *)*data + *size * elem_size, elem, elem_size);
+    *size += 1;
+}
+
+void array_push_n(void **data,
+    size_t *capacity,
+    size_t *size,
+    void *elems,
+    size_t elem_size,
+    size_t elem_count)
+{
+    for (size_t i = 0; i < elem_count; ++i) {
+        array_push(data,
+            capacity,
+            size,
+            (unsigned char *)elems + elem_size * i,
+            elem_size);
+    }
+}
+
+typedef int (*CmpFn)(const void *, const void *);
+
+static void *binary_search_internal(void *data,
+    size_t begin,
+    size_t end,
+    const void *entry,
+    size_t entry_size,
+    size_t *out_idx,
+    CmpFn cmp)
+{
+    if (begin == end) {
+        if (out_idx)
+            *out_idx = begin;
+        return NULL;
+    }
+    size_t middle_idx = (end - begin) / 2 + begin;
+    void *middle_entry = (unsigned char *)data + entry_size * middle_idx;
+    int cmp_result = cmp(entry, middle_entry);
+    if (cmp_result == 0) {
+        if (out_idx)
+            *out_idx = middle_idx;
+        return middle_entry;
+    }
+    if (cmp_result < 0)
+        end = middle_idx;
+    else
+        begin = middle_idx + 1;
+
+    return binary_search_internal(
+        data, begin, end, entry, entry_size, out_idx, cmp);
+}
+
+void *binary_search(void *data,
+    size_t count,
+    const void *entry,
+    size_t entry_size,
+    size_t *out_idx,
+    CmpFn cmp)
+{
+    return binary_search_internal(
+        data, 0, count, entry, entry_size, out_idx, cmp);
+}
+
+typedef struct RbTree RbTree;
+
+RbTree *rbtree_new(size_t key_size, size_t val_size, CmpFn cmp);
+void rbtree_free(RbTree *t);
+void rbtree_insert(RbTree *t, const void *key, void *val);
+void *rbtree_get(RbTree *t, const void *key);
+
+typedef enum { Rb_Black, Rb_Red } RbColor;
+typedef enum { Rb_Left, Rb_Right } RbDir;
+
+typedef struct RbNode RbNode;
+
+struct RbNode {
+    RbNode *parent;
+    union {
+        struct {
+            RbNode *left;
+            RbNode *right;
+        };
+        RbNode *child[2];
+    };
+    uint8_t color;
+    int8_t padding[7];
+    uint8_t key_and_val[];
+};
+
+struct RbTree {
+    size_t key_size;
+    size_t val_size;
+    CmpFn cmp;
+    size_t node_key_size;
+    size_t node_val_size;
+    size_t size_of_node;
+    RbNode *root;
+};
+
+RbTree *rbtree_new(size_t key_size, size_t val_size, CmpFn cmp)
+{
+    size_t node_key_size = (key_size + 7ul) & ~7ul;
+    size_t node_val_size = (key_size + 7ul) & ~7ul;
+    size_t size_of_node = sizeof(RbNode) + node_key_size + node_val_size;
+
+    RbTree *t = malloc(sizeof(RbTree));
+    *t = (RbTree) {
+        .key_size = key_size,
+        .val_size = val_size,
+        .cmp = cmp,
+        .node_key_size = node_key_size,
+        .node_val_size = node_val_size,
+        .size_of_node = size_of_node,
+        .root = NULL,
+    };
+    return t;
+}
+
+static void t_free_node(RbNode *node)
+{
+    if (!node)
+        return;
+
+    t_free_node(node->child[0]);
+    t_free_node(node->child[1]);
+
+    free(node);
+}
+
+void rbtree_free(RbTree *t)
+{
+    t_free_node(t->root);
+    free(t);
+}
+
+static RbNode *t_find_node(RbTree *t,
+    RbNode *node,
+    const void *key,
+    RbNode **out_parent,
+    RbDir *out_dir)
+{
+    int cmp_result = t->cmp(key, node->key_and_val);
+
+    if (cmp_result == 0)
+        return node;
+
+    RbDir dir = cmp_result < 0 ? Rb_Left : Rb_Right;
+    RbNode *next = node->child[dir];
+
+    if (next)
+        return t_find_node(t, next, key, out_parent, out_dir);
+
+    if (out_dir)
+        *out_dir = dir;
+    if (out_parent)
+        *out_parent = node;
+
+    return NULL;
+}
+
+static RbNode *t_node_new(
+    RbTree *t, RbNode *parent, RbColor color, const void *key, void *val)
+{
+    RbNode *node = malloc(t->size_of_node);
+    *node = (RbNode) {
+        .parent = parent,
+        .child = { NULL, NULL },
+        .color = color,
+    };
+
+    memcpy(node->key_and_val, key, t->key_size);
+    memcpy(node->key_and_val + t->node_key_size, val, t->val_size);
+    return node;
+}
+
+static RbNode *t_rotate_subtree(RbTree *t, RbNode *sub, RbDir dir)
+{
+    RbNode *sub_parent = sub->parent;
+    RbNode *new_root = sub->child[1 - dir];
+    RbNode *new_child = new_root->child[dir];
+
+    sub->child[1 - dir] = new_child;
+
+    if (new_child)
+        new_child->parent = sub;
+
+    new_root->child[dir] = sub;
+
+    new_root->parent = sub_parent;
+    sub->parent = new_root;
+    if (sub_parent)
+        sub_parent->child[sub == sub_parent->right] = new_root;
+    else
+        t->root = new_root;
+
+    return new_root;
+}
+
+void rbtree_insert(RbTree *t, const void *key, void *val)
+{
+    // https://en.wikipedia.org/wiki/Red%E2%80%93black_tree#Insertion
+
+    RbNode *node, *parent = NULL;
+    RbDir dir = Rb_Left;
+
+    if (t->root && (node = t_find_node(t, t->root, key, &parent, &dir))) {
+        memcpy(node->key_and_val + t->node_key_size, val, t->val_size);
+        return;
+    }
+
+    node = t_node_new(t, parent, Rb_Red, key, val);
+
+    if (!t->root) {
+        t->root = node;
+        return;
+    }
+
+    parent->child[dir] = node;
+
+    do {
+        if (parent->color == Rb_Black)
+            // case 1
+            return;
+
+        RbNode *grandparent = parent->parent;
+
+        if (!grandparent) {
+            // case 4
+            parent->color = Rb_Black;
+            return;
+        }
+
+        dir = parent == grandparent->left ? Rb_Left : Rb_Right;
+        RbNode *uncle = grandparent->child[1 - dir];
+
+        if (!uncle || uncle->color == Rb_Black) {
+            if (node == parent->child[1 - dir]) {
+                // case 5
+                t_rotate_subtree(t, parent, dir);
+                node = parent;
+                parent = grandparent->child[dir];
+            }
+
+            // case 6
+            t_rotate_subtree(t, grandparent, 1 - dir);
+            parent->color = Rb_Black;
+            grandparent->color = Rb_Red;
+            return;
+        }
+
+        // case 2
+        parent->color = Rb_Black;
+        uncle->color = Rb_Black;
+        grandparent->color = Rb_Red;
+        node = grandparent;
+
+    } while ((parent = node->parent));
+
+    // case 3
+}
+
+void *rbtree_get(RbTree *t, const void *key)
+{
+    RbNode *node;
+    if (!t->root || !(node = t_find_node(t, t->root, key, NULL, NULL)))
+        return NULL;
+
+    return node->key_and_val + t->node_key_size;
+}
+
+static bool t_red_node_children_valid(RbNode *node)
+{
+    if (!node)
+        return true;
+    if (node->parent && node->color == Rb_Red)
+        return node->parent->color != Rb_Red;
+    return t_red_node_children_valid(node->left)
+        && t_red_node_children_valid(node->right);
+}
+
+static int t_walk_count_black_nodes(RbNode *node)
+{
+    if (!node)
+        return 1;
+
+    int left = t_walk_count_black_nodes(node->left);
+    int right = t_walk_count_black_nodes(node->right);
+
+    if (left != right)
+        return 0;
+
+    return left + (node->color == Rb_Black);
+}
+
+static bool rbtree_valid(RbTree *t)
+{
+    return t_red_node_children_valid(t->root)
+        && t_walk_count_black_nodes(t->root) != 0;
+}
+
+static int int_cmp(const int *a, const int *b)
+{
+    return *a - *b;
+}
+
+void test_common(void)
+{
+    RbTree *t = rbtree_new(sizeof(int), sizeof(int), (CmpFn)int_cmp);
+
+    int keys[] = { 3, 7, 5, 1, 4, 6 };
+    int vals[] = { 0, 2, 4, 6, 8, 10 };
+    int *result;
+
+    result = rbtree_get(t, &keys[0]);
+    assert(!result);
+
+    for (size_t i = 0; i < sizeof(keys) / sizeof(keys[0]); ++i) {
+        rbtree_insert(t, &keys[i], &vals[i]);
+        assert(rbtree_valid(t));
+    }
+
+    for (size_t i = 0; i < sizeof(keys) / sizeof(keys[0]); ++i) {
+        assert(*(int *)rbtree_get(t, &keys[i]) == vals[i]);
+    }
+
+    rbtree_free(t);
+}

simple-ngram/common.h

@@ -0,0 +1,38 @@
+#ifndef COMMON_H
+#define COMMON_H
+
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdio.h>
+
+size_t next_word(char line[], size_t line_size, FILE *file, int *ch);
+
+void array_push(
+    void **data, size_t *capacity, size_t *size, void *elem, size_t elem_size);
+
+void array_push_n(void **data,
+    size_t *capacity,
+    size_t *size,
+    void *elems,
+    size_t elem_size,
+    size_t elem_count);
+
+typedef int (*CmpFn)(const void *, const void *);
+
+void *binary_search(void *data,
+    size_t count,
+    const void *entry,
+    size_t entry_size,
+    size_t *out_idx,
+    CmpFn cmp);
+
+typedef struct RbTree RbTree;
+
+RbTree *rbtree_new(size_t key_size, size_t val_size, CmpFn cmp);
+void rbtree_free(RbTree *t);
+void rbtree_insert(RbTree *t, const void *key, void *val);
+void *rbtree_get(RbTree *t, const void *key);
+
+void test_common(void);
+
+#endif

simple-ngram/make_tokens.c

@@ -0,0 +1,622 @@
+#include "common.h"
+#include <assert.h>
+#include <errno.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+typedef struct {
+    char *text;
+    uint32_t len;
+    uint32_t count;
+} MultisetEntry;
+
+typedef struct {
+    MultisetEntry *entries;
+    size_t entries_cap;
+    size_t entries_size;
+} Multiset;
+
+static int multisetentry_cmp(const MultisetEntry *a, const MultisetEntry *b)
+{
+    if (a->len != b->len)
+        return (int)a->len - (int)b->len;
+    return strcmp(a->text, b->text);
+}
+
+static int multisetentry_cmp_wordcount(
+    const MultisetEntry *a, const MultisetEntry *b)
+{
+    return (int)a->count - (int)b->count;
+}
+
+static void multimap_search(Multiset *s,
+    const MultisetEntry *entry,
+    MultisetEntry **out_entry,
+    size_t *out_idx)
+{
+    *out_entry = binary_search(s->entries,
+        s->entries_size,
+        (void *)entry,
+        sizeof(*entry),
+        out_idx,
+        (CmpFn)multisetentry_cmp);
+}
+
+void multiset_free(Multiset *s)
+{
+    for (size_t i = 0; i < s->entries_size; ++i)
+        free(s->entries[i].text);
+    if (s->entries)
+        free(s->entries);
+}
+
+void multiset_dump(Multiset *s)
+{
+    printf("entries_size = %lu\n", s->entries_size);
+    for (size_t i = 0; i < s->entries_size; ++i)
+        printf("  %s -> %u\n", s->entries[i].text, s->entries[i].count);
+}
+
+// strdup was introduced in C23, we're using C17
+#if (__STDC_VERSION__ < 202311L)
+char *strdup(const char *str)
+{
+    char *p = malloc(strlen(str) + 1);
+    strcpy(p, str);
+    return p;
+}
+#endif
+
+void multiset_insert(Multiset *s, const char *text)
+{
+    if (!s->entries_size) {
+        MultisetEntry entry = {
+            .text = strdup(text),
+            .len = strlen(text),
+            .count = 1,
+        };
+        array_push((void **)&s->entries,
+            &s->entries_cap,
+            &s->entries_size,
+            &entry,
+            sizeof(entry));
+        return;
+    }
+
+    MultisetEntry search_entry = {
+        .text = (char *)text, // NOTE: cast away const. we'll borrow it shortly
+        .len = strlen(text),
+        .count = 0,
+    };
+
+    MultisetEntry *out_entry;
+    size_t out_idx;
+    multimap_search(s, &search_entry, &out_entry, &out_idx);
+
+    if (out_entry) {
+        out_entry->count += 1;
+        return;
+    }
+
+    MultisetEntry last_entry = s->entries[s->entries_size - 1];
+    array_push((void **)&s->entries,
+        &s->entries_cap,
+        &s->entries_size,
+        &last_entry,
+        sizeof(last_entry));
+
+    memmove(&s->entries[out_idx + 1],
+        &s->entries[out_idx],
+        sizeof(s->entries[0]) * (s->entries_size - out_idx - 1));
+
+    s->entries[out_idx] = (MultisetEntry) {
+        .text = strdup(text),
+        .len = strlen(text),
+        .count = 1,
+    };
+}
+
+void test_multiset(void);
+
+void multiset_sort(Multiset *s)
+{
+    qsort(s->entries,
+        s->entries_size,
+        sizeof(s->entries[0]),
+        (CmpFn)multisetentry_cmp);
+}
+
+void multiset_sort_wordcount(Multiset *s)
+{
+    qsort(s->entries,
+        s->entries_size,
+        sizeof(s->entries[0]),
+        (CmpFn)multisetentry_cmp_wordcount);
+}
+
+static int char_cmp(const char *a, const char *b)
+{
+    return *a - *b;
+}
+
+typedef struct {
+    uint16_t left, right;
+} Tok;
+
+static int tok_cmp(const Tok *a, const Tok *b)
+{
+    return a->left != b->left ? a->left - b->left : a->right - b->right;
+}
+
+typedef struct {
+    const char *chars;
+    size_t chars_size;
+    Tok *toks;
+    size_t toks_cap;
+    size_t toks_size;
+} Vocab;
+
+uint16_t vocab_add_tok(Vocab *v, Tok tok)
+{
+    uint16_t tokval = (uint16_t)v->toks_size + v->chars_size + 1;
+
+    array_push(
+        (void **)&v->toks, &v->toks_cap, &v->toks_size, &tok, sizeof(tok));
+
+    return tokval;
+}
+
+uint16_t vocab_char_tokval(Vocab *v, char ch)
+{
+    size_t idx;
+    char *found = binary_search((void *)v->chars,
+        v->chars_size,
+        &ch,
+        sizeof(ch),
+        &idx,
+        (CmpFn)char_cmp);
+
+    assert(found);
+    return (uint16_t)idx + 1;
+}
+
+size_t vocab_tok_str(
+    Vocab *v, char str[], size_t str_size, const Tok *tok, bool subst_space);
+
+size_t vocab_tokval_str(
+    Vocab *v, char str[], size_t str_size, uint16_t tokval, bool subst_space)
+{
+    if (tokval == 0) {
+        assert(str_size >= 1);
+        str[0] = '0';
+        return 1;
+    }
+
+    if (tokval < v->chars_size + 1) {
+        char ch = v->chars[tokval - 1];
+
+        if (ch == ' ') {
+            if (subst_space) {
+                char *blank = "␣";
+                size_t blank_len = strlen(blank);
+                assert(str_size >= blank_len);
+                memcpy(str, blank, blank_len);
+                return blank_len;
+            }
+
+            assert(str_size >= 1);
+            str[0] = ' ';
+            return 1;
+        }
+
+        assert(str_size >= 1);
+        str[0] = v->chars[tokval - 1];
+        return 1;
+    }
+
+    return vocab_tok_str(
+        v, str, str_size, &v->toks[tokval - v->chars_size - 1], subst_space);
+}
+
+size_t vocab_tok_str(
+    Vocab *v, char str[], size_t str_size, const Tok *tok, bool subst_space)
+{
+    size_t len = 0;
+    len += vocab_tokval_str(v, str, str_size, tok->left, subst_space);
+    len += vocab_tokval_str(
+        v, str + len, str_size - len, tok->right, subst_space);
+    return len;
+}
+
+void vocab_print_tokval(Vocab *v, FILE *file, uint16_t tokval, bool subst_space)
+{
+    char tokstr[60];
+    size_t tokstr_len
+        = vocab_tokval_str(v, tokstr, sizeof(tokstr) - 1, tokval, subst_space);
+    tokstr[tokstr_len] = '\0';
+    fprintf(file, "%s", tokstr);
+}
+
+typedef struct {
+    size_t size_of_word;
+    uint16_t *words;
+    size_t words_size;
+} Corpus;
+
+Tok corpus_most_common_tok_pair(Corpus *c, Vocab *v)
+{
+    (void)v;
+
+    typedef struct {
+        Tok tok;
+        uint32_t count;
+    } TokPair;
+
+    TokPair *pairs = NULL;
+    size_t pairs_cap = 0;
+    size_t pairs_size = 0;
+
+    for (size_t word_idx = 0; word_idx < c->words_size; ++word_idx) {
+        uint16_t *word = &c->words[c->size_of_word * word_idx];
+
+        for (size_t i = 0; i < c->size_of_word - 1 && word[i + 1]; ++i) {
+            Tok tok = { word[i], word[i + 1] };
+
+            if (!pairs_size) {
+                TokPair pair = { tok, 1 };
+                array_push((void **)&pairs,
+                    &pairs_cap,
+                    &pairs_size,
+                    &pair,
+                    sizeof(pair));
+                continue;
+            }
+
+            size_t idx;
+            TokPair *found = binary_search(pairs,
+                pairs_size,
+                &tok,
+                sizeof(pairs[0]),
+                &idx,
+                (CmpFn)tok_cmp);
+
+            if (found) {
+                found->count += 1;
+                continue;
+            }
+
+            TokPair empty_pair = { 0 };
+            array_push((void **)&pairs,
+                &pairs_cap,
+                &pairs_size,
+                &empty_pair,
+                sizeof(empty_pair));
+
+            memmove(&pairs[idx + 1],
+                &pairs[idx],
+                sizeof(pairs[0]) * (pairs_size - idx - 1));
+
+            pairs[idx] = (TokPair) { tok, 1 };
+        }
+    }
+
+    TokPair *max_pair = &pairs[0];
+    for (size_t i = 0; i < pairs_size; ++i) {
+        if (pairs[i].count > max_pair->count)
+            max_pair = &pairs[i];
+    }
+
+    Tok max_tok = max_pair->tok;
+    free(pairs);
+    return max_tok;
+}
+
+void corpus_rewrite_with_tok(Corpus *c, Tok tok, Vocab *v)
+{
+    uint16_t tok_idx = vocab_add_tok(v, tok);
+
+    for (size_t word_idx = 0; word_idx < c->words_size; ++word_idx) {
+        uint16_t *word = &c->words[c->size_of_word * word_idx];
+
+        for (size_t i = 0; i < c->size_of_word - 1 && word[i + 1]; ++i) {
+            if (word[i] != tok.left || word[i + 1] != tok.right)
+                continue;
+
+            memmove(&word[i + 1],
+                &word[i + 2],
+                sizeof(uint16_t) * (c->size_of_word - i - 2));
+
+            word[i] = tok_idx;
+
+            // NOTE: since the token is newly introduced,
+            // we don't need to check it again,
+            // hence we don't set i back one here.
+        }
+    }
+}
+
+static void corpus_dump_word(Corpus *c, Vocab *v, size_t word_idx)
+{
+    uint16_t *word = &c->words[c->size_of_word * word_idx];
+
+    printf("[ ");
+    for (size_t i = 0; i < c->size_of_word; ++i) {
+        if (!word[i])
+            break;
+
+        vocab_print_tokval(v, stdout, word[i], true);
+        printf(" ");
+    }
+    printf("]");
+}
+
+void corpus_dump(Corpus *c, Vocab *v)
+{
+    size_t batch_size = 10;
+    for (size_t word_idx = 0;
+        word_idx < c->words_size && word_idx < batch_size * 10;
+        ++word_idx) {
+        corpus_dump_word(c, v, word_idx);
+        printf("\n");
+    }
+    printf("...\n");
+    for (size_t word_idx = c->words_size / 2 - batch_size / 2;
+        word_idx < c->words_size
+        && word_idx < c->words_size / 2 + batch_size / 2;
+        ++word_idx) {
+        corpus_dump_word(c, v, word_idx);
+        printf("\n");
+    }
+    printf("...\n");
+    // for (size_t word_idx = c->words_size - batch_size; word_idx <
+    // c->words_size;
+    //     ++word_idx) {
+    //     corpus_dump_word(c, v, word_idx);
+    //     printf("\n");
+    // }
+}
+
+int main(int argc, char *argv[])
+{
+    test_multiset();
+
+    if (argc <= 2) {
+        fprintf(stderr, "error: no files specified\n");
+        fprintf(
+            stderr, "usage: tokenizer <input text file> <output token file>\n");
+        return 1;
+    }
+
+    const char *input_filename = argv[1];
+    const char *output_filename = argv[2];
+
+    FILE *input_file = fopen(input_filename, "r");
+    if (!input_file) {
+        fprintf(stderr,
+            "error: could not open file '%s': %s\n",
+            input_filename,
+            strerror(errno));
+        return 1;
+    }
+
+    Multiset words = { NULL, 0, 0 };
+    size_t total_word_count = 0;
+    size_t max_word_len = 0;
+
+    char line[64];
+    size_t line_size = sizeof(line) / sizeof(line[0]);
+    size_t len;
+
+    printf("reading '%s'...\n", input_filename);
+
+    int ch = 0;
+    while ((len = next_word(line, line_size, input_file, &ch))) {
+        line[len] = '\0';
+        // printf("%s\n", line);
+
+        if (line[0] == '\n') {
+            assert(len == 1);
+            continue;
+        }
+
+        multiset_insert(&words, line);
+        total_word_count += 1;
+
+        if (len > max_word_len)
+            max_word_len = len;
+    }
+
+    fclose(input_file);
+
+    printf("total word count: %ld\n", total_word_count);
+    printf("unique word count: %ld\n", words.entries_size);
+    printf("max word length: %ld\n", max_word_len);
+
+    multiset_sort_wordcount(&words);
+
+    size_t selected_words_count = 50000;
+
+    struct {
+        char *data;
+        size_t cap;
+        size_t size;
+    } vocab_chars = { NULL, 0, 0 };
+
+    printf("discovering initial vocabulary...\n");
+
+    for (size_t i = 0; i < words.entries_size && i < selected_words_count;
+        ++i) {
+        MultisetEntry *entry = &words.entries[words.entries_size - 1 - i];
+
+        for (size_t j = 0; j < entry->len; ++j) {
+            char ch = entry->text[j];
+
+            size_t idx = 0;
+            char *found = binary_search(vocab_chars.data,
+                vocab_chars.size,
+                &ch,
+                sizeof(ch),
+                &idx,
+                (CmpFn)char_cmp);
+
+            if (found)
+                continue;
+
+            array_push((void **)&vocab_chars.data,
+                &vocab_chars.cap,
+                &vocab_chars.size,
+                &ch,
+                sizeof(ch));
+
+            if (vocab_chars.size == 1)
+                continue;
+
+            memmove(&vocab_chars.data[idx + 1],
+                &vocab_chars.data[idx],
+                sizeof(vocab_chars.data[0]) * (vocab_chars.size - idx - 1));
+
+            vocab_chars.data[idx] = ch;
+        }
+    }
+
+    printf("[");
+    for (size_t i = 0; i < vocab_chars.size; ++i)
+        printf("%c", vocab_chars.data[i]);
+    printf("]\n");
+
+    Vocab vocab = {
+        .chars = vocab_chars.data,
+        .chars_size = vocab_chars.size,
+        .toks = NULL,
+        .toks_cap = 0,
+        .toks_size = 0,
+    };
+
+    size_t word_size_aligned = (max_word_len + 3) & ~3ul;
+
+    size_t corpus_words_cap = selected_words_count;
+    uint16_t *corpus_words
+        = calloc(corpus_words_cap, word_size_aligned * sizeof(uint16_t));
+
+    Corpus corpus = {
+        .size_of_word = word_size_aligned,
+        .words = corpus_words,
+        .words_size = 0,
+    };
+
+    printf("using word size = %ld, word count = %ld\n",
+        corpus.size_of_word,
+        corpus.size_of_word);
+    printf("converting words to tokens...\n");
+
+    for (size_t word_idx = 0; word_idx < selected_words_count; ++word_idx) {
+
+        assert(word_idx < words.entries_size);
+
+        MultisetEntry entry = words.entries[words.entries_size - word_idx - 1];
+        assert(entry.len != 0 && entry.text[0]);
+
+        for (size_t i = 0; i < corpus.size_of_word && i < entry.len; ++i) {
+            corpus.words[word_idx * word_size_aligned + i]
+                = vocab_char_tokval(&vocab, entry.text[i]);
+        }
+
+        corpus.words_size += 1;
+    }
+
+    multiset_free(&words);
+
+    size_t comb_iters = 4000;
+
+    printf("combining tokens in %ld iterations...\n", comb_iters);
+
+    for (size_t j = 0; j < comb_iters / 100; ++j) {
+        for (size_t i = 0; i < 100; ++i) {
+            Tok tok = corpus_most_common_tok_pair(&corpus, &vocab);
+            corpus_rewrite_with_tok(&corpus, tok, &vocab);
+        }
+
+        printf("%ld / %ld, vocabulary size = %ld\n",
+            (j + 1) * 100,
+            comb_iters,
+            vocab.chars_size + vocab.toks_size + 1);
+    }
+
+    corpus_dump(&corpus, &vocab);
+
+    FILE *output_file = fopen(output_filename, "w");
+    if (!output_file) {
+        fprintf(stderr,
+            "error: could not open file '%s' for writing: %s\n",
+            output_filename,
+            strerror(errno));
+        return 1;
+    }
+
+    for (size_t i = 0; i < vocab_chars.size; ++i) {
+        fputc(vocab_chars.data[i], output_file);
+    }
+    fputc('\n', output_file);
+
+    for (size_t i = 0; i < vocab.toks_size; ++i) {
+        vocab_print_tokval(
+            &vocab, output_file, 1 + vocab.chars_size + i, false);
+        fputc('\n', output_file);
+    }
+
+    fclose(output_file);
+
+    free(vocab_chars.data);
+    free(vocab.toks);
+    free(corpus.words);
+}
+
+void test_multiset(void)
+{
+    {
+        Multiset m = { NULL, 0, 0 };
+
+        multiset_insert(&m, "a");
+        assert(m.entries_size == 1);
+        assert(strcmp(m.entries[0].text, "a") == 0);
+        assert(m.entries[0].count == 1);
+
+        multiset_insert(&m, "d");
+        assert(m.entries_size == 2);
+        assert(strcmp(m.entries[0].text, "a") == 0);
+        assert(m.entries[0].count == 1);
+        assert(strcmp(m.entries[1].text, "d") == 0);
+        assert(m.entries[1].count == 1);
+
+        multiset_insert(&m, "d");
+        assert(m.entries_size == 2);
+        assert(m.entries[1].count == 2);
+
+        multiset_insert(&m, "b");
+        assert(m.entries_size == 3);
+        assert(strcmp(m.entries[0].text, "a") == 0);
+        assert(strcmp(m.entries[1].text, "b") == 0);
+        assert(strcmp(m.entries[2].text, "d") == 0);
+
+        multiset_insert(&m, ".");
+        assert(m.entries_size == 4);
+        assert(strcmp(m.entries[0].text, ".") == 0);
+
+        multiset_insert(&m, "aa");
+        assert(m.entries_size == 5);
+        assert(strcmp(m.entries[4].text, "aa") == 0);
+
+        multiset_insert(&m, "c");
+        multiset_insert(&m, "c");
+        multiset_insert(&m, "c");
+        assert(m.entries_size == 6);
+        assert(strcmp(m.entries[3].text, "c") == 0);
+        assert(m.entries[3].count == 3);
+
+        multiset_free(&m);
+    }
+}

simple-ngram/profile_1.sh

@@ -0,0 +1,7 @@
+#!/bin/bash
+
+set -xe
+
+sudo perf record -F 1000 -g --call-graph dwarf -- $@
+
+

simple-ngram/profile_2.sh

@@ -0,0 +1,17 @@
+#!/bin/bash
+
+set -xe
+
+sudo chown $USER perf.data
+
+perf script | stackcollapse-perf.pl > out.perf-folded
+rm perf.data
+
+flamegraph.pl out.perf-folded > perf.svg
+rm out.perf-folded
+
+readlink -f perf.svg
+
+# firefox --file perf.svg
+
+

simple-ngram/run.sh

@@ -0,0 +1,4 @@
+#!/bin/bash
+
+gcc make_tokens.c -o make_tokens -O3 -g -Wall -Wextra && ./make_tokens dataset_lenin.txt tokens.txt
+

simple-ngram/train_model.c

@@ -0,0 +1,767 @@
+#include "common.h"
+#include <assert.h>
+#include <errno.h>
+#include <stdbool.h>
+#include <stddef.h>
+#include <stdint.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+
+#define seq_len 5
+
+typedef struct {
+    char ch;
+    uint16_t idx;
+} TokMatchOpt;
+
+typedef struct {
+    TokMatchOpt *opts;
+    uint16_t opts_cap;
+    uint16_t opts_size;
+    uint16_t id;
+} TokMatchNode;
+
+typedef struct {
+    TokMatchNode *nodes;
+    size_t nodes_cap;
+    size_t nodes_size;
+} TokMatcher;
+
+void tokmatcher_init(TokMatcher *m)
+{
+    *m = (TokMatcher) {
+        .nodes = NULL,
+        .nodes_cap = 0,
+        .nodes_size = 0,
+    };
+
+    TokMatchNode head = { NULL, 0, 0, 0 };
+
+    array_push(
+        (void **)&m->nodes, &m->nodes_cap, &m->nodes_size, &head, sizeof(head));
+}
+
+void tokmatcher_deinit(TokMatcher *m)
+{
+    for (size_t i = 0; i < m->nodes_size; ++i)
+        if (m->nodes[i].opts)
+            free(m->nodes[i].opts);
+    if (m->nodes)
+        free(m->nodes);
+}
+
+static TokMatchNode *tokmatcher_make_path(
+    TokMatcher *m, size_t node_idx, const char *key)
+{
+    TokMatchNode *node = &m->nodes[node_idx];
+    if (!*key)
+        return node;
+
+    for (size_t i = 0; i < node->opts_size; ++i)
+        if (node->opts[i].ch == *key)
+            return tokmatcher_make_path(m, node->opts[i].idx, key + 1);
+
+    size_t idx = m->nodes_size;
+
+    TokMatchNode new_node = { NULL, 0, 0, 0 };
+    array_push((void **)&m->nodes,
+        &m->nodes_cap,
+        &m->nodes_size,
+        &new_node,
+        sizeof(new_node));
+
+    // WARNING: m->nodes might be relocated here,
+    // hence we need to re-index node.
+    node = &m->nodes[node_idx];
+
+    assert(idx <= UINT16_MAX);
+    TokMatchOpt new_opt = {
+        .ch = *key,
+        .idx = (uint16_t)idx,
+    };
+
+    size_t opts_cap = node->opts_cap;
+    size_t opts_size = node->opts_size;
+    array_push(
+        (void **)&node->opts, &opts_cap, &opts_size, &new_opt, sizeof(new_opt));
+
+    assert(opts_cap <= UINT16_MAX && opts_size <= UINT16_MAX);
+    node->opts_cap = opts_cap;
+    node->opts_size = opts_size;
+
+    return tokmatcher_make_path(m, idx, key + 1);
+}
+
+void tokmatcher_add(TokMatcher *m, const char *key, uint16_t id)
+{
+    TokMatchNode *node = tokmatcher_make_path(m, 0, key);
+    node->id = id;
+}
+
+typedef struct {
+    uint32_t len;
+    uint16_t id;
+} TokRecMatch;
+
+static TokRecMatch tokmatcher_match_rec(TokMatcher *m,
+    TokMatchNode *node,
+    TokRecMatch cand,
+    size_t len,
+    const char *text,
+    size_t text_len)
+{
+    if (node->id != 0) {
+        cand = (TokRecMatch) { len, node->id };
+    }
+
+    if (!*text || text_len == 0)
+        return cand;
+
+    for (size_t i = 0; i < node->opts_size; ++i) {
+        if (node->opts[i].ch != *text)
+            continue;
+
+        TokMatchNode *next = &m->nodes[node->opts[i].idx];
+
+        return tokmatcher_match_rec(
+            m, next, cand, len + 1, text + 1, text_len - 1);
+    }
+
+    return cand;
+}
+
+uint16_t tokmatcher_match(
+    TokMatcher *m, const char *text, size_t text_len, size_t *chars_matched)
+{
+    TokRecMatch cand = { .len = 1, .id = 0 };
+
+    TokRecMatch match
+        = tokmatcher_match_rec(m, &m->nodes[0], cand, 0, text, text_len);
+
+    *chars_matched += match.len;
+    return match.id;
+}
+
+static void tokmatcher_dump_rec(
+    TokMatcher *m, TokMatchNode *node, FILE *file, int depth)
+{
+    fprintf(file, " -> %d", node->id);
+
+    for (size_t i = 0; i < node->opts_size; ++i) {
+        fprintf(file, "\n");
+        for (int j = 0; j < depth + 1; ++j)
+            fprintf(file, "       | ");
+
+        fprintf(file, "'%c'", node->opts[i].ch);
+
+        tokmatcher_dump_rec(m, &m->nodes[node->opts[i].idx], file, depth + 1);
+    }
+}
+
+void tokmatcher_dump(TokMatcher *m, FILE *file)
+{
+    printf(" * ");
+    tokmatcher_dump_rec(m, &m->nodes[0], file, 0);
+}
+
+typedef enum {
+    SpecialTok_Unk,
+    SpecialTok_Begin,
+    SpecialTok_End,
+    SpecialTokCount,
+} SpecialTok;
+
+typedef struct {
+    char *toks_strs;
+    size_t toks_strs_cap;
+    size_t toks_strs_size;
+
+    size_t *toks;
+    size_t toks_cap;
+    size_t toks_size;
+
+    TokMatcher tok_matcher;
+    uint16_t special_toks[SpecialTokCount];
+
+    int next_word_ch;
+
+    FILE *file;
+    char *word;
+    size_t word_cap;
+    size_t word_len;
+    size_t chars_matched;
+} Tokenizer;
+
+static uint16_t t_add_token_str(Tokenizer *t, const char *str)
+{
+    size_t tok_str_idx = t->toks_strs_size;
+
+    array_push_n((void **)&t->toks_strs,
+        &t->toks_strs_cap,
+        &t->toks_strs_size,
+        (void *)str,
+        1,
+        strlen(str) + 1);
+
+    array_push((void **)&t->toks,
+        &t->toks_cap,
+        &t->toks_size,
+        &tok_str_idx,
+        sizeof(tok_str_idx));
+
+    assert(tok_str_idx <= UINT16_MAX);
+    uint16_t id = tok_str_idx;
+    return id;
+}
+
+void tokenizer_init(Tokenizer *t)
+{
+    *t = (Tokenizer) {
+        .toks_strs = NULL,
+        .toks_strs_cap = 0,
+        .toks_strs_size = 0,
+
+        .toks = NULL,
+        .toks_cap = 0,
+        .toks_size = 0,
+
+        .tok_matcher = { 0 },
+        .special_toks = { 0 },
+
+        .next_word_ch = 0,
+
+        .file = NULL,
+        .word = NULL,
+        .word_cap = 0,
+        .word_len = 0,
+        .chars_matched = 0,
+    };
+
+    tokmatcher_init(&t->tok_matcher);
+
+    t->special_toks[SpecialTok_Unk] = t_add_token_str(t, "<unk>");
+    t->special_toks[SpecialTok_Begin] = t_add_token_str(t, "<s>");
+    t->special_toks[SpecialTok_End] = t_add_token_str(t, "</s>");
+}
+
+void tokenizer_deinit(Tokenizer *t)
+{
+    if (t->toks_strs)
+        free(t->toks_strs);
+    if (t->toks)
+        free(t->toks);
+    tokmatcher_deinit(&t->tok_matcher);
+    if (t->file)
+        fclose(t->file);
+    if (t->word)
+        free(t->word);
+}
+
+int tokenizer_load_tokens_file(Tokenizer *t, const char *filename)
+{
+    FILE *file = fopen(filename, "r");
+    if (!file) {
+        fprintf(stderr,
+            "error: could not open file '%s': %s\n",
+            filename,
+            strerror(errno));
+        return 1;
+    }
+
+    char chars_line[34] = "";
+    fgets(chars_line, 34, file);
+    size_t chars_line_len = strlen(chars_line);
+
+    // This is just a sanity check that 1) we have read a
+    // correct file, and 2) that the chars are what we expect.
+    assert(strcmp(chars_line, " !',.?abcdefghijklmnopqrstuvwxyz"));
+    assert(chars_line_len == 33);
+
+    for (size_t i = 0; i < chars_line_len; ++i) {
+        char str[2] = { chars_line[i], '\0' };
+
+        uint16_t id = t_add_token_str(t, str);
+        tokmatcher_add(&t->tok_matcher, str, id);
+    }
+
+    while (!feof(file)) {
+        char line[41] = "";
+        fgets(line, 41, file);
+        size_t line_len = strlen(line);
+        if (line_len == 0)
+            break;
+        assert(line[line_len - 1] == '\n');
+        line[line_len - 1] = '\0';
+
+        uint16_t id = t_add_token_str(t, line);
+        tokmatcher_add(&t->tok_matcher, line, id);
+    }
+
+    fclose(file);
+
+    return 0;
+}
+
+int tokenizer_open(Tokenizer *t, const char *filename)
+{
+    FILE *file = fopen(filename, "r");
+    if (!file) {
+        fprintf(stderr,
+            "error: could not open file '%s': %s\n",
+            filename,
+            strerror(errno));
+        return 1;
+    }
+
+    t->file = file;
+
+    t->word_cap = 64;
+    t->word = malloc(t->word_cap);
+    t->word_len = 0;
+    t->chars_matched = 0;
+
+    return 0;
+}
+
+int32_t tokenizer_next(Tokenizer *t)
+{
+    if (t->chars_matched >= t->word_len) {
+        size_t len = next_word(t->word, t->word_cap, t->file, &t->next_word_ch);
+
+        if (len == 0)
+            return -1;
+        if (t->word[0] == '\n')
+            return -2;
+
+        t->word[len] = '\0';
+        t->word_len = len;
+        t->chars_matched = 0;
+    }
+
+    uint16_t id = tokmatcher_match(&t->tok_matcher,
+        t->word + t->chars_matched,
+        t->word_len - t->chars_matched,
+        &t->chars_matched);
+
+    return id;
+}
+
+typedef struct {
+    uint16_t tok;
+    uint32_t node_idx;
+} SeqIdxMapOpt;
+
+int seqidxmapopt_cmp(const SeqIdxMapOpt *a, const SeqIdxMapOpt *b)
+{
+    if (a->tok != b->tok)
+        return (int)a->tok - (int)b->tok;
+    if (a->node_idx != b->node_idx)
+        return a->node_idx < b->node_idx ? 1 : -1;
+    return 0;
+}
+
+typedef struct {
+    uint32_t val;
+    uint32_t opts_cap_exp : 8, opts_size : 24;
+    SeqIdxMapOpt opts[];
+} SeqIdxMapNode;
+
+typedef struct {
+    SeqIdxMapNode **nodes;
+    size_t nodes_size;
+    size_t nodes_cap;
+} SeqIdxMap;
+
+void seqidxmap_init(SeqIdxMap *m)
+{
+    *m = (SeqIdxMap) {
+        .nodes = NULL,
+        .nodes_cap = 0,
+        .nodes_size = 0,
+    };
+
+    SeqIdxMapNode *head
+        = malloc(sizeof(SeqIdxMapNode) + sizeof(SeqIdxMapOpt) * 1);
+    *head = (SeqIdxMapNode) { .val = 0, .opts_cap_exp = 0, .opts_size = 0 };
+    array_push((void **)&m->nodes,
+        &m->nodes_cap,
+        &m->nodes_size,
+        &head,
+        sizeof(SeqIdxMapNode *));
+}
+
+void seqidxmap_deinit(SeqIdxMap *m)
+{
+    for (size_t i = 0; i < m->nodes_size; ++i)
+        free(m->nodes[i]);
+    free(m->nodes);
+}
+
+uint32_t *seqidxmap_get(SeqIdxMap *m, uint16_t seq[seq_len])
+{
+    size_t node_idx = 0;
+
+    for (size_t seq_idx = 0; seq_idx < seq_len; ++seq_idx) {
+        assert(node_idx < m->nodes_size);
+        SeqIdxMapNode *node = m->nodes[node_idx];
+
+        size_t opt_idx;
+        SeqIdxMapOpt *opt = binary_search(&node->opts,
+            node->opts_size,
+            &seq[seq_idx],
+            sizeof(uint16_t),
+            &opt_idx,
+            (CmpFn)seqidxmapopt_cmp);
+
+        if (opt) {
+            node_idx = opt->node_idx;
+            continue;
+        }
+
+        if (node->opts_size >= (1 << node->opts_cap_exp)) {
+            assert(node->opts_cap_exp >= 0 && node->opts_cap_exp <= 255);
+            node->opts_cap_exp += 1;
+
+            SeqIdxMapNode *new_ptr = realloc(m->nodes[node_idx],
+                sizeof(SeqIdxMapNode)
+                    + sizeof(SeqIdxMapOpt) * (1ul << node->opts_cap_exp));
+            assert(new_ptr);
+            m->nodes[node_idx] = new_ptr;
+        }
+
+        size_t new_node_idx = m->nodes_size;
+        SeqIdxMapNode *new_node
+            = malloc(sizeof(SeqIdxMapNode) + sizeof(SeqIdxMapOpt) * 1);
+        *new_node
+            = (SeqIdxMapNode) { .val = 0, .opts_cap_exp = 0, .opts_size = 0 };
+        array_push((void **)&m->nodes,
+            &m->nodes_cap,
+            &m->nodes_size,
+            &new_node,
+            sizeof(SeqIdxMapNode *));
+
+        assert(node_idx < m->nodes_size);
+        node = m->nodes[node_idx];
+
+        assert(node->opts_size < (1 << 24) - 1);
+        node->opts[node->opts_size] = (SeqIdxMapOpt) {
+            .node_idx = new_node_idx,
+            .tok = seq[seq_idx],
+        };
+        node->opts_size += 1;
+
+        node_idx = new_node_idx;
+    }
+
+    return &m->nodes[node_idx]->val;
+}
+
+int seq_cmp(const uint16_t a[seq_len], const uint16_t b[seq_len])
+{
+    for (size_t i = 0; i < seq_len; ++i)
+        if (a[i] != b[i])
+            return (int)a[i] - (int)b[i];
+    return 0;
+}
+
+typedef struct {
+    uint32_t seq_idx;
+    uint16_t id;
+    uint16_t count;
+} SeqPred;
+
+int seqpred_cmp(const SeqPred *a, const SeqPred *b)
+{
+    if (a->seq_idx != b->seq_idx)
+        return a->seq_idx < b->seq_idx ? 1 : -1;
+    if (a->id != b->id)
+        return (int)a->id - (int)b->id;
+    return 0;
+}
+
+typedef uint64_t SeqPredKey;
+
+SeqPredKey seqpredkey_from_seqpred(const SeqPred *pred)
+{
+    SeqPredKey key = (uint64_t)pred->seq_idx | (uint64_t)pred->id << 32ul;
+    return key;
+}
+
+int seqpred_key_cmp(const SeqPredKey *a, const SeqPredKey *b)
+{
+    if (*a != *b)
+        return *a < *b ? 1 : -1;
+    return 0;
+}
+
+typedef struct {
+    uint16_t *seqs;
+    size_t seqs_cap;
+    size_t seqs_size;
+
+    SeqIdxMap seqidxmap;
+
+    SeqPred *preds;
+    size_t preds_cap;
+    size_t preds_size;
+
+    RbTree *pred_map;
+} Model;
+
+void model_init(Model *m)
+{
+    *m = (Model) {
+        .seqs = NULL,
+        .seqs_cap = 0,
+        .seqs_size = 0,
+
+        .seqidxmap = { 0 },
+
+        .preds = NULL,
+        .preds_cap = 0,
+        .preds_size = 0,
+
+        .pred_map = rbtree_new(
+            sizeof(SeqPredKey), sizeof(SeqPred *), (CmpFn)seqpred_key_cmp),
+    };
+
+    seqidxmap_init(&m->seqidxmap);
+}
+
+void model_deinit(Model *m)
+{
+    if (m->seqs)
+        free(m->seqs);
+    seqidxmap_deinit(&m->seqidxmap);
+
+    if (m->preds)
+        free(m->preds);
+
+    rbtree_free(m->pred_map);
+}
+
+size_t model_get_seq_or_insert(Model *m, uint16_t seq[seq_len])
+{
+    uint32_t *idx = seqidxmap_get(&m->seqidxmap, seq);
+
+    if (*idx)
+        return *idx;
+
+    assert(m->seqs_size <= UINT32_MAX);
+    uint32_t new_idx = m->seqs_size;
+
+    array_push((void **)&m->seqs,
+        &m->seqs_cap,
+        &m->seqs_size,
+        seq,
+        sizeof(uint16_t) * seq_len);
+
+    *idx = new_idx;
+
+    return new_idx;
+}
+
+uint32_t *model_get_seq(Model *m, uint16_t seq[seq_len])
+{
+    return seqidxmap_get(&m->seqidxmap, seq);
+}
+
+SeqPred *model_find_seqpred(Model *m, SeqPred *pred)
+{
+    SeqPredKey key = (uint64_t)pred->seq_idx | (uint64_t)pred->id << 32ul;
+    return rbtree_get(m->pred_map, &key);
+}
+
+SeqPred *model_insert_seqpred(Model *m, SeqPred *pred)
+{
+    array_push(
+        (void **)&m->preds, &m->preds_cap, &m->preds_size, pred, sizeof(*pred));
+
+    SeqPred *ptr = &m->preds[m->preds_size - 1];
+
+    SeqPredKey key = seqpredkey_from_seqpred(pred);
+    rbtree_insert(m->pred_map, &key, &ptr);
+
+    return ptr;
+}
+
+int main(int argc, char *argv[])
+{
+    test_common();
+    if (argc < 2) {
+        fprintf(stderr, "error: invalid arguments\n");
+        fprintf(stderr, "usage: ./train_model <training text> <tokens>\n");
+        return 1;
+    }
+
+    const char *training_text_filename = argv[1];
+    const char *tokens_filename = argv[2];
+
+    Tokenizer tokenizer;
+    tokenizer_init(&tokenizer);
+
+    printf("reading tokens...\n");
+    int result = tokenizer_load_tokens_file(&tokenizer, tokens_filename);
+    if (result != 0)
+        return 1;
+
+    result = tokenizer_open(&tokenizer, training_text_filename);
+    if (result != 0)
+        return 1;
+
+    printf("training...\n");
+
+    Model model;
+    model_init(&model);
+
+    const uint16_t begin_tok = tokenizer.special_toks[SpecialTok_Begin];
+    const uint16_t end_tok = tokenizer.special_toks[SpecialTok_End];
+
+    uint16_t current_seq[seq_len] = {
+        begin_tok,
+        begin_tok,
+        begin_tok,
+        begin_tok,
+        begin_tok,
+    };
+
+    int toks_counter = 0;
+    int iters = 0;
+
+    int32_t id;
+    while ((id = tokenizer_next(&tokenizer)) != -1) {
+        id = id == -2 ? end_tok : id;
+
+        size_t seq_idx = model_get_seq_or_insert(&model, current_seq);
+
+        SeqPred cand_seq = { .seq_idx = seq_idx, .id = id, .count = 0 };
+
+        SeqPred *pred = model_find_seqpred(&model, &cand_seq);
+
+        if (!pred) {
+            pred = model_insert_seqpred(&model, &cand_seq);
+        }
+
+        pred->count += 1;
+
+        if (id == end_tok) {
+
+            for (size_t i = 0; i < seq_len; ++i)
+                current_seq[i] = begin_tok;
+        } else {
+            for (size_t i = 0; i < seq_len - 1; ++i)
+                current_seq[i] = current_seq[i + 1];
+            current_seq[seq_len - 1] = id;
+        }
+
+        toks_counter += 1;
+        iters += 1;
+        if (iters >= 1000000) {
+            iters = 0;
+
+            printf(
+                "  toks_counter: %d, seqs: %ld, seq nodes: %ld, preds: %ld\n",
+                toks_counter,
+                model.seqs_size,
+                model.seqidxmap.nodes_size,
+                model.preds_size);
+            break;
+        }
+    }
+
+    printf("done!\n\n");
+
+    for (size_t j = 0; j < model.preds_size; j += 100) {
+        for (size_t i = j; i < j + 100; ++i) {
+            SeqPred *pred = &model.preds[i];
+            uint16_t *seq = &model.seqs[pred->seq_idx];
+            printf(" '");
+            for (size_t i = 0; i < seq_len; ++i) {
+                printf("%s", &tokenizer.toks_strs[seq[i]]);
+            }
+            printf("' -> %d x '%s'\n",
+                pred->count,
+                &tokenizer.toks_strs[pred->id]);
+        }
+        getc(stdin);
+    }
+
+    //
+    // printf("> ");
+    //
+    // for (size_t i = 0; i < seq_len; ++i)
+    //     current_seq[i] = begin_tok;
+    //
+    // {
+    //     char word[64] = "";
+    //     size_t word_cap = sizeof(word);
+    //     size_t word_len = 0;
+    //     size_t chars_matched = 0;
+    //
+    //     int ch = 0;
+    //     while (true) {
+    //         if (chars_matched >= word_len) {
+    //             size_t len = next_word(word, word_cap, stdin, &ch);
+    //
+    //             if (len == 0)
+    //                 break;
+    //             if (word[0] == '\n')
+    //                 break;
+    //
+    //             word[len] = '\0';
+    //             word_len = len;
+    //             chars_matched = 0;
+    //         }
+    //
+    //         uint16_t id = tokmatcher_match(&tokenizer.tok_matcher,
+    //             word + chars_matched,
+    //             word_len - chars_matched,
+    //             &chars_matched);
+    //
+    //         for (size_t i = 0; i < seq_len - 1; ++i)
+    //             current_seq[i] = current_seq[i + 1];
+    //         current_seq[seq_len - 1] = id;
+    //     }
+    // }
+    //
+    // printf("generating prompt...\n");
+    //
+    // printf("\n");
+    //
+    // for (size_t i = 0; i < seq_len; ++i) {
+    //     printf("%s", &tokenizer.toks_strs[current_seq[i]]);
+    // }
+    //
+    // for (int iters = 0; iters < 1000; ++iters) {
+    //
+    //     uint32_t *seq_idx = model_get_seq(&model, current_seq);
+    //     if (!seq_idx) {
+    //         printf("no sequence\n");
+    //         break;
+    //     }
+    //
+    //     SeqPred *pred = NULL;
+    //
+    //     for (size_t i = 0; i < model.preds_size; ++i) {
+    //         if (model.preds[i].seq_idx != *seq_idx)
+    //             continue;
+    //         if (!pred || model.preds[i].count > pred->count)
+    //             pred = &model.preds[i];
+    //     }
+    //
+    //     if (!pred) {
+    //         printf("no prediction\n");
+    //         break;
+    //     }
+    //
+    //     printf("%s", &tokenizer.toks_strs[pred->id]);
+    //
+    //     for (size_t i = 0; i < seq_len - 1; ++i)
+    //         current_seq[i] = current_seq[i + 1];
+    //     current_seq[seq_len - 1] = pred->id;
+    // }
+
+    tokenizer_deinit(&tokenizer);
+    model_deinit(&model);
+}