Build a Tiny Language Model

Table of Contents

1. A Language Model Is a Table of Odds
2. Counting Which Word Follows Which
3. From Counts to Probabilities
4. Random Numbers You Can Repeat
5. Rolling a Weighted Die
6. Generating Text
7. Why the Text Sounds Odd
8. More Context, Better Text
9. The Price of More Context
10. How This Differs From a Real LLM
11. Recipe: The bigram and rng Modules
12. Try It Yourself
13. Quick Review: Common Questions

1. A Language Model Is a Table of Odds

In What Is an LLM? you saw the loop at the heart of every language model: predict the next piece, append it, repeat. In Tokens and Tokenizers you saw that text is cut into small pieces called tokens.

Now we build the middle part: the thing that does the predicting. We will build the smallest honest version of it, a bigram model. A bigram is a pair of neighbouring words. The model's whole knowledge is one question:

Given the word I just saw, which word usually comes next, and how often?

That is it. No neural network, no math beyond counting and dividing. But it is a real language model: a thing that assigns a probability to each possible next word. The big models you use every day do the same job. They just look at far more context, use far more clever machinery, and learned from far more text.

Make sure you have the llmlab folder from the earlier chapters. This chapter adds src/rng.js and src/bigram.js (plus a small src/ngram.js at the end). All code is plain ESM JavaScript, so run every file with node, for example node demo-counts.js.

2. Counting Which Word Follows Which

