tiny_gen.c

/*
 * tiny_gen.c - Tiny generative LLM for L1VM code generation.
 *
 * A small autoregressive next-token model. The transformer body is a
 * frozen random projection (the same trick the classifier uses); the
 * token/position embeddings and a 2-layer readout head are trained with
 * teacher forcing so the model memorizes (prompt -> program) pairs and
 * can re-generate valid L1VM code for prompts close to its training set.
 *
 * (c) Copyright Stefan Pietzonke (info@midnight-coding.de), 2026
 *
 * This file is part of brackets-code.
 */

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <ctype.h>
#include <math.h>

#include "tiny_gen.h"

/* ==================== Small helpers ==================== */

static void tgn_tolower(char *s) {
    for (; *s; s++) *s = (char)tolower((unsigned char)*s);
}

/* ==================== Vocabulary ==================== */

void tgn_vocab_init(TgnVocab *v) {
    memset(v, 0, sizeof(TgnVocab));
}

int tgn_vocab_add(TgnVocab *v, const char *word) {
    if (v->count >= TGN_MAX_VOCAB) return -1;
    for (int i = 0; i < v->count; i++) {
        if (strcmp(v->words[i], word) == 0) return i;
    }
    snprintf(v->words[v->count], 64, "%s", word);
    return v->count++;
}

int tgn_vocab_find(const TgnVocab *v, const char *word) {
    for (int i = 0; i < v->count; i++) {
        if (strcmp(v->words[i], word) == 0) return i;
    }
    return -1;
}

/* Word-level tokenizer: splits on runs of spaces/tabs, keeps "\n" as a
   token, lowercases everything. This is lossless for L1VM source because
   L1VM is whitespace-delimited with parens attached to tokens. */
int tgn_tokenize(const TgnVocab *v, const char *text, int *tokens, int max_tokens) {
    char buf[TGN_RAW_CODE];
    snprintf(buf, sizeof(buf), "%s", text);
    tgn_tolower(buf);

    char cur[128];
    int cur_len = 0;
    int n = 0;

    for (const char *p = buf; ; p++) {
        char c = *p;
        if (c == '\n' || c == '\0' || c == ' ' || c == '\t' || c == '\r') {
            if (cur_len > 0) {
                cur[cur_len] = '\0';
                int id = tgn_vocab_find(v, cur);
                if (n < max_tokens) tokens[n] = (id < 0) ? TGN_UNK : id;
                n++;
                cur_len = 0;
            }
            if (c == '\n') {
                int id = tgn_vocab_find(v, "\n");
                if (n < max_tokens) tokens[n] = (id < 0) ? TGN_UNK : id;
                n++;
            }
            if (c == '\0') break;
        } else {
            if (cur_len < (int)sizeof(cur) - 1) cur[cur_len++] = c;
        }
    }
    return n;
}

int tgn_vocab_save(const TgnVocab *v, const char *path) {
    char vocab_path[1024];
    snprintf(vocab_path, sizeof(vocab_path), "%s.vocab", path);
    FILE *f = fopen(vocab_path, "wb");
    if (!f) return -1;

    int32_t magic = 0x564F4342; /* "VOCB" */
    fwrite(&magic, sizeof(int32_t), 1, f);
    fwrite(&v->count, sizeof(int32_t), 1, f);
    for (int i = 0; i < v->count; i++) {
        int32_t len = (int32_t)strlen(v->words[i]);
        fwrite(&len, sizeof(int32_t), 1, f);
        fwrite(v->words[i], 1, (size_t)len, f);
    }
    fclose(f);
    return 0;
}

int tgn_vocab_load(TgnVocab *v, const char *path) {
    char vocab_path[1024];
    snprintf(vocab_path, sizeof(vocab_path), "%s.vocab", path);
    FILE *f = fopen(vocab_path, "rb");
    if (!f) return -1;

    int32_t magic;
    fread(&magic, sizeof(int32_t), 1, f);
    if (magic != 0x564F4342) { fclose(f); return -1; }

    int32_t count;
    fread(&count, sizeof(int32_t), 1, f);
    if (count > TGN_MAX_VOCAB) { fclose(f); return -1; }

    v->count = 0;
    for (int i = 0; i < count; i++) {
        int32_t len;
        fread(&len, sizeof(int32_t), 1, f);
        if (len > 63) len = 63;
        fread(v->words[i], 1, (size_t)len, f);
        v->words[i][len] = '\0';
        v->count++;
    }
    fclose(f);
    return 0;
}

/* ==================== Model ==================== */

