train_tiny.c

/*
 * This file train_tiny.c is part of L1vm.
 *
 * (c) Copyright Stefan Pietzonke (info@midnight-coding.de), 2026
 *
 * L1vm is free software: you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * L1vm is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with L1vm.  If not, see <http://www.gnu.org/licenses/>.
 */

/* train_tiny.c — Standalone training tool for the tiny transformer. */
/* Reads DSL files, trains a classification model, saves weights. */
/* Also trains the tiny generative LLM (--gen-train / --gen-predict). */

#include "tiny_transformer.h"
#include "tiny_gen.h"
#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <math.h>
#include <time.h>

static void usage(const char *prog) {
    printf("Usage: %s [options]\n", prog);
    printf("Options:\n");
    printf("  --dsl-dir <path>     DSL directory (default: dsl/)\n");
    printf("  --model <path>       Output model file (default: tiny_model.tiny)\n");
    printf("  --epochs <n>         Training epochs (default: 50)\n");
    printf("  --lr <float>         Learning rate (default: 0.01)\n");
    printf("  --predict <prompt>   Predict emitter for a prompt (requires trained model)\n");
    printf("  --gen-train          Train the tiny generative LLM\n");
    printf("  --gen-predict <prompt>  Generate L1VM code for a prompt (requires trained gen model)\n");
    printf("  --gen-dsl-dir <path> DSL directory for gen training (default: dsl/)\n");
    printf("  --gen-examples-dir <path> Example programs dir (default: l1vm-example-code/prog/)\n");
    printf("  --gen-model <path>   Gen model file (default: tiny_gen.tiny)\n");
    printf("  --gen-epochs <n>     Gen training epochs (default: 400)\n");
    printf("  --gen-lr <float>     Gen learning rate (default: 0.2)\n");
    printf("  --gen-temp <float>   Sampling temperature (default: 0.7)\n");
    printf("  --gen-topk <n>       Top-k sampling (default: 16)\n");
    printf("  --gen-rep <float>    Repetition penalty >1 (default: 1.2)\n");
    printf("  --gen-stream         Stream generated tokens as they are sampled\n");
    printf("  --help               Show this help\n");
}

static const char *emitter_names[] = {
    "math","input_loop","loop","for_sum","print_even","find_max",
    "countdown","fib_seq","input_sort","median","string_cat",
    "string_compare","array_assign","array_reverse","array_find",
    "input_fact","array_vmath","read_file","write_file","string_to_num",
    "timer","factorial","fizzbuzz","primes","even_odd","power",
    "mult_table","guess","gcd","hello_name","random","array_min_max",
    "bool_demo","bit_check","fann_create","fann_train","fann_run",
    "average","selection_sort","palindrome","lcm","collatz",
    "sum_of_digits","reverse_string","armstrong","perfect_number",
    "count_vowels","anagram_check","string_to_upper","string_to_lower",
    "caesar_cipher","palindrome_string","bubble_sort","binary_search",
    "square_root","prime_factorization","standard_deviation",
    "compound_interest","decimal_to_binary","dice_roll","double_math",
    "double_circle_area","double_average","double_compound_interest",
    "double_pythagoras","double_temp_convert","double_sqrt","function",
    "string_length","stack","queue","insertion_sort","calculator",
    "unit_converter","rock_paper_scissors","pyramid","temp_converter_menu",
    "sort_stats","string_analyzer","number_analyzer","filter_numbers",
    "random_generator","math_menu","quiz_game","bmi_calculator",
    "statistics_suite","linked_list","binary_search_tree","tree_traversal",
    "graph_bfs_dfs","n_queens","sudoku","levenshtein","maze_generator",
    "maze_solver","monte_carlo","matrix_mul","matrix_transpose",
    "numerical_integration","complex_numbers","linear_regression",
    "base_converter","freq_analysis","shuffle","weighted_random",
    "ascii_table","bignum_math","password_card","chess_problem",
    "shell_repl","webserver","sdl_window","sdl_button","thread",
    "scheduler","shell_exec","json","crypto","bluetooth_ble",
    "serial_rs232","gpio","gps","timer_date","sdl_sound","sdl_joystick",
    "sdl_mouse","fractal","cluster_3x1","reload","coordinate_grid",
    "turmite","crossword","linter","double_power","double_volume_sphere",
    "double_discount","double_simple_interest","double_bmi",
    "double_standard_deviation","double_kinetic_energy","hello_world",
    "string_find","string_split","switch_demo","type_convert",
    "iterative_factorial","random_walk","bar_chart","hanoi_tower",
    "ascii_art","number_to_words","temperature_table","loop_demo",
    "pointer","struct","hex_binary","shell_args","time"
};
#define NUM_EMITTER_NAMES (sizeof(emitter_names) / sizeof(emitter_names[0]))

