brackets-code/tiny_gen.h at main · koder77/brackets-code · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
/*
 * tiny_gen.h - Tiny generative LLM for L1VM code generation.
 *
 * A small autoregressive next-token model: given a natural-language
 * prompt it generates L1VM source tokens one at a time (like a tiny
 * language model), conditioned on the prompt and its own output so far.
 *
 * (c) Copyright Stefan Pietzonke (info@midnight-coding.de), 2026
 *
 * This file is part of brackets-code.
 */

#ifndef TINY_GEN_H
#define TINY_GEN_H

#include "tiny_transformer.h"

/* ==================== Configuration ==================== */

#define TGN_MAX_VOCAB   2048
#define TGN_MAX_SEQ     512
#define TGN_MAX_CODE    448
#define TGN_MAX_PROMPT  16
#define TGN_HIDDEN      64
#define TGN_NUM_LAYERS  TT_NUM_LAYERS
#define TGN_EMBED       TT_EMBED_DIM

#define TGN_MAX_PAIRS  4096
#define TGN_RAW_CODE   8192

/* Special token ids */
#define TGN_UNK  0
#define TGN_SEP  1
#define TGN_EOS  2

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

typedef struct {
    char  words[TGN_MAX_VOCAB][64];
    int   count;
} TgnVocab;

void tgn_vocab_init(TgnVocab *v);
int  tgn_vocab_add(TgnVocab *v, const char *word);
int  tgn_vocab_find(const TgnVocab *v, const char *word);
int  tgn_tokenize(const TgnVocab *v, const char *text, int *tokens, int max_tokens);
int  tgn_vocab_save(const TgnVocab *v, const char *path);
int  tgn_vocab_load(TgnVocab *v, const char *path);

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

typedef struct {
    Matrix *token_emb;          /* [vocab, embed] */
    Matrix *pos_emb;            /* [max_seq, embed] */
    Matrix *head_w;             /* [embed, hidden] */
    Matrix *head_b;             /* [hidden] */
    Matrix *out_w;              /* [vocab, hidden] (row-major, cache friendly) */
    Matrix *out_b;              /* [vocab] */
    TransformerBlock *blocks[TGN_NUM_LAYERS];   /* causal, frozen */
    int vocab_size;
    int embed;
    int max_seq;
    int hidden;
} TinyGenModel;

TinyGenModel *tgn_model_create(const TgnVocab *v);
void          tgn_model_free(TinyGenModel *m);
int           tgn_model_save(const TinyGenModel *m, const char *path);
int           tgn_model_load(TinyGenModel *m, const char *path);

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

typedef struct {
    char prompt[256];
    char code[TGN_RAW_CODE];
} TgnRawPair;

typedef struct {
    int prompt_tokens[TGN_MAX_PROMPT];
    int prompt_len;
    int code_tokens[TGN_MAX_CODE];
    int code_len;
} TgnPair;

/* Collect (prompt, code) pairs from a DSL directory and an examples
   directory. Pairs stay raw until the vocab is built. */
int tgn_collect_dsl(TgnRawPair *pairs, int max_pairs, const char *dsl_dir);
int tgn_collect_examples(TgnRawPair *pairs, int max_pairs, const char *examples_dir);
int tgn_build_vocab(TgnVocab *v, const TgnRawPair *pairs, int num_pairs);
int tgn_tokenize_pairs(const TgnVocab *v, const TgnRawPair *pairs, int num_pairs,
                       TgnPair *out, int max_out);

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

/* Teacher-forced next-token training over the code portion of each pair.
   Returns the final average loss. */
float tgn_train(TinyGenModel *m, const TgnPair *pairs, int num_pairs,
                int epochs, float lr);

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

/* Autoregressively generate code for a prompt. If stream != 0, each token
   is printed as it is sampled. Returns 0 on success. rep_penalty (>1)
   suppresses tokens seen in the recent generation window. */
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);

/* Load the whole generator (model + vocab) in one call. Returns 0 on success. */
int tgn_load(const char *model_path, TinyGenModel **m, TgnVocab *v);

#endif /* TINY_GEN_H */