TinyGenModel *tgn_model_create(const TgnVocab *v) {
    TinyGenModel *m = (TinyGenModel *)calloc(1, sizeof(TinyGenModel));
    if (!m) return NULL;

    m->vocab_size = v->count;
    m->embed = TGN_EMBED;
    m->max_seq = TGN_MAX_SEQ;
    m->hidden = TGN_HIDDEN;

    m->token_emb = mat_create(m->vocab_size, m->embed);
    m->pos_emb = mat_create(m->max_seq, m->embed);
    m->head_w = mat_create(m->embed, m->hidden);
    m->head_b = mat_create(1, m->hidden);
    m->out_w = mat_create(m->vocab_size, m->hidden);
    m->out_b = mat_create(1, m->vocab_size);

    float tscale = 0.5f;
    for (int i = 0; i < m->token_emb->rows * m->token_emb->cols; i++)
        m->token_emb->data[i] = ((float)rand() / (float)RAND_MAX - 0.5f) * tscale;
    for (int i = 0; i < m->pos_emb->rows * m->pos_emb->cols; i++)
        m->pos_emb->data[i] = ((float)rand() / (float)RAND_MAX - 0.5f) * 0.5f;
    float hscale = 0.5f;
    for (int i = 0; i < m->head_w->rows * m->head_w->cols; i++)
        m->head_w->data[i] = ((float)rand() / (float)RAND_MAX - 0.5f) * hscale;
    for (int i = 0; i < m->head_b->rows * m->head_b->cols; i++)
        m->head_b->data[i] = 0.0f;
    float oscale = sqrtf(2.0f / (float)m->hidden);
    for (int i = 0; i < m->out_w->rows * m->out_w->cols; i++)
        m->out_w->data[i] = ((float)rand() / (float)RAND_MAX - 0.5f) * oscale;
    for (int i = 0; i < m->out_b->rows * m->out_b->cols; i++)
        m->out_b->data[i] = 0.0f;

    for (int i = 0; i < TGN_NUM_LAYERS; i++)
        m->blocks[i] = block_create();

    return m;
}

void tgn_model_free(TinyGenModel *m) {
    if (!m) return;
    mat_free(m->token_emb);
    mat_free(m->pos_emb);
    mat_free(m->head_w);
    mat_free(m->head_b);
    mat_free(m->out_w);
    mat_free(m->out_b);
    for (int i = 0; i < TGN_NUM_LAYERS; i++)
        block_free(m->blocks[i]);
    free(m);
}

