Function main

examples/slm_basic_train_ex.cpp:95–317

// ----------------------------------------------------------------------------------------

int main(int argc, char** argv)
{
    try
    {
        command_line_parser parser;
        parser.add_option("train", "Train a small transformer on the built-in Shakespeare text");
        parser.add_option("generate", "Generate text from a previously trained model (needs shakespeare_prompt)");
        parser.add_option("learning-rate", "Set the learning rate for training (default: 1e-4)", 1);
        parser.add_option("batch-size", "Set the mini-batch size for training (default: 64)", 1);
        parser.add_option("generation-length", "Set the length of generated text (default: 400)", 1);
        parser.add_option("alpha", "Set the weight decay for Adam optimizer (default: 0.004)", 1);
        parser.add_option("beta1", "Set the first moment coefficient (default: 0.9)", 1);
        parser.add_option("beta2", "Set the second moment coefficient (default: 0.999)", 1);
        parser.add_option("max-samples", "Set the maximum number of training samples (default: 50000)", 1);
        parser.add_option("shuffle", "Shuffle training sequences and labels before training (default: false)");
        parser.parse(argc, argv);

        if (parser.number_of_arguments() == 0 && !parser.option("train") && !parser.option("generate"))
        {
            parser.print_options();
            return 0;
        }

        // Default values
        const double learning_rate = get_option(parser, "learning-rate", 1e-4);
        const long batch_size = get_option(parser, "batch-size", 64);
        const int generation_length = get_option(parser, "generation-length", 400);
        const double alpha = get_option(parser, "alpha", 0.004);       // Initial learning rate for Adam
        const double beta1 = get_option(parser, "beta1", 0.9);         // Decay rate for the first moment estimate
        const double beta2 = get_option(parser, "beta2", 0.999);       // Decay rate for the second moment estimate
        const size_t max_samples = get_option(parser, "max-samples",50000); // Default maximum number of training samples

        // We define a minimal config for demonstration
        const long vocab_size = MAX_TOKEN_ID + 1 + 1;   // 256 for chars + 1 pad token
        const long num_layers = 3;
        const long num_heads = 4;
        const long embedding_dim = 64;
        const long max_seq_len = 80;   // a small sequence length for the example
        const bool use_squeezing = false;

        using my_transformer_cfg = transformer::transformer_config<
            vocab_size,
            num_layers,
            num_heads,
            embedding_dim,
            max_seq_len,
            use_squeezing,
            gelu,
            dropout_10
        >;

        // For GPU usage (if any), set gpus = {0} for a single GPU, etc.
        std::vector<int> gpus{ 0 };

        // The model file to store or load
        const std::string model_file = "shakespeare_lm_char_model.dat";

        // ----------------------------------------------------------------------------------------