exTransformer.cpp

1#include <oa/oa.h>
2
3namespace {
4constexpr oa::I32 VocabSize = 300;
5constexpr oa::I32 ContextLength = 16;
6constexpr oa::I32 ModelWidth = 32;
7constexpr oa::I32 HiddenWidth = 64;
8constexpr oa::I32 BatchSize = 64;
9constexpr oa::I32 TrainingSteps = 300;
10constexpr const char* Corpus =
11	"to be or not to be that is the question whether tis nobler in the mind "
12	"to suffer the slings and arrows of outrageous fortune or to take arms "
13	"against a sea of troubles and by opposing end them to be or not to be "
14	"that is the question whether tis nobler in the mind to suffer the slings "
15	"and arrows of outrageous fortune or to take arms against a sea of "
16	"troubles and by opposing end them to be or not to be that is the question "
17	"whether tis nobler in the mind to suffer the slings and arrows of "
18	"outrageous fortune or to take arms against a sea of troubles and by "
19	"opposing end them ";
20
21oa::Matrix tokenMatrix(const oa::Vector<oa::I32>& tokens, oa::I32 batch) {
22	return oa::FnMatrix::fromInt32(
23		oa::Span<const oa::I32>(tokens.data(), tokens.size()),
24		{batch, ContextLength}, oa::ScalarType::UInt32
25	);
26}
27
28void nextBatch(const oa::Vector<oa::I32>& tokens, oa::I64& cursor,
29	oa::Matrix& input, oa::Matrix& target) {
30	const oa::I64 limit = static_cast<oa::I64>(tokens.size()) - ContextLength - 1;
31	oa::Vector<oa::I32> x(static_cast<oa::Usize>(BatchSize * ContextLength));
32	oa::Vector<oa::I32> y(x.size());
33	for (oa::I32 batch = 0; batch < BatchSize; ++batch) {
34		const oa::I64 start = (cursor + static_cast<oa::I64>(batch) * 7) % limit;
35		for (oa::I32 position = 0; position < ContextLength; ++position) {
36			const auto output = static_cast<oa::Usize>(batch * ContextLength + position);
37			x[output] = tokens[static_cast<oa::Usize>(start + position)];
38			y[output] = tokens[static_cast<oa::Usize>(start + position + 1)];
39		}
40	}
41	cursor = (cursor + BatchSize) % limit;
42	input = tokenMatrix(x, BatchSize);
43	target = tokenMatrix(y, BatchSize);
44}
45
46oa::String generate(oa::NnTransformer& model, const oa::BpeTokenizer& tokenizer) {
47	const oa::String prompt = "to be";
48	auto promptTokens = tokenizer.encode(prompt.cStr());
49	oa::Vector<oa::I32> context(static_cast<oa::Usize>(ContextLength), 0);
50	const auto copied = oa::min<oa::Usize>(promptTokens.size(), context.size());
51	for (oa::Usize index = 0; index < copied; ++index) context[index] = promptTokens[index];
52	oa::I32 filled = oa::max<oa::I32>(1, static_cast<oa::I32>(copied));
53	oa::String result = prompt;
54	for (oa::I32 index = 0; index < 32; ++index) {
55		auto logits = model.forward(tokenMatrix(context, 1));
56		auto row = oa::FnMatrix::slice(logits, 0, filled - 1, filled);
57		const auto next = static_cast<oa::I32>(oa::FnMatrix::argmax(row.reshape({VocabSize})));
58		result += tokenizer.decode({next});
59		if (filled < ContextLength) {
60			context[static_cast<oa::Usize>(filled++)] = next;
61		} else {
62			for (oa::Usize token = 1; token < context.size(); ++token) context[token - 1] = context[token];
63			context[ContextLength - 1] = next;
64		}
65	}
66	return result;
67}
68} // namespace
69
70OA_MAIN("ExampleMlTransformer") {
71	oa::FnMatrix::setRngSeed(20260714ULL);
72	oa::BpeTokenizer tokenizer(VocabSize);
73	tokenizer.train(Corpus, VocabSize - 256);
74	if (tokenizer.vocabSize() != VocabSize) return 1;
75	const auto corpusTokens = tokenizer.encode(Corpus);
76
77	oa::NnTransformer model(VocabSize, ContextLength, ModelWidth, HiddenWidth);
78	auto parameters = model.allParameterPtrs();
79	oa::AdamW optimizer(parameters, 0.01F);
80	oa::MetricLoss lossMetric;
81	oa::CbProgressBar progress;
82	oa::CbSummary summary;
83	progress.addMetric(&lossMetric);
84	oa::ItTrainingConfig config;
85	config.totalSteps = TrainingSteps;
86	config.batchSize = BatchSize;
87	config.sequenceLength = ContextLength;
88	config.sequenceUnit = "token";
89	config.timerName = "example_transformer_step";
90	oa::ItTraining training(engine, optimizer, config);
91	training.addMetric(&lossMetric);
92	training.addCallback(&progress);
93	training.addCallback(&summary);
94
95	oa::print("\nOA SDK Example — BPE Transformer · all-position LM");
96	oa::print("Tokenizer: byte BPE · vocab={} · context={}", VocabSize, ContextLength);
97	oa::print("Model: NnTransformer(width={}, hidden={}, layers=1, heads=1)",
98		ModelWidth,
99		HiddenWidth
100	);
101	oa::print("Params: {} · AdamW(lr=0.01)", static_cast<long long>(model.numParameters()));
102	oa::print("Training: {} steps · batch={} · sequence={} tokens",
103		TrainingSteps,
104		BatchSize,
105		ContextLength
106	);
107
108	oa::I64 cursor = 0;
109	oa::Matrix input;
110	oa::Matrix target;
111	oa::F32 initialLoss = 0.0F;
112	while (not training.isDone()) {
113		nextBatch(corpusTokens, cursor, input, target);
114		optimizer.zeroGrad();
115		oa::GradientTape tape;
116		auto logits = model.forward(input);
117		auto loss = oa::FnLoss::crossEntropy(logits, target.reshape({target.numElements()}));
118		tape.backward(loss);
119		training.next(loss);
120		if (training.index() == 1) {
121			initialLoss = training.lastLoss();
122		}
123	}
124	if (not training.finish().isOk()) {
125		return 1;
126	}
127
128	const oa::F32 finalLoss = training.lastLoss();
129	if (not (finalLoss < initialLoss)) {
130		return 1;
131	}
132	const auto generated = generate(model, tokenizer);
133	oa::print("Transformer training verified: vocab=300, steps=300");
134	oa::print("Loss: {:.4f} -> {:.4f}", initialLoss, finalLoss);
135	oa::print("Generated: {}", generated.cStr());
136	return 0;
137}
138