Remember src/corpus.js from chapter 1: it exports corpus (a short text about Sokha's coffee shop) and words(text), which lowercases the text and splits it into a list of words. That list is our training data.

To "train" a bigram model, walk through the words once. For every pair (current word, next word) add one to a counter. Store the counters in a Map of Maps: the outer key is the current word, the inner map holds each next word and how many times it followed.

// src/bigram.js (first part)
export function trainBigram(tokens) {
  const model = new Map();
  for (let i = 0; i < tokens.length - 1; i++) {
    const current = tokens[i];
    const next = tokens[i + 1];
    if (!model.has(current)) model.set(current, new Map());
    const counts = model.get(current);
    counts.set(next, (counts.get(next) ?? 0) + 1);
  }
  return model;
}

The loop stops at length - 1 because the last word has no next word to learn from. The ?? 0 means "if we have never seen this pair, start from zero". Try it. Save this as demo-counts.js in the llmlab folder:

import { corpus, words } from "./src/corpus.js";
import { trainBigram } from "./src/bigram.js";

const tokens = words(corpus);
const model = trainBigram(tokens);

console.log(tokens.length);
console.log(new Set(tokens).size);
console.log(model.get("coffee"));
console.log(model.get("sokha"));
// 182
// 67
// Map(5) {
//   'shop' => 7,
//   'with' => 1,
//   'and' => 2,
//   'sokha' => 1,
//   'after' => 1
// }
// Map(5) {
//   'opens' => 1,
//   'makes' => 1,
//   'likes' => 1,
//   'says' => 1,
//   'cleans' => 1
// }

Our corpus has 182 words, but only 67 different ones. After "coffee" the word "shop" came 7 times, "and" twice, and three other words once each. After "sokha" five different verbs appeared, once each. All of this output is stable, because there is no randomness yet.

That Map is the model. If you printed all of it, you would see everything the model knows. This is the first big lesson: a language model is learned statistics about text, stored in some data structure. In a bigram model the structure is a table you can read with your eyes. In a real LLM it is billions of numbers you cannot read, but the job is the same.

3. From Counts to Probabilities

Counts are not probabilities yet. "shop appeared 7 times" means little until you divide by the total: after "coffee" the corpus has 7 + 2 + 1 + 1 + 1 = 12 follow-ups, so "shop" has a probability of 7 / 12, about 58%.

Add a function that does this division and sorts the answers from most likely to least likely:

// src/bigram.js (second part)
export function nextWordProbs(model, word) {
  const counts = model.get(word);
  if (!counts) return [];
  let total = 0;
  for (const c of counts.values()) total += c;
  return [...counts.entries()]
    .map(([next, count]) => ({ word: next, count, prob: count / total }))
    .sort((a, b) => b.count - a.count || a.word.localeCompare(b.word));
}

If the word is not in the model at all, we return an empty list: the model has never seen it, so it has no opinion. The sort puts the highest count first, and uses the word itself to break ties, so the order is the same on every computer. Save demo-probs.js:

import { corpus, words } from "./src/corpus.js";
import { trainBigram, nextWordProbs } from "./src/bigram.js";

const model = trainBigram(words(corpus));
for (const word of ["coffee", "dara"]) {
  console.log(`after "${word}":`);
  for (const o of nextWordProbs(model, word)) {
    const pct = (o.prob * 100).toFixed(0).padStart(3);
    console.log(`  ${o.word.padEnd(8)} ${o.count} time(s)  ${pct}%`);
  }
}
console.log(nextWordProbs(model, "banana"));
console.log(nextWordProbs(model, "app"));
// after "coffee":
//   shop     7 time(s)   58%
//   and      2 time(s)   17%
//   after    1 time(s)    8%
//   sokha    1 time(s)    8%
//   with     1 time(s)    8%
// after "dara":
//   writes   2 time(s)   29%
//   drinks   1 time(s)   14%
//   is       1 time(s)   14%
//   says     1 time(s)   14%
//   thinks   1 time(s)   14%
//   walks    1 time(s)   14%
// []
// [
//   { word: 'and', count: 1, prob: 0.2 },
//   { word: 'because', count: 1, prob: 0.2 },
//   { word: 'helps', count: 1, prob: 0.2 },
//   { word: 'needs', count: 1, prob: 0.2 },
//   { word: 'saves', count: 1, prob: 0.2 }
// ]

Look at what the model "believes". After "coffee" it expects "shop" more than half the time. After "dara" it thinks "writes" is the most likely, but five other verbs are possible too. A word it has never seen ("banana") gets an empty list: this is a real weakness, and a first hint at why real models use sub-word tokens (from the tokens chapter), where nothing is ever truly unknown.

All of this output is stable. The percentages are rounded to whole numbers, so they add up to roughly 100, not always exactly.

This list of probabilities is the output of a language model. Every LLM, from this toy to the biggest hosted ones, ends each step with a list like this: for every possible next token, a probability. The next chapter, Probabilities, Temperature and Sampling, is all about what you do with that list.

4. Random Numbers You Can Repeat

To pick the next word we need randomness: if the model picks "shop" 58% of the time and "and" 17% of the time, we need a random number to decide. Math.random() works, but every run gives different text, and that makes tutorials, tests, and bug reports painful. You cannot say "look at the third line of output" when the output changes every time.

The fix is a seeded random number generator. It looks random, but you give it a starting number (the seed), and the same seed always produces the same sequence. We use a tiny, well-known generator called mulberry32. You do not need to understand its bit tricks. Treat it like a box: put a seed in, get a function out that returns numbers from 0 up to (but not including) 1.

// src/rng.js
export function mulberry32(seed) {
  let a = seed >>> 0;
  return function () {
    a = (a + 0x6d2b79f5) >>> 0;
    let t = a;
    t = Math.imul(t ^ (t >>> 15), t | 1);
    t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}

It returns a function, so each call to that function gives the next number in the sequence. Save demo-rng.js:

import { mulberry32 } from "./src/rng.js";

const a = mulberry32(42);
console.log(a(), a(), a());

const b = mulberry32(42);
console.log(b(), b(), b());

const c = mulberry32(43);
console.log(c());

// 0.6011037519201636 0.44829055899754167 0.8524657934904099
// 0.6011037519201636 0.44829055899754167 0.8524657934904099
// 0.9998110907617956

The first two lines are identical because both generators started from seed 42. Seed 43 gives a completely different number. Every line of output here is stable: this is the whole point of a seed.

Real LLM services work the same way in spirit. They use a random number generator to pick tokens. Some let you set the seed, to make results more repeatable, and we will come back to that in the next chapter.

5. Rolling a Weighted Die

Now we turn a random number into a choice. Picture a ruler 1 unit long cut into pieces. After "coffee" the ruler has a 0.58 piece for "shop", a 0.17 piece for "and", and three 0.08 pieces. Throw a dart at a random spot between 0 and 1 and see which piece it hits. Bigger pieces get hit more often, which is exactly what "58% likely" means.

In code, walk through the options adding up their probabilities (a running total). The first time the running total goes above the random number, that option wins:

function pickWeighted(items, r) {
  let running = 0;
  for (const item of items) {
    running += item.prob;
    if (r < running) return item.word;
  }
  return items[items.length - 1].word; // guards against rounding error
}

The last line only matters when rounding makes the probabilities add up to 0.9999999 instead of 1, and r happens to land in that sliver. It is a safety net, and a good habit with floating point numbers.

This tiny function (a "weighted random choice") is called sampling. You will meet it again in chapter 4 under that name.

6. Generating Text

Everything is ready for the loop from chapter 1. Start with a word, look up its next-word probabilities, roll the weighted die, append the winner, and repeat with the new word as the current word:

// src/bigram.js (third part)
export function generate(model, startWord, length, rng = Math.random) {
  const out = [startWord];
  let current = startWord;
  for (let i = 1; i < length; i++) {
    const options = nextWordProbs(model, current);
    if (options.length === 0) break; // dead end: word never had a successor
    current = pickWeighted(options, rng());
    out.push(current);
  }
  return out.join(" ");
}

The rng parameter has a default of Math.random, so in a real app you get different text each time, but in our demos we pass a seeded generator. Save demo-generate.js:

import { corpus, words } from "./src/corpus.js";
import { trainBigram, generate } from "./src/bigram.js";
import { mulberry32 } from "./src/rng.js";

const model = trainBigram(words(corpus));

for (const seed of [1, 2, 3]) {
  console.log(seed, generate(model, "the", 14, mulberry32(seed)));
}
console.log(generate(model, "dara", 14, mulberry32(7)));
console.log(generate(model, "banana", 5, mulberry32(7)));
console.log(generate(model, "app", 5, mulberry32(7)));
// 1 the river and the small coffee shop needs good app saves time every morning
// 2 the morning and the app saves time every morning the coffee after lunch sokha
// 3 the morning and the coffee after lunch sokha cleans the river and dara says
// dara writes code for a phone sokha makes iced coffee and dara writes code
// banana
// app and dara walks home

You just built a text generator. Three seeds, three different sentences, and each one is repeatable: run the file again and you get the same lines. The "banana" line shows the dead end: the model does not know the word, so it stops after the start word. If you remove the seed (use Math.random), your output will vary on every run, but each sentence will still be built from the same table.

7. Why the Text Sounds Odd

Read the first output slowly: "the river and the small coffee shop needs good app saves time every morning". Every pair of neighbouring words is fine: "the river", "small coffee", "needs good", "saves time". But the sentence as a whole makes no sense. "Coffee shop needs good app saves time" is not a grammatical English sentence.

The reason is that the model has a one-word memory. When it picks the word after "shop", it only looks at "shop". It has completely forgotten that the sentence began with "the river", or that a sentence needs a verb and then an end. Each step is locally plausible, but nothing keeps the whole thing on track. The text drifts, like a person telling a story who forgets what they said two words ago.

Three problems show up in the output:

1. No long-range structure. Sentences wander and never finish.
2. Topic drift. The text jumps from the river to an app to lunch.
3. Copying and stitching. Everything comes from the training text. The model can only recombine what it saw. It cannot invent a word it never met.

These are exactly the problems that more context helps with, which is the next idea.

8. More Context, Better Text

What if the model remembers the last two words instead of one? That is a trigram model. Or three words: a 4-gram. In general, an n-gram model looks at the last n words. The code barely changes. Instead of using one word as the key, join the last n words into one string. Create src/ngram.js:

// src/ngram.js
export function trainNgram(tokens, n) {
  const model = new Map();
  for (let i = 0; i + n < tokens.length; i++) {
    const context = tokens.slice(i, i + n).join(" ");
    const next = tokens[i + n];
    if (!model.has(context)) model.set(context, new Map());
    const counts = model.get(context);
    counts.set(next, (counts.get(next) ?? 0) + 1);
  }
  return model;
}

export function generateNgram(model, n, startWords, length, rng = Math.random) {
  const out = [...startWords];
  while (out.length < length) {
    const context = out.slice(-n).join(" ");
    const counts = model.get(context);
    if (!counts) break; // never seen this context
    const entries = [...counts.entries()];
    const total = entries.reduce((sum, [, c]) => sum + c, 0);
    let r = rng() * total;
    let chosen = entries[entries.length - 1][0];
    for (const [word, count] of entries) {
      r -= count;
      if (r < 0) { chosen = word; break; }
    }
    out.push(chosen);
  }
  return out.join(" ");
}

Here trainNgram(tokens, 1) is exactly our bigram model, just with string keys, and trainNgram(tokens, 2) is a trigram model. (This version picks a word by subtracting counts from r * total, which is the same weighted die as before, in a different shape.) Compare them in demo-ngram.js:

import { corpus, words } from "./src/corpus.js";
import { trainNgram, generateNgram } from "./src/ngram.js";
import { mulberry32 } from "./src/rng.js";

const tokens = words(corpus);

for (const n of [1, 2, 3]) {
  const model = trainNgram(tokens, n);
  let options = 0;
  for (const counts of model.values()) options += counts.size;
  const avg = (options / model.size).toFixed(2);
  console.log(`n=${n}: ${model.size} contexts, ${avg} choices per context`);
}

const one = trainNgram(tokens, 1);
const two = trainNgram(tokens, 2);
console.log(generateNgram(one, 1, ["the"], 16, mulberry32(11)));
console.log(generateNgram(two, 2, ["the", "coffee"], 16, mulberry32(11)));
console.log(generateNgram(two, 2, ["the", "students"], 16, mulberry32(3)));
console.log(generateNgram(two, 2, ["dara", "writes"], 16, mulberry32(3)));

// How many 4-word windows in the output appear word-for-word in the corpus?
const text = tokens.join(" ");
function copied(model, n, start) {
  let hit = 0, total = 0;
  for (let seed = 1; seed <= 20; seed++) {
    const out = generateNgram(model, n, start, 20, mulberry32(seed)).split(" ");
    for (let i = 0; i + 4 <= out.length; i++) {
      total++;
      if (text.includes(out.slice(i, i + 4).join(" "))) hit++;
    }
  }
  return `${hit} of ${total}`;
}
console.log("bigram copied 4-word runs:", copied(one, 1, ["the"]));
console.log("trigram copied 4-word runs:", copied(two, 2, ["the", "coffee"]));
// n=1: 67 contexts, 1.94 choices per context
// n=2: 129 contexts, 1.25 choices per context
// n=3: 160 contexts, 1.07 choices per context
// the students order more iced tea the river is a phone sokha cleans the table by
// the coffee shop after school the students order iced coffee and the students laugh and the
// the students share a table dara says that a good table dara is a student and
// dara writes code at the table by the window dara writes code at the table by
// bigram copied 4-word runs: 130 of 340
// trigram copied 4-word runs: 294 of 340

Read the numbers carefully, because they tell the real story:

• With more context the model has more contexts (67, then 129, then 160) and fewer choices in each (1.94, then 1.25, then 1.07). The model becomes more certain.
• The trigram text is more fluent: "the coffee shop after school the students order iced coffee and the students laugh" reads almost like real sentences, and "dara writes code at the table by the window" is a clean, sensible line.
• But the last two lines show the cost. Of the 4-word runs the trigram produced, 294 of 340 (about 86%) are copied word for word from the training text. The bigram copied only 130 of 340 (about 38%). The more context the model uses, the more it just repeats its training data.

All the lines above are stable because every generator is seeded.

9. The Price of More Context

Why not use 10 words of context, then? Look at the table again: n=3 has 160 contexts and only 1.07 choices each. Almost every context was seen exactly once, which means the model can only say one thing: what followed it in the training text. That is memorization, not learning.

This is the central trade-off of language modelling:

With long contexts there is a second problem: most contexts never appeared in training. Ask a 10-word-context table about a sentence it has never seen, and model.get(context) returns undefined. The n-gram model has no way to say "this is similar to something I saw". It only matches exact words, so the table grows enormous and most of its rows are empty or single-use.

This wall is exactly why n-gram models stopped at small contexts and why researchers moved to neural networks. A neural network does not store a row per context. It learns a compact set of numbers that lets it say "this context is like those contexts I saw", even if the exact words are new. That is how modern models read thousands of words of context without a table that fills the universe. We build the beginning of that idea in How a Model Is Trained, and see how words become comparable numbers in Embeddings and Attention.

10. How This Differs From a Real LLM

Our toy and a real LLM have the same interface: context in, probability for each possible next token out. Inside they differ a lot. Here is an honest comparison:

Whole-word models like ours are a simplification. Real models use sub-word tokens, so they can handle any word, even a new one. Real models are also not literally n-gram tables, and calling a Transformer "a bigger bigram model" would be wrong. The point of the toy is to show the shape of the job (context in, a probability list out, then sampling) and to let you see every number.

11. Recipe: The bigram and rng Modules

Add this chapter's tools to llmlab. Create src/rng.js:

// src/rng.js
// Seeded random number generator (mulberry32).
// Same seed -> same sequence, so demos are reproducible.
export function mulberry32(seed) {
  let a = seed >>> 0;
  return function () {
    a = (a + 0x6d2b79f5) >>> 0;
    let t = a;
    t = Math.imul(t ^ (t >>> 15), t | 1);
    t ^= t + Math.imul(t ^ (t >>> 7), t | 61);
    return ((t ^ (t >>> 14)) >>> 0) / 4294967296;
  };
}

and src/bigram.js, which combines the three parts from this chapter:

// src/bigram.js
// A bigram language model: count which word follows which.
export function trainBigram(tokens) {
  const model = new Map();
  for (let i = 0; i < tokens.length - 1; i++) {
    const current = tokens[i];
    const next = tokens[i + 1];
    if (!model.has(current)) model.set(current, new Map());
    const counts = model.get(current);
    counts.set(next, (counts.get(next) ?? 0) + 1);
  }
  return model;
}

// Turn the counts for one word into probabilities, biggest first.
export function nextWordProbs(model, word) {
  const counts = model.get(word);
  if (!counts) return [];
  let total = 0;
  for (const c of counts.values()) total += c;
  return [...counts.entries()]
    .map(([next, count]) => ({ word: next, count, prob: count / total }))
    .sort((a, b) => b.count - a.count || a.word.localeCompare(b.word));
}

// Pick one item using probabilities and a number r in [0, 1).
function pickWeighted(items, r) {
  let running = 0;
  for (const item of items) {
    running += item.prob;
    if (r < running) return item.word;
  }
  return items[items.length - 1].word;
}

// Predict, append, repeat.
export function generate(model, startWord, length, rng = Math.random) {
  const out = [startWord];
  let current = startWord;
  for (let i = 1; i < length; i++) {
    const options = nextWordProbs(model, current);
    if (options.length === 0) break; // dead end: word never had a successor
    current = pickWeighted(options, rng());
    out.push(current);
  }
  return out.join(" ");
}

The public functions are trainBigram(tokens), nextWordProbs(model, word), and generate(model, startWord, length, rng). If you have an src/index.js, re-export them (export * from "./bigram.js"; and export * from "./rng.js";). Chapter 4 imports both modules, so keep these files exactly as shown.

12. Try It Yourself

1. Read the model. Print nextWordProbs(model, "iced") and nextWordProbs(model, "good"). Which word is most likely in each case? Do the percentages add up to 100?
2. Find the dead end. A word is a dead end if it has no entry in the model. Compare the last word of words(corpus) with the keys of the model. Is it a dead end in this corpus? What would you have to change in the text to create one, and what does generate do then?
3. Same seed, same text. Call generate(model, "the", 20, mulberry32(5)) three times with a fresh generator each time. Then call it three times with one shared generator (const rng = mulberry32(5)) and compare. Explain the difference.
4. Add your own text. Replace the corpus with a paragraph of your own (about 200 words) and train again. Does the generated text sound like your paragraph? What changes if you use only 20 words?
5. Find the "best" start word. Write a loop that generates 50 sentences of 10 words from each start word in the model and counts how many sentences stop early. Which start words are riskiest?
6. Try a 4-gram. Use trainNgram(tokens, 3) and generate from ["the", "coffee", "shop"]. How many 20-word outputs does it make before it hits a context it has never seen?

13. Quick Review: Common Questions

Questions beginners often ask about this chapter. Try answering each one yourself first, then click the question to open the answer.

Where to go next. You have a model that outputs a list of probabilities, and you have seen it pick from that list with a weighted die. The next chapter opens that step up: how scores become probabilities with softmax, what temperature really does, and how top-k and top-p sampling shape the output. Continue with Probabilities, Temperature and Sampling.