jgpt

A tiny GPT, written from scratch in plain Java.

This is a small educational project — an excuse to build a working transformer-based language model end to end without any machine learning libraries. Everything (the autograd engine, the optimizer, the tokenizer and the transformer itself) is implemented by hand so you can read it and follow along.

It is inspired by Andrej Karpathy's microGPT.

What it does

The model is a character-level language model. For example, it can be trained on the included list of Ukrainian settlement names (app/data/ua-settlements-geonames.txt) and then generate new, made-up names in the same style. You can also give it a prefix and let it complete the name.

Under the hood it includes a small handwritten autograd engine, an Adam optimizer, a Unicode code-point tokenizer, and a configurable stack of transformer blocks with multi-head attention, RMS normalization, residual connections, and MLPs — all kept intentionally small.

Requirements

• Java 25 (the Gradle build will download a matching toolchain if needed)
• No other setup — the Gradle wrapper (./gradlew) is included

Quick start

The project ships without a pre-trained model. Train one from the included settlement dataset, then enter the interactive prompt:

./gradlew run --console=plain -q \
  --args="--data data/ua-settlements-geonames.txt --seq-length 30"

(On Windows, use gradlew.bat in place of ./gradlew.)

Paths passed through Gradle are resolved from the app/ subproject. This is why the command uses data/..., while paths passed to the native binary below use app/data/... when it is launched from the repository root.

The --console=plain -q flags keep Gradle's own progress bar and logging out of the way, so the model's training output and the interactive prompt stay clean and readable. They are recommended for the best experience.

During training the app prints the current and moving-average loss. Every --checkpoint-frequency iterations it saves a checkpoint and prints samples; it also saves the final model after training. With Gradle, the default output directory is app/checkpoints/ because checkpoints is relative to the subproject. Checkpoints default to the safetensors format (.safetensors). The tensor payload follows the standard format and can be read by tools such as PyTorch or NumPy; jgpt-specific architecture and tokenizer information is stored as safetensors metadata.

Once training is done, type a prefix at the |> prompt and press Enter to get a generated name, or just press Enter on an empty line (or type /q) to quit.

Example for the included Ukrainian settlement names dataset sampling:

|> за
ЗАРОНА
|> пере
ПЕРЕЗІВКА
|> над
НАДІНКИ
|> тест
ТЕСТОЛО
|> ін
ІНКОРОВО
|> по
ПОРОВИНЕ
|> /q
Good bye!

Reusing a trained model

To skip training and load an existing model, pass either its path or its bare file name. A bare name is also searched for under --checkpoints-dir:

./gradlew run --console=plain -q --args="--model jgpt-vocab-43-seq-30-emb-16-trans-1-attn-4.safetensors"

The repository also includes editable helper scripts:

• ./train_settlements.sh trains on Ukrainian settlement names.
• ./train_proverbs.sh uses a larger configuration for Ukrainian proverbs and defaults to a 16 GiB JVM heap.
• ./run.sh loads a trained model; edit its MODEL value to point at one of your checkpoints first.

More included datasets and their provenance are documented in app/data/README.md.

Command-line options

All options are passed to the app itself. With Gradle, wrap them in --args="..." (e.g. ./gradlew run -q --args="--data names.txt --iteration-count 500"); with the native binary, pass them directly. Every option has a sensible default, so the only one you normally need is --data (to train) or --model (to load). Run with --help to see this list at any time.

Input / output

Option	Default	Description
--data <file>	—	Training data file, one document per line. Required unless --model is given.
--model <file>	—	Load an existing model instead of training. The format is inferred from the file extension (.safetensors / .json / .bin).
--type <format>	safetensors	Save format for the trained model: safetensors, json, or bin (case-insensitive).
--checkpoints-dir <dir>	checkpoints	Directory where checkpoints and the final model are written. Relative to app/ under Gradle.
--checkpoint-frequency <n>	20	Save a checkpoint every n training iterations.
--sample-count <n>	10	Number of sample names printed at each checkpoint.

Sampling

Option	Default	Description
--temperature <t>	0.5	Sampling temperature. 0 uses greedy decoding; positive values sample, with higher values producing more randomness.

Training

Option	Default	Description
--iteration-count <n>	1000	Number of training iterations.
--batch-size <n>	32	Training batch size.
--learning-rate <r>	0.01	Adam optimizer learning rate.
--seed <n>	1234	RNG seed, for reproducible training.

Model architecture

These take effect only when training a new model (they are read from the file when loading with --model).

Option	Default	Description
--embedding-dim <n>	16	Embedding dimension.
--seq-length <n>	20	Maximum sequence length.
--transformer-blocks <n>	1	Number of transformer blocks.
--attention-heads <n>	4	Number of attention heads.

For a conventional even head split, choose an embedding dimension divisible by the number of attention heads.

Help

Option	Description
-h, --help	Show usage help and exit.
-V, --version	Print version information and exit.

Building a native binary

The app can be compiled ahead-of-time into a single, self-contained native executable with GraalVM Native Image — no JVM needed to run it, and it starts instantly.

This needs a GraalVM JDK 25 with Native Image on the build machine. The Gradle build selects a matching GraalVM toolchain when one is available. If it cannot find one, point JAVA_HOME at it (for example, export JAVA_HOME=/path/to/graalvm-jdk-25).

./gradlew nativeCompile

The binary is written to app/build/native/nativeCompile/jgpt. Run it like the CLI above — train a fresh model, or load an existing one:

# Train from scratch, then drop into the interactive prompt
./app/build/native/nativeCompile/jgpt --data app/data/ua-settlements-geonames.txt --seq-length 30

# Or load a previously trained model and sample from it
./app/build/native/nativeCompile/jgpt --model checkpoints/jgpt-vocab-43-seq-30-emb-16-trans-1-attn-4.safetensors

All three model formats (.safetensors, .json, .bin) work in the native binary. The picocli command metadata is generated at compile time by the picocli-codegen annotation processor, and the Java-serialization metadata for the .bin format lives in app/src/main/resources/META-INF/native-image/dev.alvo/jgpt/reachability-metadata.json.

Note: run the produced binary directly for the interactive prompt — Gradle's nativeRun task runs the executable with an empty stdin, so the REPL would exit immediately.

Tests

./gradlew test