For a while now I've been trying to understand how language models actually work — not just how to call an API — and I wanted to do it in Swift rather than the usual Python. I'm not an ML expert; this was mostly me learning in the open.
Along the way I wrote a lot of small, heavily-commented code and notes, so I've cleaned them up and put them online as a kind of runnable textbook, in case it's useful to anyone else.
It's the same task solved four times — "predict what comes next" — each time with a little more machinery:
n-gram — a character-level counting model (a single self-contained script).
neural-char — a tiny reverse-mode autograd engine written from scratch
(micrograd-style), then a neural bigram and an MLP trained on top of it. Pure
Swift, no dependencies.
mini-gpt — a small char-level GPT (Transformer) on MLX Swift. Attention, causal
masking, residuals, LayerNorm and AdamW are written out explicitly rather
than calling a fused op, so each step is visible.
swift-rag — a small retrieval-augmented generation pipeline: BM25 over an
inverted index (from scratch), feeding a local Llama-2. The inference engine
is a vendored, unmodified llama2.swift.
A few Swift-relevant notes:
Three of the four parts are pure Swift + Foundation and build/run on Linux
and macOS; only the Transformer needs MLX (Apple Silicon).
Clarity is prioritised over performance everywhere — e.g. the autograd engine
differentiates individual scalars, which is slow but keeps the math visible.
It's a teaching tool, not a benchmark.
Each part has a README, a chapter-by-chapter walkthrough of the real code, and
SVG diagrams.
Time flies like an arrow, fruit flies like a banana.
Got the output:
Time flies like an arrow, fruit flies like an arrow, fruit flies like a banana.
Also, I have noticed that the output is not deterministic; that is, if I run the program multiple times I may get different output.
What does all this mean?
Refactored n-gram
// [https://forums.swift.org/t/a-small-build-it-to-understand-it-llm-textbook-written-in-swift/87804]
//
// ngram.swift
// LLM
//
// Created by ibex on 29/6/2026.
//
// SPDX-License-Identifier: Apache-2.0
import Foundation
struct ngram {
// MARK: - Settings
// ----------------------------------------------------------------------------
// Tunable knobs. Feel free to experiment — each has a clear, visible effect.
/// How many previous characters form the context (the "n" in n-gram, minus one
/// for the predicted character). Larger `order` means the model conditions on
/// more history, so the output looks more like real, memorized phrases — but it
/// also generalizes less and is likelier to hit "dead-end" contexts that were
/// never seen during training. Good values to try: 2 through 8.
let order: Int
/// How many characters of new text to generate after training.
let outputLength: Int
/// Controls how "adventurous" sampling is by reshaping the count distribution
/// before we draw from it (see `sample(from:)` for the exact math):
/// * temperature < 1 → sharpen toward the most frequent next character
/// (more predictable, more repetitive output).
/// * temperature = 1 → sample in proportion to the raw counts.
/// * temperature > 1 → flatten the distribution (more random, more chaotic).
let temperature: Double
// MARK: - Training: context -> counts of next characters
// ----------------------------------------------------------------------------
// "Training" here is pure bookkeeping: slide a window of length `order` across
// the text and, for each window, record which character came next and how often.
//
// The data structure is a nested dictionary:
//
// model[context][nextCharacter] = number of times that next character
// followed that context in the text.
//
// The context key is a `String` of exactly `order` characters. The inner
// dictionary is effectively a histogram over possible next characters.
private var model: [String: [Character: Int]] = [:]
let sample: String = """
Time flies like an arrow, fruit flies like a banana.
"""
init (order: Int = 4, outputLength: Int = 1000, temperature: Double = 0.8) {
self.order = order
self.outputLength = outputLength
self.temperature = temperature
}
}
private extension ngram {
// ============================================================================
// ngram.swift — a character-level n-gram language model in pure Swift.
// ============================================================================
//
// WHAT THIS PROGRAM IS
// --------------------
// A complete, self-contained character-level language model that you can run
// as a single-file Swift script (no package, no dependencies):
//
// swift ngram.swift # train on the built-in sample
// swift ngram.swift input.txt # train on your own UTF-8 text file
//
// It reads some text, learns the statistics of which character tends to follow
// the previous few characters, and then generates brand-new text with the same
// "flavor" as the training data.
//
// WHAT IS AN N-GRAM MODEL?
// ------------------------
// An "n-gram" is a contiguous run of n items — here, n characters. This model
// is built on a simple assumption (the "Markov assumption"): the next character
// depends only on the previous `order` characters, not on the entire history.
// So to predict the next character we only ever look at a short, fixed-length
// window of recent characters, which we call the "context".
//
// Crucially, this model does NOT use neural networks, gradients, or training in
// the machine-learning sense. It just COUNTS. For every context of length
// `order` that appears in the text, it tallies how often each possible next
// character follows it. Those raw counts ARE the model: a context that was
// followed by "e" 40 times and "a" 10 times will, when we generate, pick "e"
// about four times as often as "a".
//
// HOW THIS RELATES TO THE NEURAL PROJECT
// --------------------------------------
// The sibling project in the `neural-char/` folder solves the very same task —
// predicting the next character — but LEARNS the probabilities with gradient
// descent on top of a hand-written automatic-differentiation engine, instead of
// counting them directly. Comparing the two is illuminating:
//
// * This file: count frequencies → normalize → sample. Instant,
// exact, but cannot generalize beyond contexts it has
// literally seen (a never-before-seen context is a
// dead end).
// * neural-char/: start from random weights and nudge them to minimize
// a loss. Slower and approximate, but it learns smooth,
// shareable representations and degrades gracefully on
// unseen contexts.
//
// Both are character-level models; they differ only in HOW the next-character
// distribution is obtained.
// MARK: - Load text
// ----------------------------------------------------------------------------
// Obtain the training text: from a file path passed on the command line, or
// from a small built-in sample if no (readable) path was given.
/// Returns the full training text as a `String`.
///
/// `CommandLine.arguments` holds the process arguments; index 0 is always the
/// script/executable path itself, so the first user-supplied argument (a file
/// path, if any) is at index 1. We attempt to read that file as UTF-8 and fall
/// back to the built-in sample on any failure.
func loadText() -> String {
if CommandLine.arguments.count > 1 {
let path = CommandLine.arguments[1]
// `try?` turns a read error into nil rather than crashing, so a bad path
// simply falls through to the built-in sample below.
if let text = try? String(contentsOfFile: path, encoding: .utf8) {
return text
}
print("Could not read \(path) — using the built-in sample.")
}
// A tiny sample in case no file is provided.
// Replace with any text: a book, code, logs — anything.
return sample
}
// MARK: - Sample the next character from a distribution
// ----------------------------------------------------------------------------
/// Randomly choose one character from a histogram of counts, after reshaping the
/// histogram according to the global `temperature`.
///
/// THE TEMPERATURE MATH. Each raw count `c` is transformed into a weight
///
/// w = c ^ (1 / temperature)
///
/// and then we draw a character with probability proportional to its weight.
/// Why this formula?
/// * temperature = 1 → exponent 1 → weights equal the raw counts → we sample
/// exactly in proportion to observed frequency.
/// * temperature < 1 → exponent > 1 → large counts grow disproportionately,
/// so the most common next characters dominate even more
/// (sharper, more deterministic).
/// * temperature > 1 → exponent < 1 → counts are compressed toward each
/// other, giving rarer characters a better chance
/// (flatter, more random).
///
/// THE SAMPLING METHOD ("roulette wheel" / inverse-transform sampling). We sum
/// all weights to get `total`, draw a uniform random number `r` in [0, total),
/// then walk through the weights subtracting each from `r` until it goes
/// non-positive. The character whose weight tipped `r` past zero is our pick;
/// each character is chosen with probability weight/total, exactly as desired.
func sample(from counts: [Character: Int]) -> Character {
var weighted: [(Character, Double)] = []
var total = 0.0
for (ch, c) in counts {
let w = pow(Double(c), 1.0 / temperature)
weighted.append((ch, w))
total += w
}
var r = Double.random(in: 0..<total)
for (ch, w) in weighted {
r -= w
if r <= 0 { return ch }
}
// Fallback for floating-point rounding: if `r` never quite reached zero,
// return the last character we considered.
return weighted.last!.0
}
mutating func bootstrap (text: [Character]) {
if text.count > order {
// Start at index `order` so there are always `order` characters of history
// available behind the current position `i`.
for i in order..<text.count {
// The `order` characters immediately before position `i`.
let context = String(text[(i - order)..<i])
// The character at position `i` is what actually followed that context.
let next = text[i]
// Increment the count for (context → next). The `default:` subscripts
// create the dictionary entry on first use, so we never special-case the
// "first time we see this context/character" branch by hand.
model[context, default: [:]][next, default: 0] += 1
}
}
print("Unique contexts learned: \(model.count)")
}
mutating func generate (text: [Character]) {
if text.count > order {
// Start at index `order` so there are always `order` characters of history
// available behind the current position `i`.
for i in order..<text.count {
// The `order` characters immediately before position `i`.
let context = String(text[(i - order)..<i])
// The character at position `i` is what actually followed that context.
let next = text[i]
// Increment the count for (context → next). The `default:` subscripts
// create the dictionary entry on first use, so we never special-case the
// "first time we see this context/character" branch by hand.
model[context, default: [:]][next, default: 0] += 1
}
}
print("Unique contexts learned: \(model.count)")
// MARK: - Generation
// ----------------------------------------------------------------------------
// Generate new text autoregressively: repeatedly look at the last `order`
// characters produced so far, sample the next character from that context's
// histogram, append it, and continue. The model's own output becomes its next
// input — which is exactly what "autoregressive" means.
// Seed the generator with the first `order` characters of the training text, so
// we always begin from a context the model has actually seen.
var generated = Array(text.prefix(order))
for _ in 0..<outputLength {
// The current context is the last `order` characters generated so far.
let context = String(generated.suffix(order))
guard let counts = model[context] else {
// Dead end: this exact context never appeared in the training data, so
// the model has no idea what could come next. With pure counting there
// is nothing sensible to do but stop. (A neural model would not have
// this problem — it always produces SOME distribution. Larger `order`
// values make dead ends more likely.)
break
}
generated.append(sample(from: counts))
}
print("\n--- Generated text ---\n")
print(String(generated))
}
}
extension ngram {
mutating func run () {
// We convert the text to an `Array<Character>` once, up front. Swift `String`s
// are not randomly indexable by integer (a character can span several bytes),
// but an array of `Character` is — and we need to slice fixed-length windows out
// of it efficiently below.
let text = Array(loadText()) // [Character]
print("Characters loaded: \(text.count)")
bootstrap (text: text)
generate (text: text)
}
}
Thanks for running it — and for the refactor! Nothing's broken: you've hit the two built-in limits of a pure n-gram, which is exactly what it's meant to show.
Changes every run: it samples the next character via an unseeded Double.random, so each run takes different branches. Seed the RNG for reproducible output (the neural-char/ sibling shows how); temperature won't settle a tie where both options have count 1.
"arrow" repeats, "banana" rarely: ~50 characters with order = 4 means the model has basically memorised the sentence. It loops at the shared "flies like a…", while "banana." runs into a dead-end and generation stops (the break). Try order = 2 vs 8 to watch it shift.
Really like your struct ngram version — happy to add it with credit if you're up for it!
Nothing's broken: you've hit the two built-in limits of a pure n-gram, which is exactly what it's meant to show.