int tgn_model_save(const TinyGenModel *m, const char *path) {
    FILE *f = fopen(path, "wb");
    if (!f) return -1;

    int32_t magic = 0x54474E31; /* "TGN1" */
    fwrite(&magic, sizeof(int32_t), 1, f);
    int32_t v = m->vocab_size, e = m->embed, s = m->max_seq, h = m->hidden, nl = TGN_NUM_LAYERS;
    fwrite(&v, sizeof(int32_t), 1, f);
    fwrite(&e, sizeof(int32_t), 1, f);
    fwrite(&s, sizeof(int32_t), 1, f);
    fwrite(&h, sizeof(int32_t), 1, f);
    fwrite(&nl, sizeof(int32_t), 1, f);

    fwrite(m->token_emb->data, sizeof(float), (size_t)m->vocab_size * m->embed, f);
    fwrite(m->pos_emb->data, sizeof(float), (size_t)m->max_seq * m->embed, f);
    fwrite(m->head_w->data, sizeof(float), (size_t)m->embed * m->hidden, f);
    fwrite(m->head_b->data, sizeof(float), (size_t)m->hidden, f);
    fwrite(m->out_w->data, sizeof(float), (size_t)m->hidden * m->vocab_size, f);
    fwrite(m->out_b->data, sizeof(float), (size_t)m->vocab_size, f);

    for (int i = 0; i < TGN_NUM_LAYERS; i++) {
        TransformerBlock *b = m->blocks[i];
        fwrite(b->ln1_gamma->data, sizeof(float), TGN_EMBED, f);
        fwrite(b->ln1_beta->data, sizeof(float), TGN_EMBED, f);
        fwrite(b->ln2_gamma->data, sizeof(float), TGN_EMBED, f);
        fwrite(b->ln2_beta->data, sizeof(float), TGN_EMBED, f);
        fwrite(b->attn->wq->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fwrite(b->attn->wk->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fwrite(b->attn->wv->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fwrite(b->attn->wo->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fwrite(b->ff->w1->data, sizeof(float), TGN_EMBED * TT_FF_DIM, f);
        fwrite(b->ff->b1->data, sizeof(float), TT_FF_DIM, f);
        fwrite(b->ff->w2->data, sizeof(float), TT_FF_DIM * TGN_EMBED, f);
        fwrite(b->ff->b2->data, sizeof(float), TGN_EMBED, f);
    }

    fclose(f);
    return 0;
}

int tgn_model_load(TinyGenModel *m, const char *path) {
    FILE *f = fopen(path, "rb");
    if (!f) return -1;

    int32_t magic;
    fread(&magic, sizeof(int32_t), 1, f);
    if (magic != 0x54474E31) { fclose(f); return -1; }

    int32_t v, e, s, h, nl;
    fread(&v, sizeof(int32_t), 1, f);
    fread(&e, sizeof(int32_t), 1, f);
    fread(&s, sizeof(int32_t), 1, f);
    fread(&h, sizeof(int32_t), 1, f);
    fread(&nl, sizeof(int32_t), 1, f);

    if (v != m->vocab_size || e != m->embed || s != m->max_seq || h != m->hidden) {
        fclose(f);
        return -1;
    }

    fread(m->token_emb->data, sizeof(float), (size_t)m->vocab_size * m->embed, f);
    fread(m->pos_emb->data, sizeof(float), (size_t)m->max_seq * m->embed, f);
    fread(m->head_w->data, sizeof(float), (size_t)m->embed * m->hidden, f);
    fread(m->head_b->data, sizeof(float), (size_t)m->hidden, f);
    fread(m->out_w->data, sizeof(float), (size_t)m->hidden * m->vocab_size, f);
    fread(m->out_b->data, sizeof(float), (size_t)m->vocab_size, f);

    for (int i = 0; i < TGN_NUM_LAYERS && i < nl; i++) {
        TransformerBlock *b = m->blocks[i];
        fread(b->ln1_gamma->data, sizeof(float), TGN_EMBED, f);
        fread(b->ln1_beta->data, sizeof(float), TGN_EMBED, f);
        fread(b->ln2_gamma->data, sizeof(float), TGN_EMBED, f);
        fread(b->ln2_beta->data, sizeof(float), TGN_EMBED, f);
        fread(b->attn->wq->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fread(b->attn->wk->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fread(b->attn->wv->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fread(b->attn->wo->data, sizeof(float), TGN_EMBED * TGN_EMBED, f);
        fread(b->ff->w1->data, sizeof(float), TGN_EMBED * TT_FF_DIM, f);
        fread(b->ff->b1->data, sizeof(float), TT_FF_DIM, f);
        fread(b->ff->w2->data, sizeof(float), TT_FF_DIM * TGN_EMBED, f);
        fread(b->ff->b2->data, sizeof(float), TGN_EMBED, f);
    }

    fclose(f);
    return 0;
}

/* ==================== Training data collection ==================== */

static void tgn_trim_line(char *s) {
    char *end = s + strlen(s);
    while (end > s && (end[-1] == '\n' || end[-1] == '\r' || end[-1] == ' ' || end[-1] == '\t'))
        end--;
    *end = '\0';
    char *p = s;
    while (*p == ' ' || *p == '\t') p++;
    if (p != s) memmove(s, p, strlen(p) + 1);
}

static int tgn_is_dsl_key(const char *line) {
    static const char *keys[] = {
        "parser:", "token:", "result:", "include:", "include-post:", "var:", "desc:",
        "array:", "param:", "require:", "example:", "category:", "version:",
        "complexity:", "alias:", "test:", "help:", "validate:", "compose:", NULL
    };
    for (int i = 0; keys[i]; i++) {
        if (strncmp(line, keys[i], strlen(keys[i])) == 0) return 1;
    }
    return 0;
}

int tgn_collect_dsl(TgnRawPair *pairs, int max_pairs, const char *dsl_dir) {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "ls %s/*.l1dsl 2>/dev/null", dsl_dir);
    FILE *pipe = popen(cmd, "r");
    if (!pipe) return 0;

    char filepath[512];
    int num = 0;

    while (fgets(filepath, sizeof(filepath), pipe) && num < max_pairs) {
        tgn_trim_line(filepath);
        if (filepath[0] == '\0') continue;

        FILE *f = fopen(filepath, "r");
        if (!f) continue;

        char parser[512] = "";
        char code[TGN_RAW_CODE] = "";
        int in_code = 0;

        char line[4096];
        while (fgets(line, sizeof(line), f)) {
            tgn_trim_line(line);
            if (strcmp(line, "code:") == 0 || strcmp(line, "code:|") == 0) {
                in_code = 1;
                continue;
            }
            if (in_code) {
                if (line[0] == '\0' || tgn_is_dsl_key(line) ||
                    (strncmp(line, "//", 2) == 0 && line[2] != '/')) {
                    in_code = 0;
                } else if (line[0] == '/' && line[1] == '/') {
                    /* inline comment inside code block: skip */
                    continue;
                } else {
                    if (strlen(code) + strlen(line) + 1 < sizeof(code)) {
                        if (code[0]) strncat(code, "\n", sizeof(code) - strlen(code) - 1);
                        strncat(code, line, sizeof(code) - strlen(code) - 1);
                    }
                }
            } else if (strncmp(line, "parser:", 7) == 0) {
                snprintf(parser, sizeof(parser), "%s", line + 7);
            }
        }
        fclose(f);

        /* Skip if no code was collected */
        if (code[0] == '\0') continue;

        /* One pair per comma-separated parser variant */
        char *save = NULL;
        char *tok = strtok_r(parser, ",", &save);
        int variants = 0;
        while (tok && num < max_pairs) {
            char *pv = tok;
            while (*pv == ' ' || *pv == '\t' || *pv == '"') pv++;
            char *end = pv + strlen(pv);
            while (end > pv && (end[-1] == ' ' || end[-1] == '\t' || end[-1] == '"'))
                end--;
            *end = '\0';
            if (pv[0] && variants < 6) {
                snprintf(pairs[num].prompt, sizeof(pairs[num].prompt), "%s", pv);
                snprintf(pairs[num].code, sizeof(pairs[num].code), "%s", code);
                num++;
                variants++;
            }
            tok = strtok_r(NULL, ",", &save);
        }
    }

    pclose(pipe);
    return num;
}

int tgn_collect_examples(TgnRawPair *pairs, int max_pairs, const char *examples_dir) {
    char cmd[1024];
    snprintf(cmd, sizeof(cmd), "ls %s/*.l1com 2>/dev/null", examples_dir);
    FILE *pipe = popen(cmd, "r");
    if (!pipe) return 0;

    char filepath[512];
    int num = 0;

    while (fgets(filepath, sizeof(filepath), pipe) && num < max_pairs) {
        tgn_trim_line(filepath);
        if (filepath[0] == '\0') continue;

        /* prompt = filename stem (e.g. "prog/fizz-buzz") */
        const char *base = strrchr(filepath, '/');
        base = base ? base + 1 : filepath;
        char stem[256];
        snprintf(stem, sizeof(stem), "%s", base);
        char *dot = strrchr(stem, '.');
        if (dot) *dot = '\0';

        FILE *f = fopen(filepath, "r");
        if (!f) continue;
        char code[TGN_RAW_CODE] = "";
        char line[512];
        while (fgets(line, sizeof(line), f)) {
            char *t = line;
            while (*t == ' ' || *t == '\t') t++;
            /* Skip comments and the (main func)/(funcend) wrappers so example
               programs match the bare-body style of the DSL code blocks. */
            if (t[0] == '/' && t[1] == '/') continue;
            if (strncmp(t, "(main func)", 11) == 0) continue;
            if (strncmp(t, "(funcend)", 9) == 0) continue;
            if (strlen(code) + strlen(line) + 1 < sizeof(code)) {
                if (code[0]) strncat(code, "\n", sizeof(code) - strlen(code) - 1);
                strncat(code, line, sizeof(code) - strlen(code) - 1);
            }
        }
        fclose(f);

        snprintf(pairs[num].prompt, sizeof(pairs[num].prompt), "%s", stem);
        snprintf(pairs[num].code, sizeof(pairs[num].code), "%s", code);
        num++;
    }

    pclose(pipe);
    return num;
}

int tgn_build_vocab(TgnVocab *v, const TgnRawPair *pairs, int num_pairs) {
    tgn_vocab_init(v);
    tgn_vocab_add(v, "<unk>");
    tgn_vocab_add(v, "<sep>");
    tgn_vocab_add(v, "<eos>");
    tgn_vocab_add(v, "\n");

    for (int i = 0; i < num_pairs; i++) {
        char buf[TGN_RAW_CODE + 300];
        snprintf(buf, sizeof(buf), "%s\n%s", pairs[i].prompt, pairs[i].code);
        /* add each token to vocab */
        char lower[1024];
        snprintf(lower, sizeof(lower), "%s", buf);
        tgn_tolower(lower);
        char cur[128];
        int cur_len = 0;
        for (const char *p = lower; ; p++) {
            char c = *p;
            if (c == '\n' || c == '\0' || c == ' ' || c == '\t' || c == '\r') {
                if (cur_len > 0) {
                    cur[cur_len] = '\0';
                    tgn_vocab_add(v, cur);
                    cur_len = 0;
                }
                if (c == '\n') tgn_vocab_add(v, "\n");
                if (c == '\0') break;
            } else {
                if (cur_len < 127) cur[cur_len++] = c;
            }
        }
    }
    return v->count;
}

int tgn_tokenize_pairs(const TgnVocab *v, const TgnRawPair *pairs, int num_pairs,
                       TgnPair *out, int max_out) {
    int n = 0;
    int code_tmp[TGN_MAX_SEQ];
    int prompt_tmp[TGN_MAX_PROMPT];
    for (int i = 0; i < num_pairs && n < max_out; i++) {
        TgnPair *p = &out[n];
        /* Tokenize into temp buffers first so oversized programs are skipped
           instead of silently truncated to a partial program. */
        p->prompt_len = tgn_tokenize(v, pairs[i].prompt, prompt_tmp, TGN_MAX_PROMPT);
        p->code_len = tgn_tokenize(v, pairs[i].code, code_tmp, TGN_MAX_SEQ);
        if (p->code_len == 0 || p->code_len > TGN_MAX_CODE) continue;
        if (p->prompt_len > TGN_MAX_PROMPT) p->prompt_len = TGN_MAX_PROMPT;
        memcpy(p->prompt_tokens, prompt_tmp, (size_t)p->prompt_len * sizeof(int));
        memcpy(p->code_tokens, code_tmp, (size_t)p->code_len * sizeof(int));
        n++;
    }
    return n;
}

/* ==================== Training ==================== */

#define TGN_NEG_SAMPLES 96

static float tgn_train_pair(TinyGenModel *m, const int *seq, int len, int prompt_len, float lr) {
    const int E = m->embed, H = m->hidden, V = m->vocab_size;

    /* Forward: embeddings + causal blocks (save per-block inputs for backprop) */
    Matrix *x = mat_create(len, E);
    for (int i = 0; i < len; i++) {
        int t = seq[i];
        if (t < 0 || t >= V) t = TGN_UNK;
        for (int d = 0; d < E; d++)
            x->data[i * E + d] = m->token_emb->data[t * E + d] + m->pos_emb->data[i * E + d];
    }
    Matrix *blk_in[TGN_NUM_LAYERS];
    Matrix *cur = x;
    for (int l = 0; l < TGN_NUM_LAYERS; l++) {
        blk_in[l] = cur;
        cur = block_forward_mask(m->blocks[l], cur, 1);
    }

    /* Head MLP: z = cur @ head_w + head_b ; h = relu(z) */
    float *z = malloc((size_t)len * H * sizeof(float));
    float *h = malloc((size_t)len * H * sizeof(float));
    for (int i = 0; i < len; i++) {
        for (int k = 0; k < H; k++) {
            float v = m->head_b->data[k];
            for (int d = 0; d < E; d++)
                v += cur->data[i * E + d] * m->head_w->data[d * H + k];
            z[i * H + k] = v;
            h[i * H + k] = v > 0 ? v : 0;
        }
    }

    int cands[TGN_NEG_SAMPLES + 1];
    float logits[TGN_NEG_SAMPLES + 1];
    float dzh[H];
    float loss = 0;
    int positions = 0;

    /* Accumulated gradients */
    float *g_out_w = calloc((size_t)V * H, sizeof(float));
    float *g_out_b = calloc((size_t)V, sizeof(float));
    float *g_head_w = calloc((size_t)E * H, sizeof(float));
    float *g_head_b = calloc((size_t)H, sizeof(float));
    float *g_te = calloc((size_t)V * E, sizeof(float));
    float *g_pe = calloc((size_t)TGN_MAX_SEQ * E, sizeof(float));
    /* dL/d block output per position */
    Matrix *dcur = mat_create(len, E);
    /* block grads are accumulated into m->blocks[l]->*_grad (zeroed per pair) */
    for (int l = 0; l < TGN_NUM_LAYERS; l++) {
        TransformerBlock *b = m->blocks[l];
        mat_fill(b->ln1_gamma_grad, 0.0f);
        mat_fill(b->ln1_beta_grad, 0.0f);
        mat_fill(b->ln2_gamma_grad, 0.0f);
        mat_fill(b->ln2_beta_grad, 0.0f);
        mat_fill(b->attn->wq_grad, 0.0f);
        mat_fill(b->attn->wk_grad, 0.0f);
        mat_fill(b->attn->wv_grad, 0.0f);
        mat_fill(b->attn->wo_grad, 0.0f);
        mat_fill(b->ff->w1_grad, 0.0f);
        mat_fill(b->ff->b1_grad, 0.0f);
        mat_fill(b->ff->w2_grad, 0.0f);
        mat_fill(b->ff->b2_grad, 0.0f);
    }

    /* Sampled softmax: at each position, train against the target class and
       a small set of uniformly sampled negatives instead of all V classes. */
    for (int i = prompt_len; i <= len - 2; i++) {
        int target = seq[i + 1];
        if (target < 0 || target >= V) target = TGN_UNK;

        cands[0] = target;
        for (int c = 1; c <= TGN_NEG_SAMPLES; c++) {
            int n = rand() % V;
            if (n == target) n = (target + c) % V;
            cands[c] = n;
        }

        /* logits over candidates */
        float maxv = -INFINITY;
        for (int c = 0; c <= TGN_NEG_SAMPLES; c++) {
            float v = m->out_b->data[cands[c]];
            for (int k = 0; k < H; k++)
                v += h[i * H + k] * m->out_w->data[cands[c] * H + k];
            logits[c] = v;
            if (v > maxv) maxv = v;
        }
        float sum = 0;
        for (int c = 0; c <= TGN_NEG_SAMPLES; c++) { logits[c] = expf(logits[c] - maxv); sum += logits[c]; }
        for (int c = 0; c <= TGN_NEG_SAMPLES; c++) logits[c] /= sum;
        loss -= logf(logits[0] + 1e-9f);
        positions++;

        /* Readout gradients */
        for (int c = 0; c <= TGN_NEG_SAMPLES; c++) {
            float g = logits[c] - (c == 0 ? 1.0f : 0.0f);
            g_out_b[cands[c]] += g;
            for (int k = 0; k < H; k++)
                g_out_w[cands[c] * H + k] += g * h[i * H + k];
        }

        /* dzh = (dL/dh) through relu */
        for (int k = 0; k < H; k++) {
            float acc = 0;
            for (int c = 0; c <= TGN_NEG_SAMPLES; c++) {
                float g = logits[c] - (c == 0 ? 1.0f : 0.0f);
                acc += g * m->out_w->data[cands[c] * H + k];
            }
            dzh[k] = acc * (z[i * H + k] > 0 ? 1.0f : 0.0f);
        }

        /* Head MLP grads: g_head_w += dzh x cur ; dcur[i] = dzh @ head_w^T */
        for (int k = 0; k < H; k++) {
            g_head_b[k] += dzh[k];
            for (int d = 0; d < E; d++) {
                g_head_w[d * H + k] += dzh[k] * cur->data[i * E + d];
                dcur->data[i * E + d] += dzh[k] * m->head_w->data[d * H + k];
            }
        }
    }

    if (positions > 0) {
        /* Backward through the transformer blocks (accumulates block grads) */
        Matrix *dblk = dcur;
        for (int l = TGN_NUM_LAYERS - 1; l >= 0; l--) {
            Matrix *dx = block_forward_mask_backward(m->blocks[l], blk_in[l], 1, dblk);
            if (l < TGN_NUM_LAYERS - 1) mat_free(dblk);
            dblk = dx;
        }
        /* dblk is dL/d(embedding input) */

        /* Embedding gradients */
        for (int i = 0; i < len; i++) {
            int t = seq[i];
            if (t < 0 || t >= V) t = TGN_UNK;
            for (int d = 0; d < E; d++) {
                g_te[t * E + d] += dblk->data[i * E + d];
                if (i < TGN_MAX_SEQ) g_pe[i * E + d] += dblk->data[i * E + d];
            }
        }
        mat_free(dblk);
        mat_free(dcur);

        /* Normalize by position count and clip element-wise */
        for (int q = 0; q < V * H; q++) g_out_w[q] = (g_out_w[q] / positions > 1.0f) ? 1.0f : (g_out_w[q] / positions < -1.0f) ? -1.0f : g_out_w[q] / positions;
        for (int q = 0; q < V; q++) g_out_b[q] = (g_out_b[q] / positions > 1.0f) ? 1.0f : (g_out_b[q] / positions < -1.0f) ? -1.0f : g_out_b[q] / positions;
        for (int q = 0; q < E * H; q++) g_head_w[q] = (g_head_w[q] / positions > 1.0f) ? 1.0f : (g_head_w[q] / positions < -1.0f) ? -1.0f : g_head_w[q] / positions;
        for (int q = 0; q < H; q++) g_head_b[q] = (g_head_b[q] / positions > 1.0f) ? 1.0f : (g_head_b[q] / positions < -1.0f) ? -1.0f : g_head_b[q] / positions;
        for (int q = 0; q < V * E; q++) g_te[q] = (g_te[q] / positions > 1.0f) ? 1.0f : (g_te[q] / positions < -1.0f) ? -1.0f : g_te[q] / positions;
        for (int q = 0; q < TGN_MAX_SEQ * E; q++) g_pe[q] = (g_pe[q] / positions > 1.0f) ? 1.0f : (g_pe[q] / positions < -1.0f) ? -1.0f : g_pe[q] / positions;

        /* Apply readout + head */
        for (int q = 0; q < V * H; q++) m->out_w->data[q] -= lr * g_out_w[q];
        for (int q = 0; q < V; q++) m->out_b->data[q] -= lr * g_out_b[q];
        for (int q = 0; q < E * H; q++) m->head_w->data[q] -= lr * g_head_w[q];
        for (int q = 0; q < H; q++) m->head_b->data[q] -= lr * g_head_b[q];
        for (int q = 0; q < V * E; q++) m->token_emb->data[q] -= lr * g_te[q];
        for (int q = 0; q < TGN_MAX_SEQ * E; q++) m->pos_emb->data[q] -= lr * g_pe[q];

        /* Apply block grads (clip element-wise) */
        for (int l = 0; l < TGN_NUM_LAYERS; l++) {
            TransformerBlock *b = m->blocks[l];
            int n = TT_EMBED_DIM * TT_EMBED_DIM;
            for (int q = 0; q < n; q++) {
                float v = b->attn->wq_grad->data[q] / positions;
                b->attn->wq->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->attn->wk_grad->data[q] / positions;
                b->attn->wk->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->attn->wv_grad->data[q] / positions;
                b->attn->wv->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->attn->wo_grad->data[q] / positions;
                b->attn->wo->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
            }
            n = TT_EMBED_DIM * TT_FF_DIM;
            for (int q = 0; q < n; q++) {
                float v = b->ff->w1_grad->data[q] / positions;
                b->ff->w1->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->ff->w2_grad->data[q] / positions;
                b->ff->w2->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
            }
            for (int q = 0; q < TT_FF_DIM; q++) {
                float v = b->ff->b1_grad->data[q] / positions;
                b->ff->b1->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
            }
            for (int q = 0; q < TT_EMBED_DIM; q++) {
                float v = b->ff->b2_grad->data[q] / positions;
                b->ff->b2->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->ln1_gamma_grad->data[q] / positions;
                b->ln1_gamma->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->ln1_beta_grad->data[q] / positions;
                b->ln1_beta->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->ln2_gamma_grad->data[q] / positions;
                b->ln2_gamma->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
                v = b->ln2_beta_grad->data[q] / positions;
                b->ln2_beta->data[q] -= lr * (v > 1.0f ? 1.0f : v < -1.0f ? -1.0f : v);
            }
        }
    } else {
        mat_free(dcur);
    }

    /* free block inputs (x and subsequent block outputs) */
    for (int l = 0; l < TGN_NUM_LAYERS; l++) {
        mat_free(blk_in[l]);
    }
    /* blk_in[0] == x; block outputs between layers were moved: blk_in[l+1] == block l output.
       The last block output is `cur` — free it too. */
    mat_free(cur);

    free(g_out_w); free(g_out_b); free(g_head_w); free(g_head_b);
    free(g_te); free(g_pe);
    free(z);
    free(h);
    return loss;
}

float tgn_train(TinyGenModel *m, const TgnPair *pairs, int num_pairs,
                int epochs, float lr) {
    int *seq = malloc((size_t)TGN_MAX_SEQ * sizeof(int));
    float last_loss = 0;

    for (int ep = 0; ep < epochs; ep++) {
        float total = 0;
        int count = 0;
        for (int p = 0; p < num_pairs; p++) {
            const TgnPair *pair = &pairs[p];
            int L = 0;
            for (int i = 0; i < pair->prompt_len && L < TGN_MAX_SEQ; i++)
                seq[L++] = pair->prompt_tokens[i];
            if (L >= TGN_MAX_SEQ) continue;
            seq[L++] = TGN_SEP;
            int prompt_len = pair->prompt_len;   /* index of sep */
            for (int i = 0; i < pair->code_len && L < TGN_MAX_SEQ; i++)
                seq[L++] = pair->code_tokens[i];
            if (L >= TGN_MAX_SEQ) continue;
            if (L + 1 >= TGN_MAX_SEQ) continue;
            seq[L++] = TGN_EOS;

            total += tgn_train_pair(m, seq, L, prompt_len, lr);
            count++;
        }
        last_loss = count > 0 ? total / count : 0;
        if ((ep + 1) % 10 == 0 || ep == 0 || ep == epochs - 1) {
            printf("  epoch %4d/%d  loss=%.4f  (pairs=%d)\n", ep + 1, epochs, last_loss, count);
            fflush(stdout);
        }
    }

    free(seq);
    return last_loss;
}

/* ==================== Inference ==================== */

static float tgn_sample_logits(const float *logits, int V, float temperature,
                               int top_k, const int *recent, int recent_n,
                               float rep_penalty) {
    float *p = malloc((size_t)V * sizeof(float));
    memcpy(p, logits, (size_t)V * sizeof(float));

    /* Repetition penalty: lower the score of any token seen in the recent
       window so a stuck model is pushed out of degenerate loops. */
    if (rep_penalty > 1.0f) {
        int *seen = calloc((size_t)V, sizeof(int));
        for (int i = 0; i < recent_n; i++) {
            int t = recent[i];
            if (t >= 0 && t < V) seen[t] = 1;
        }
        for (int j = 0; j < V; j++)
            if (seen[j]) p[j] /= rep_penalty;
        free(seen);
    }

    if (temperature > 0)
        for (int j = 0; j < V; j++) p[j] /= temperature;

    if (top_k > 0 && top_k < V) {
        int *idx = malloc((size_t)V * sizeof(int));
        for (int j = 0; j < V; j++) idx[j] = j;
        /* partial selection sort of top_k */
        for (int a = 0; a < top_k; a++) {
            int best = a;
            for (int b = a + 1; b < V; b++)
                if (p[idx[b]] > p[idx[best]]) best = b;
            int tmp = idx[a]; idx[a] = idx[best]; idx[best] = tmp;
        }
        float thr = p[idx[top_k - 1]];
        for (int j = 0; j < V; j++)
            if (p[j] < thr) p[j] = -INFINITY;
        free(idx);
    }

    float maxv = -INFINITY;
    for (int j = 0; j < V; j++) if (p[j] > maxv) maxv = p[j];
    float sum = 0;
    for (int j = 0; j < V; j++) { p[j] = expf(p[j] - maxv); sum += p[j]; }
    float r = sum * ((float)rand() / (float)RAND_MAX);
    float cum = 0;
    int tok = 0;
    for (int j = 0; j < V; j++) {
        cum += p[j];
        if (r < cum || j == V - 1) { tok = j; break; }
    }
    free(p);
    return tok;
}

int tgn_generate(TinyGenModel *m, const TgnVocab *v, const char *prompt,
                 float temperature, int top_k, float rep_penalty, char *out,
                 int out_size, int stream) {
    int pt[TGN_MAX_PROMPT];
    int pn = tgn_tokenize(v, prompt, pt, TGN_MAX_PROMPT);
    if (pn > TGN_MAX_PROMPT) pn = TGN_MAX_PROMPT;

    int seq[TGN_MAX_SEQ];
    int L = 0;
    for (int i = 0; i < pn && L < TGN_MAX_SEQ; i++)
        seq[L++] = pt[i];
    if (L < TGN_MAX_SEQ) seq[L++] = TGN_SEP;
    int code_start = L;

    const int E = m->embed, H = m->hidden, V = m->vocab_size;
    float *logits = malloc((size_t)V * sizeof(float));

    /* Track consecutive blank lines so a stuck model stops instead of padding
       the output with newline tokens. */
    int consec_nl = 0;
    int nl_id = (v && tgn_vocab_find(v, "\n") >= 0) ? tgn_vocab_find(v, "\n") : -1;

    for (int gen = 0; gen < TGN_MAX_CODE; gen++) {
        if (L >= TGN_MAX_SEQ) break;

        /* Forward over current sequence (causal), keep last hidden row */
        Matrix *x = mat_create(L, E);
        for (int i = 0; i < L; i++) {
            int t = seq[i];
            if (t < 0 || t >= V) t = TGN_UNK;
            for (int d = 0; d < E; d++)
                x->data[i * E + d] = m->token_emb->data[t * E + d] + m->pos_emb->data[i * E + d];
        }
        Matrix *cur = x;
        for (int l = 0; l < TGN_NUM_LAYERS; l++) {
            Matrix *nxt = block_forward_mask(m->blocks[l], cur, 1);
            if (l > 0) mat_free(cur);
            cur = nxt;
        }

        int last = L - 1;
        for (int j = 0; j < V; j++) {
            float vv = m->out_b->data[j];
            for (int k = 0; k < H; k++) {
                float z = 0;
                for (int d = 0; d < E; d++)
                    z += cur->data[last * E + d] * m->head_w->data[d * H + k];
                vv += (z > 0 ? z : 0) * m->out_w->data[j * H + k];
            }
            logits[j] = vv;
        }
        mat_free(cur);

        /* Repetition window: last few generated code tokens. */
        int wstart = code_start > L - 24 ? code_start : L - 24;
        int wlen = L - wstart;
        int tok = (int)tgn_sample_logits(logits, V, temperature, top_k,
                                         seq + wstart, wlen, rep_penalty);
        seq[L++] = tok;

        if (nl_id >= 0 && tok == nl_id) {
            if (++consec_nl >= 3) { L--; break; }   /* three blank lines: stop */
        } else {
            consec_nl = 0;
        }

        if (stream && v && tok >= 0 && tok < v->count) {
            if (strcmp(v->words[tok], "\n") == 0) printf("\n");
            else printf("%s ", v->words[tok]);
            fflush(stdout);
        }

        if (tok == TGN_EOS) break;
    }

    /* Reconstruct code text from generated tokens */
    out[0] = '\0';
    int pos = 0;
    for (int i = code_start; i < L; i++) {
        int t = seq[i];
        if (t == TGN_UNK || t == TGN_SEP || t < 0 || t >= v->count) continue;
        if (t == TGN_EOS) break;
        const char *w = v->words[t];
        if (strcmp(w, "\n") == 0) {
            if (pos > 0 && out[pos - 1] == ' ') out[--pos] = '\0';
            if (pos < out_size - 1) out[pos++] = '\n';
        } else {
            int wlen = (int)strlen(w);
            if (pos + wlen + 1 < out_size) {
                memcpy(out + pos, w, (size_t)wlen);
                pos += wlen;
                out[pos++] = ' ';
            }
        }
    }
    if (pos > 0 && out[pos - 1] == ' ') out[pos - 1] = '\0';
    else out[pos] = '\0';

    free(logits);
    return 0;
}

int tgn_load(const char *model_path, TinyGenModel **m, TgnVocab *v) {
    if (tgn_vocab_load(v, model_path) != 0) return -1;
    TinyGenModel *model = tgn_model_create(v);
    if (!model) return -1;
    if (tgn_model_load(model, model_path) != 0) {
        tgn_model_free(model);
        return -1;
    }
    *m = model;
    return 0;
}