/*
This program demonstrates a minimal example of a Very Small Language Model (VSLM)
using dlib's deep learning tools. It includes two modes:
1) --train : Train a small Transformer-based language model on a character-based
corpus extracted from "slm_data.h" (named shakespeare_text).
2) --generate: Generate new text from a trained model, given an initial prompt
extracted from "slm_data.h" (named shakespeare_prompt).
The "slm_dels.h" header is expected to provide a comprehensive Transformer
definition with the following key elements:
- A configurable transformer_config
- The use of classification_head to output a single token
- The network_type<true> or network_type<false> for training vs inference
- The typical dlib constructs (input<matrix<int>>, etc.)
Character-level tokenization is used here. Each character is directly transformed
into an integer token. The model attempts to learn the sequence of characters in
shakespeare_text. Then you can ask the model to generate new text from a short
prompt.
This model is intentionally kept small (few neurons/parameters) to ensure
simplicity and efficiency. As a result, it may not generalize well to unseen
patterns or concepts. However, it effectively illustrates the principle of
attention and the ability to perfectly memorize and reproduce sequences from
the training data. This makes it a useful educational tool for understanding
the mechanics of Transformer models, even if it lacks the capacity for
sophisticated language understanding.
*/
#include <iostream>
#include <string>
#include <vector>
#include <algorithm>
#include <cmath>
#include <random>
#include <dlib/data_io.h>
#include <dlib/cmd_line_parser.h>
#include <dlib/misc_api.h>
// Include Transformer definitions
#include "slm_defs.h"
// This header "slm_data.h" is assumed to contain:
// const std::string shakespeare_text;
// const std::string shakespeare_prompt;
#include "slm_data.h"
// ----------------------------------------------------------------------------------------
using namespace std;
using namespace dlib;
// We treat each character as a token ID in [0..255].
const int MAX_TOKEN_ID = 255;
const int PAD_TOKEN = 256; // an extra "pad" token if needed
// For simplicity, we assume each line from shakespeare_text is appended, ignoring them.
std::vector<int> char_based_tokenize(const std::string& text)
{
std::vector<int> tokens;
tokens.reserve(text.size());
for (const int c : text)
{
tokens.push_back(std::min(c, MAX_TOKEN_ID));
}
return tokens;
}
// Function to shuffle samples and labels in sync
void shuffle_samples_and_labels(std::vector<matrix<int, 0, 1>>& samples, std::vector<unsigned long>& labels) {
std::vector<size_t> indices(samples.size());
std::iota(indices.begin(), indices.end(), 0); // Fill with 0, 1, 2, ..., N-1
std::shuffle(indices.begin(), indices.end(), std::default_random_engine{});
// Create temporary vectors to hold shuffled data
std::vector<matrix<int, 0, 1>> shuffled_samples(samples.size());
std::vector<unsigned long> shuffled_labels(labels.size());
// Apply the shuffle
for (size_t i = 0; i < indices.size(); ++i)
{
shuffled_samples[i] = samples[indices[i]];
shuffled_labels[i] = labels[indices[i]];
}
// Replace the original data with shuffled data
samples = std::move(shuffled_samples);
labels = std::move(shuffled_labels);
}
// ----------------------------------------------------------------------------------------
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";
// ----------------------------------------------------------------------------------------
// Train mode
// ----------------------------------------------------------------------------------------
if (parser.option("train"))
{
cout << "=== TRAIN MODE ===\n";
// 1) Prepare training data (simple approach)
// We will store characters from shakespeare_text into a vector
// and then produce training samples of length (max_seq_len+1),
// where the last token is the label to predict from the preceding max_seq_len.
auto full_tokens = char_based_tokenize(shakespeare_text);
if (full_tokens.empty())
{
cerr << "ERROR: The Shakespeare text is empty. Please provide a valid training text.\n";
return 0;
}
// Calculate the maximum number of sequences
size_t max_sequences = (full_tokens.size() > (size_t)max_seq_len + 1)
? (full_tokens.size() - ((size_t)max_seq_len + 1))
: 0;
// Display the size of the training text and the number of sequences
cout << "Training text size: " << full_tokens.size() << " characters\n";
cout << "Maximum number of sequences: " << max_sequences << "\n";
// Check if the text is too short
if (max_sequences == 0)
{
cerr << "ERROR: The Shakespeare text is too short for training. It must contain at least "
<< (max_seq_len + 1) << " characters.\n";
return 0;
}
std::vector<matrix<int, 0, 1>> samples;
std::vector<unsigned long> labels;
// Let's create a training set of about (N) samples from the text
// Each sample: [x0, x1, ..., x_(max_seq_len-1)] -> y
// We'll store them in "samples" and "labels".
const size_t N = (max_sequences < max_samples) ? max_sequences : max_samples;
for (size_t start = 0; start < N; ++start)
{
matrix<int, 0, 1> seq(max_seq_len, 1);
for (long t = 0; t < max_seq_len; ++t)
seq(t, 0) = full_tokens[start + t];
samples.push_back(seq);
labels.push_back(full_tokens[start + max_seq_len]);
}
// Shuffle samples and labels if the --shuffle option is enabled
if (parser.option("shuffle"))
{
cout << "Shuffling training sequences and labels...\n";
shuffle_samples_and_labels(samples, labels);
}
// 3) Construct the network in training mode
using net_type = my_transformer_cfg::network_type<true>;
net_type net;
if (file_exists(model_file))
deserialize(model_file) >> net;
// 4) Create dnn_trainer
dnn_trainer<net_type, adam> trainer(net, adam(alpha, beta1, beta2), gpus);
trainer.set_learning_rate(learning_rate);
trainer.set_min_learning_rate(1e-6);
trainer.set_mini_batch_size(batch_size);
trainer.set_iterations_without_progress_threshold(15000);
trainer.set_max_num_epochs(400);
trainer.be_verbose();
// 5) Train
trainer.train(samples, labels);
// 6) Evaluate quickly on the training set
auto predicted = net(samples);
size_t correct = 0;
for (size_t i = 0; i < labels.size(); ++i)
if (predicted[i] == labels[i])
correct++;
double accuracy = (double)correct / labels.size();
cout << "Training accuracy (on this sample set): " << accuracy << "\n";
// 7) Save the model
net.clean();
serialize(model_file) << net;
cout << "Model saved to " << model_file << "\n";
}
// ----------------------------------------------------------------------------------------
// Generate mode
// ----------------------------------------------------------------------------------------
if (parser.option("generate"))
{
cout << "=== GENERATE MODE ===\n";
// 1) Load the trained model
using net_infer = my_transformer_cfg::network_type<false>;
net_infer net;
if (file_exists(model_file))
{
deserialize(model_file) >> net;
cout << "Loaded model from " << model_file << "\n";
}
else
{
cerr << "Error: model file not found. Please run --train first.\n";
return 0;
}
cout << my_transformer_cfg::model_info::describe() << endl;
cout << "Model parameters: " << count_parameters(net) << endl << endl;
// 2) Get the prompt from the included slm_data.h
std::string prompt_text = shakespeare_prompt;
if (prompt_text.empty())
{
cerr << "No prompt found in slm_data.h.\n";
return 0;
}
// If prompt is longer than max_seq_len, we keep only the first window
if (prompt_text.size() > (size_t)max_seq_len)
prompt_text.erase(prompt_text.begin() + max_seq_len, prompt_text.end());
// Convert prompt to a token sequence
const auto prompt_tokens = char_based_tokenize(prompt_text);
// Put into a dlib matrix
matrix<int, 0, 1> input_seq(max_seq_len, 1);
// Fill with pad if prompt is shorter than max_seq_len
for (long i = 0; i < max_seq_len; ++i)
{
if ((size_t)i < prompt_tokens.size())
input_seq(i, 0) = prompt_tokens[i];
else
input_seq(i, 0) = PAD_TOKEN;
}
cout << "\nInitial prompt:\n" << prompt_text << " (...)\n\n\nGenerated text:\n" << prompt_text;
// 3) Generate new text
// We'll predict one character at a time, then shift the window
for (int i = 0; i < generation_length; ++i)
{
const int next_char = net(input_seq); // single inference
// Print the generated character
cout << static_cast<char>(std::min(next_char, MAX_TOKEN_ID)) << flush;
// Shift left by 1
for (long i = 0; i < max_seq_len - 1; ++i)
input_seq(i, 0) = input_seq(i + 1, 0);
input_seq(max_seq_len - 1, 0) = std::min(next_char, MAX_TOKEN_ID);
}
cout << "\n\n(end of generation)\n";
}
return 0;
}
catch (exception& e)
{
cerr << "Exception thrown: " << e.what() << endl;
return 1;
}
}
/*
* This program demonstrates the training of a language model on about 15k sequences.
* The training process produces a data file of approximately 32MB on disk.
*
* - Transformer model configuration:
* + vocabulary size: 257
* + layers: 3
* + attention heads: 4
* + embedding dimension: 64
* + max sequence length: 80
* - Number of parameters: 8,247,496
*
* The training cab be done using the following command line:
* >./slm_basic_train_ex --train --shuffle
*
* After this phase, the model achieves perfect prediction accuracy (i.e acc=1).
* The generation option produces text that is very close to the original training data,
* as illustrated by the example below:
* > Generated text:
* > QUEEN ELIZABETH:
* > But thou didst kill my children.
* >
* > KING RICHARD III:
* > But in your daughter's womb I bury them:
* > Where in that nest of spicery they shall breed
* > Selves of themselves, to your recomforture.
* >
* > QUEEN ELIZABETH:
* > Shall I go win my daughter to thy will?
* >
* > KING RICHARD III:
* > And be a happy mother by the deed.
* >
* > QUEEN ELIZABETH:
* > I go. Write to me very shortly.
* > And you shall understand from me her mind.
*/