7 - Tokens and embeddings: how text becomes numbers¶
In chapter 7 the Minermont team wants to understand the large language models they already use every day. The first lesson at AIA is about the raw material: a neural network only works with numbers, so text has to be cut into tokens, each token gets a numerical ID, and each ID becomes a vector, its embedding. In this notebook we build each piece by hand:
- Import the libraries.
- Three ways to cut a sentence into tokens.
- Subword tokenization with BPE, from scratch.
- From tokens to IDs and embedding vectors.
- Learning embeddings from context: word2vec.
- Measuring similarity: cosine or Euclidean distance (Martha's notes, section 7.9).
Step 1: Import the libraries¶
import collections
import numpy as np
import matplotlib.pyplot as plt
np.set_printoptions(precision=3, suppress=True)
Step 2: Three ways to cut a sentence¶
The book compares three tokenization strategies with the sentence "The patient has a critical condition". Splitting by words is simple but needs a huge vocabulary and cannot handle words it has never seen. Splitting by characters handles any word but produces very long sequences and loses the notion of a word. Modern models use an intermediate solution, subwords.
sentence = "The patient has a critical condition"
word_tokens = sentence.split()
char_tokens = list(sentence)
print(f"Word tokens ({len(word_tokens)}): {word_tokens}")
print(f"Character tokens ({len(char_tokens)}): {char_tokens}")
Word tokens (6): ['The', 'patient', 'has', 'a', 'critical', 'condition'] Character tokens (36): ['T', 'h', 'e', ' ', 'p', 'a', 't', 'i', 'e', 'n', 't', ' ', 'h', 'a', 's', ' ', 'a', ' ', 'c', 'r', 'i', 't', 'i', 'c', 'a', 'l', ' ', 'c', 'o', 'n', 'd', 'i', 't', 'i', 'o', 'n']
Step 3: Subword tokenization with BPE¶
Byte Pair Encoding (BPE) learns its subwords from a corpus. It starts with every word split into characters (plus an end-of-word mark, _), and then repeats a single step: find the pair of adjacent symbols that appears most often in the corpus and merge it into a new symbol. Frequent words end up as a single token; rare words stay split into frequent pieces.
Our corpus is a list of medical words with how many times each one appears:
corpus_counts = {
'hypertension': 30, 'hypotension': 18, 'hyperglycemia': 22, 'hypoglycemia': 20, 'glycemia': 8, 'hyperthermia': 10,
'hypothermia': 12, 'tachycardia': 15, 'bradycardia': 12, 'cardiology': 20, 'cardiologist': 14, 'neurology': 16,
'neurologist': 12, 'dermatology': 14, 'dermatologist': 10, 'fibrillation': 16, 'defibrillator': 8, 'defibrillation': 6,
'patient': 60, 'patients': 40, 'treatment': 30, 'the': 120, 'of': 60, 'and': 60, 'with': 40, 'pain': 30, 'chest': 20,
'fever': 25, 'dose': 20, 'doses': 10, 'insulin': 15, 'test': 25, 'tests': 15, 'monitor': 20, 'ventilator': 12,
'doctor': 40, 'indicator': 8, 'dehydration': 12, 'decompensation': 6, 'detection': 10,
}
test_words = ['hyperglycemia', 'defibrillator', 'hypertensive', 'cardiomyopathy', 'patients']
def merge_pair(symbols, pair):
out, i = [], 0
while i < len(symbols):
if i < len(symbols) - 1 and (symbols[i], symbols[i + 1]) == pair:
out.append(symbols[i] + symbols[i + 1])
i += 2
else:
out.append(symbols[i])
i += 1
return tuple(out)
def train_bpe(counts, n_merges):
words = {tuple(w) + ('_',): c for w, c in counts.items()} # every word as a list of characters
merges = []
for _ in range(n_merges):
pairs = collections.Counter()
for symbols, c in words.items():
for pair in zip(symbols, symbols[1:]):
pairs[pair] += c # how often each adjacent pair appears
best = max(pairs, key=pairs.get)
merges.append(best)
words = {merge_pair(symbols, best): c for symbols, c in words.items()}
return merges
def tokenize(word, merges):
symbols = tuple(word) + ('_',)
for pair in merges: # apply the merges in the order they were learned
symbols = merge_pair(symbols, pair)
return list(symbols)
merges = train_bpe(corpus_counts, n_merges=80)
print("First 10 merges:", [a + b for a, b in merges[:10]])
subword_vocab = sorted({s for w in corpus_counts for s in tokenize(w, merges)})
print(f"Subword vocabulary: {len(subword_vocab)} tokens (the corpus has {len(corpus_counts)} different words)\n")
for w in test_words:
status = 'in the corpus' if w in corpus_counts else 'NEVER SEEN'
print(f"{w:<16} ({status:<13}) -> {tokenize(w, merges)}")
First 10 merges: ['at', 'en', 'th', 't_', 'n_', 'the', 'ati', 'hy', 'the_', 'r_'] Subword vocabulary: 66 tokens (the corpus has 40 different words) hyperglycemia (in the corpus) -> ['hyper', 'glycemia_'] defibrillator (in the corpus) -> ['de', 'fibrill', 'at', 'or_'] hypertensive (NEVER SEEN ) -> ['hyper', 'tens', 'i', 'v', 'e', '_'] cardiomyopathy (NEVER SEEN ) -> ['cardi', 'o', 'm', 'y', 'o', 'p', 'at', 'hy', '_'] patients (in the corpus) -> ['patients_']
The first merges are frequent pairs such as "at", "en" and "th", which soon grow into "the". With 80 merges, BPE splits hyperglycemia into [hyper, glycemia], exactly the book's example, and defibrillator into [de, fibrill, at, or]. The interesting cases are the words it has never seen. Hypertensive becomes [hyper, tens, i, v, e], reusing the prefix hyper-, and even cardiomyopathy can be written with known pieces (cardi, at, hy) plus single characters. No word is ever "unknown", while a frequent word like patients is a single token.
Step 4: From tokens to IDs and vectors¶
Each token of the vocabulary gets a number, its ID, and each ID selects a row of the embedding matrix $E$, of size $|V| \times d$ (vocabulary size by embedding dimension). As Emma explains in the book, this matrix is simply one more layer of the network: its values start at random and are adjusted by gradient descent together with all the other weights. Here we use the book's example with $d = 7$ and a still untrained, random matrix:
vocabulary = sorted(set(sentence.lower().split()))
token_to_id = {tok: i for i, tok in enumerate(vocabulary)}
ids = [token_to_id[t] for t in sentence.lower().split()]
rng = np.random.default_rng(0)
E = rng.normal(0, 0.5, size=(len(vocabulary), 7)) # embedding matrix |V| x d
vectors = E[ids] # looking up rows: that is all an embedding layer does
print(f"IDs: {dict(zip(sentence.lower().split(), ids))}")
print(f"Shape of the embedded sentence: {vectors.shape} ({vectors.shape[0]} tokens x {vectors.shape[1]} dimensions)")
print(vectors)
IDs: {'the': 5, 'patient': 4, 'has': 3, 'a': 0, 'critical': 2, 'condition': 1}
Shape of the embedded sentence: (6, 7) (6 tokens x 7 dimensions)
[[ 0.178 -0.327 -0.065 0.392 0.747 -0.63 0.757]
[-0.229 0.11 -0.505 -0.105 -0.08 0.27 0.107]
[ 0.683 -0.333 0.176 0.452 0.047 -0.372 -0.461]
[ 0.063 -0.066 0.32 0.052 -0.268 0.181 0.652]
[-0.623 -0.366 -0.272 -0.158 0.206 0.521 -0.064]
[ 0.474 -0.352 -0.633 -0.312 0.021 -1.163 -0.109]]
Step 5: Learning embeddings from context, word2vec¶
How do the vectors come to mean something? Noah summarizes it in the book: during training, tokens that appear in similar contexts end up close to each other. Word2vec (skip-gram with negative sampling) is the classic way to learn embeddings from that idea alone. For each word of a text, it looks at the words around it (its context) and adjusts the vectors so that:
- a word and a real neighbour get a high dot product, $\sigma(u_o \cdot v_c) \to 1$;
- a word and a few random words (negative samples) get a low one, $\sigma(u_k \cdot v_c) \to 0$.
Nobody tells the model what any word means. We train it on 4000 synthetic sentences about patients, staff, places, dates, symptoms and transport:
groups = {
'staff': ['doctor', 'physician', 'nurse', 'surgeon', 'pediatrician'],
'place': ['hospital', 'clinic', 'ward', 'pharmacy', 'laboratory'],
'month': ['january', 'march', 'may', 'july', 'october', 'december'],
'number': ['two', 'three', 'four', 'five', 'seven', 'ten'],
'symptom': ['fever', 'cough', 'headache', 'nausea', 'dizziness'],
'vehicle': ['car', 'bus', 'bicycle', 'train', 'taxi'],
'doc': ['doctor', 'physician'],
}
templates = [
'the {staff} examined the patient at the {place}', 'the {staff} works at the {place} since {month}',
'the patient saw the {staff} because of {symptom}', 'the patient had {symptom} for {number} days',
'the appointment is in {month} at the {place}', 'she came to the {place} by {vehicle} with {symptom}',
'he waited {number} hours for the {staff}', 'the {vehicle} left the {place} after {number} minutes',
'in {month} the {staff} treated {number} patients with {symptom}', 'the {staff} took the {vehicle} to the {place}',
'the {doc} prescribed antibiotics for the {symptom}', 'the {doc} signed the discharge report in {month}',
]
rng = np.random.default_rng(0)
def make_sentence():
template = templates[rng.integers(len(templates))]
return template.format(**{k: v[rng.integers(len(v))] for k, v in groups.items()}).split()
corpus = [make_sentence() for _ in range(4000)]
words = sorted({w for s in corpus for w in s})
word_id = {w: i for i, w in enumerate(words)}
print(f"Example: {' '.join(corpus[0])}")
print(f"Vocabulary: {len(words)} words, corpus: {sum(map(len, corpus))} tokens")
# (centre, context) pairs with a window of 2 words on each side
pairs = np.array([(word_id[s[i]], word_id[s[j]]) for s in corpus for i in range(len(s))
for j in range(max(0, i - 2), min(len(s), i + 3)) if j != i])
counts = np.bincount([word_id[w] for s in corpus for w in s], minlength=len(words))
noise = counts ** 0.75 / np.sum(counts ** 0.75) # word2vec's distribution for the negative samples
print(f"Training pairs: {len(pairs)}")
Example: the doctor prescribed antibiotics for the fever Vocabulary: 66 words, corpus: 31614 tokens Training pairs: 102456
def sigmoid(z):
return 1 / (1 + np.exp(-z))
def train_word2vec(pairs, n_words, dim=16, negatives=5, learning_rate=0.01, epochs=8, batch=64, seed=0):
rng = np.random.default_rng(seed)
W = rng.normal(0, 0.1, (n_words, dim)) # the embeddings we want (one row per word)
C = np.zeros((n_words, dim)) # auxiliary "context" vectors
losses = []
for _ in range(epochs):
order, total = rng.permutation(len(pairs)), 0.0
for start in range(0, len(pairs), batch):
centre, context = pairs[order[start:start + batch]].T
negative = rng.choice(n_words, (len(centre), negatives), p=noise)
v, u_pos, u_neg = W[centre], C[context], C[negative]
s_pos = sigmoid(np.sum(v * u_pos, axis=1)) # should go to 1
s_neg = sigmoid(np.einsum('bd,bkd->bk', v, u_neg)) # should go to 0
total -= np.sum(np.log(s_pos)) + np.sum(np.log(1 - s_neg))
g_pos, g_neg = (s_pos - 1)[:, None], s_neg[:, :, None] # derivatives of the loss
np.add.at(C, context, -learning_rate * g_pos * v)
np.add.at(C, negative.ravel(), -learning_rate * (g_neg * v[:, None, :]).reshape(-1, dim))
np.add.at(W, centre, -learning_rate * (g_pos * u_pos + np.sum(g_neg * u_neg, axis=1)))
losses.append(total / len(pairs))
return W, losses
W, losses = train_word2vec(pairs, len(words))
print("Mean loss per epoch:", np.round(losses, 3))
Mean loss per epoch: [2.334 1.768 1.725 1.717 1.711 1.708 1.708 1.709]
Now we ask each word for its nearest neighbours, measured with the cosine similarity of their vectors (Step 6 explains why cosine), and we project the vectors of the group words onto 2 dimensions to look at them:
unit = W / np.linalg.norm(W, axis=1, keepdims=True)
def cosine(a, b):
return float(unit[word_id[a]] @ unit[word_id[b]])
for w in ['doctor', 'march', 'three', 'fever', 'car', 'hospital']:
sims = unit @ unit[word_id[w]]
nearest = [words[i] for i in np.argsort(-sims)[1:5]]
print(f"{w:<9} nearest: {nearest}")
print()
for a, b in [('doctor', 'physician'), ('doctor', 'nurse'), ('doctor', 'car'), ('march', 'july'), ('fever', 'car')]:
print(f"cosine({a}, {b}) = {cosine(a, b):.3f}")
group_words = {g: [w.split()[-1] for w in ws] for g, ws in groups.items() if g != 'doc'}
X = np.array([unit[word_id[w]] for ws in group_words.values() for w in ws])
Xc = X - X.mean(axis=0)
coords = Xc @ np.linalg.svd(Xc, full_matrices=False)[2][:2].T # projection on the 2 main directions (PCA)
plt.figure(figsize=(8, 6))
k = 0
for g, ws in group_words.items():
pts = coords[k:k + len(ws)]
plt.scatter(*pts.T, s=40, label=g)
for (x, y), w in zip(pts, ws):
plt.annotate(w, (x, y), fontsize=8, xytext=(3, 3), textcoords='offset points')
k += len(ws)
plt.legend()
plt.title('Word2vec embeddings projected onto 2 dimensions')
plt.show()
doctor nearest: ['physician', 'surgeon', 'pediatrician', 'nurse'] march nearest: ['october', 'may', 'january', 'july'] three nearest: ['five', 'seven', 'two', 'ten'] fever nearest: ['nausea', 'dizziness', 'headache', 'cough'] car nearest: ['train', 'bus', 'bicycle', 'taxi'] hospital nearest: ['pharmacy', 'ward', 'laboratory', 'clinic'] cosine(doctor, physician) = 0.993 cosine(doctor, nurse) = 0.799 cosine(doctor, car) = 0.181 cosine(march, july) = 0.989 cosine(fever, car) = 0.258
Without being told anything about medicine or the calendar, word2vec has grouped the words by how they are used: the nearest neighbours of doctor are the rest of the staff, those of march are months, those of three are numbers, and so on. The similarities are graded. Doctor and physician, which appear in exactly the same contexts, have a cosine of 0.993; doctor and nurse, which share many contexts but not all, 0.799; doctor and car only 0.181. In the projection the six groups form separate clouds, as the months and numbers do in the book's figure. Real models learn the same kind of structure from billions of sentences and with hundreds of dimensions.
Step 6: Measuring similarity¶
Martha's notes compare two ways of measuring how similar two vectors are. The Euclidean distance measures the straight-line distance between the points; the cosine similarity measures the angle between the vectors, whatever their length:
$$d(\vec{v}_1, \vec{v}_2) = \sqrt{\sum_i (v_{1i} - v_{2i})^2} \qquad\qquad \cos(\vec{v}_1, \vec{v}_2) = \frac{\vec{v}_1 \cdot \vec{v}_2}{|\vec{v}_1|\,|\vec{v}_2|}$$
Her example uses three 2D vectors: "doctor" $(0.8, 0.6)$, "physician" $(0.7, 0.5)$ and "car" $(0.2, 0.9)$.
def euclidean(a, b):
return np.sqrt(np.sum((a - b) ** 2))
def cosine_similarity(a, b):
return a @ b / (np.linalg.norm(a) * np.linalg.norm(b))
doctor, physician, car = np.array([0.8, 0.6]), np.array([0.7, 0.5]), np.array([0.2, 0.9])
print(f"doctor-physician: Euclidean {euclidean(doctor, physician):.3f}, cosine {cosine_similarity(doctor, physician):.3f}")
print(f"doctor-car: Euclidean {euclidean(doctor, car):.3f}, cosine {cosine_similarity(doctor, car):.3f}")
long_physician = 3 * physician # same direction, three times longer
print(f"\ndoctor-(3 x physician): Euclidean {euclidean(doctor, long_physician):.3f}, "
f"cosine {cosine_similarity(doctor, long_physician):.3f}")
doctor-physician: Euclidean 0.141, cosine 1.000 doctor-car: Euclidean 0.671, cosine 0.759 doctor-(3 x physician): Euclidean 1.581, cosine 1.000
These are the numbers of Martha's notes: with both measures, doctor is closer to physician than to car (Euclidean distance 0.141 against 0.671; cosine 1.000 against 0.759). But if we make the physician vector three times longer, the Euclidean distance jumps from 0.141 to 1.581, as if the two words had nothing in common, while the cosine stays at 1.000: the direction, which is what carries the meaning, has not changed.
Martha adds a second argument: in high-dimensional spaces, distances "concentrate", and most points end up almost equidistant. We can check it with random vectors: for each dimension we measure how much farther the farthest point is than the nearest one, seen from a reference point.
rng = np.random.default_rng(1)
for dim in [2, 10, 100, 1000]:
points = rng.normal(size=(1000, dim))
d = np.linalg.norm(points[1:] - points[0], axis=1)
print(f"{dim:>5} dimensions: farthest / nearest distance = {d.max() / d.min():8.2f}")
2 dimensions: farthest / nearest distance = 91.05 10 dimensions: farthest / nearest distance = 3.78 100 dimensions: farthest / nearest distance = 1.50 1000 dimensions: farthest / nearest distance = 1.15
In 2 dimensions the farthest point is 91 times farther away than the nearest one; in 1000 dimensions, only 1.15 times. When almost all points are at about the same distance, the Euclidean distance can hardly tell near from far, which is one more reason why language models compare directions.
Conclusion¶
- Tokenization turns text into a sequence of IDs. Subword algorithms like BPE keep a manageable vocabulary and can represent any word, even one never seen, by splitting it into known pieces.
- An embedding layer is a matrix: each token ID selects a row, a vector of real numbers that the network adjusts during training.
- Meaning emerges from context: word2vec, which only learns to tell real neighbours from random words, places words used in similar ways close together.
- To compare embeddings, the cosine similarity is usually better than the Euclidean distance: it looks at directions, not lengths, and it copes better with the high-dimensional spaces of language models.
- The next notebook uses these vectors as the input of the mechanism that made large language models possible: attention.