int main(int argc, char **argv) {
    const char *dsl_dir = "dsl";
    const char *model_path = "tiny_model.tiny";
    int epochs = 50;
    float lr = 0.01f;
    const char *predict_prompt = NULL;

    int gen_train = 0;
    const char *gen_predict_prompt = NULL;
    const char *gen_dsl_dir = "dsl";
    const char *gen_examples_dir = "l1vm-example-code/prog";
    const char *gen_model = "tiny_gen.tiny";
    int gen_epochs = 400;
    float gen_lr = 0.2f;
    float gen_temp = 0.7f;
    int gen_topk = 16;
    float gen_rep = 1.2f;
    int gen_stream = 0;

    for (int i = 1; i < argc; i++) {
        if (strcmp(argv[i], "--help") == 0) {
            usage(argv[0]);
            return 0;
        } else if (strcmp(argv[i], "--dsl-dir") == 0 && i + 1 < argc) {
            dsl_dir = argv[++i];
        } else if (strcmp(argv[i], "--model") == 0 && i + 1 < argc) {
            model_path = argv[++i];
        } else if (strcmp(argv[i], "--epochs") == 0 && i + 1 < argc) {
            epochs = atoi(argv[++i]);
        } else if (strcmp(argv[i], "--lr") == 0 && i + 1 < argc) {
            lr = (float)atof(argv[++i]);
        } else if (strcmp(argv[i], "--predict") == 0 && i + 1 < argc) {
            predict_prompt = argv[++i];
        } else if (strcmp(argv[i], "--gen-train") == 0) {
            gen_train = 1;
        } else if (strcmp(argv[i], "--gen-predict") == 0 && i + 1 < argc) {
            gen_predict_prompt = argv[++i];
        } else if (strcmp(argv[i], "--gen-dsl-dir") == 0 && i + 1 < argc) {
            gen_dsl_dir = argv[++i];
        } else if (strcmp(argv[i], "--gen-examples-dir") == 0 && i + 1 < argc) {
            gen_examples_dir = argv[++i];
        } else if (strcmp(argv[i], "--gen-model") == 0 && i + 1 < argc) {
            gen_model = argv[++i];
        } else if (strcmp(argv[i], "--gen-epochs") == 0 && i + 1 < argc) {
            gen_epochs = atoi(argv[++i]);
        } else if (strcmp(argv[i], "--gen-lr") == 0 && i + 1 < argc) {
            gen_lr = (float)atof(argv[++i]);
        } else if (strcmp(argv[i], "--gen-temp") == 0 && i + 1 < argc) {
            gen_temp = (float)atof(argv[++i]);
        } else if (strcmp(argv[i], "--gen-topk") == 0 && i + 1 < argc) {
            gen_topk = atoi(argv[++i]);
        } else if (strcmp(argv[i], "--gen-rep") == 0 && i + 1 < argc) {
            gen_rep = (float)atof(argv[++i]);
        } else if (strcmp(argv[i], "--gen-stream") == 0) {
            gen_stream = 1;
        } else {
            printf("Unknown option: %s\n", argv[i]);
            usage(argv[0]);
            return 1;
        }
    }

    /* ===== Tiny generative LLM modes ===== */
    if (gen_predict_prompt) {
        TinyGenModel *m = NULL;
        TgnVocab v;
        if (tgn_load(gen_model, &m, &v) != 0) {
            printf("Error: cannot load gen model '%s' (train with --gen-train first)\n", gen_model);
            return 1;
        }
        printf("Loaded gen model: %s (vocab %d words)\n", gen_model, v.count);
        char out[TGN_RAW_CODE + 4096];
        printf("\n=== Tiny LLM generating: \"%s\" ===\n", gen_predict_prompt);
        if (gen_stream) printf("\n");
        if (tgn_generate(m, &v, gen_predict_prompt, gen_temp, gen_topk,
                         gen_rep, out, sizeof(out), gen_stream) != 0) {
            printf("Generation failed.\n");
            tgn_model_free(m);
            return 1;
        }
        printf("\n\n%s\n", out);
        tgn_model_free(m);
        return 0;
    }

    if (gen_train) {
        srand(42);
        printf("=== Tiny Generative LLM Trainer ===\n");
        TgnRawPair *pairs = (TgnRawPair *)calloc(TGN_MAX_PAIRS, sizeof(TgnRawPair));
        if (!pairs) { printf("Out of memory.\n"); return 1; }
        int num = 0;
        printf("Collecting DSL pairs from %s ...\n", gen_dsl_dir);
        num += tgn_collect_dsl(pairs + num, TGN_MAX_PAIRS - num, gen_dsl_dir);
        printf("Collecting example programs from %s ...\n", gen_examples_dir);
        num += tgn_collect_examples(pairs + num, TGN_MAX_PAIRS - num, gen_examples_dir);
        printf("Collected %d (prompt, code) pairs\n", num);
        if (num == 0) {
            printf("No training data found.\n");
            free(pairs);
            return 1;
        }

        TgnVocab v;
        int vsize = tgn_build_vocab(&v, pairs, num);
        printf("Vocabulary: %d tokens\n", vsize);

        TgnPair *tokenized = (TgnPair *)calloc(TGN_MAX_PAIRS, sizeof(TgnPair));
        if (!tokenized) { printf("Out of memory.\n"); free(pairs); return 1; }
        int np = tgn_tokenize_pairs(&v, pairs, num, tokenized, TGN_MAX_PAIRS);
        free(pairs);
        printf("Tokenized %d pairs\n", np);

        TinyGenModel *m = tgn_model_create(&v);
        if (!m) { printf("Out of memory.\n"); free(tokenized); return 1; }

        printf("\nTraining for %d epochs, lr=%.3f\n", gen_epochs, gen_lr);
        float final = tgn_train(m, tokenized, np, gen_epochs, gen_lr);
        free(tokenized);

        printf("\nFinal loss: %.4f\n", final);
        if (tgn_model_save(m, gen_model) != 0) {
            printf("Error: cannot save model.\n");
        } else {
            if (tgn_vocab_save(&v, gen_model) != 0)
                printf("Error: cannot save vocab.\n");
            else
                printf("Model saved to %s (+ .vocab)\n", gen_model);
        }
        tgn_model_free(m);
        return 0;
    }

    /* Predict mode */
    if (predict_prompt) {
        printf("Loading model from %s...\n", model_path);
        if (tt_init(model_path) != 0) {
            /* tt_init returns 1 if no model found, but still initializes */
        }
        float score = 0;
        int emitter = tt_predict(predict_prompt, &score);
        if (emitter >= 0 && emitter < (int)NUM_EMITTER_NAMES) {
            printf("Prompt: \"%s\"\n", predict_prompt);
            printf("Predicted emitter: %s (confidence: %.2f%%)\n",
                   emitter_names[emitter], score * 100.0f);
        } else {
            printf("No prediction available for: \"%s\"\n", predict_prompt);
        }
        return 0;
    }

    /* Training mode */
    printf("=== Tiny Transformer Trainer ===\n");
    printf("DSL directory: %s\n", dsl_dir);
    printf("Model output:  %s\n", model_path);
    printf("Epochs:        %d\n", epochs);
    printf("Learning rate: %.4f\n", lr);
    printf("\n");

    int result = tt_train(dsl_dir, model_path, epochs, lr);

    if (result == 0) {
        printf("\nTraining complete! Model saved to %s\n", model_path);
        printf("Use --predict \"your prompt\" to test.\n");
    } else {
        printf("\nTraining failed!\n");
    }

    return result;
}