6.5 Embeddings

import math, random

docs = [line.strip() for line in open('names.txt') if line.strip()]
LETTERS = '.abcdefghijklmnopqrstuvwxyz'
ID = {ch: i for i, ch in enumerate(LETTERS)}
pair_counts = {}
for name in docs:
    chars = '.' + name + '.'
    for a, b in zip(chars, chars[1:]):
        pair_counts[(ID[a], ID[b])] = pair_counts.get((ID[a], ID[b]), 0) + 1
TOTAL = sum(pair_counts.values())

# Two tables of knobs. wte: 2 numbers per letter (its embedding).
# lm_head: 2 numbers per possible NEXT letter. Same names as in microgpt.
rng = random.Random(42)
wte = [[rng.gauss(0, 0.5) for _ in range(2)] for _ in range(27)]
lm_head = [[rng.gauss(0, 0.5) for _ in range(2)] for _ in range(27)]


def softmax(logits):  # from lesson 6.4
    biggest = max(logits)
    exps = [math.exp(x - biggest) for x in logits]
    total = sum(exps)
    return [e / total for e in exps]


def linear(x, w):  # microgpt line 94: one dot product per row of w
    return [sum(wi * xi for wi, xi in zip(row, x)) for row in w]


# The forward pass: letter -> embedding -> 27 scores -> 27 chances.
def embed(token_id):
    return [0.0, 0.0]  # TODO: look up this letter's row in wte


def logits_for(token_id):
    return [0.0] * 27  # TODO: linear() of the embedding with lm_head


def surprise_of_pair(a, b):
    return -math.log(softmax(logits_for(a))[b])


def average_loss():
    return sum(pair_counts[(a, b)] * surprise_of_pair(a, b) for (a, b) in pair_counts) / TOTAL


# Training (given): gradient descent on both tables. The gradient formula `d`
# below was worked out by hand and you can take it on trust for now. It gives
# the same numbers 6.3's nudging would, just much faster. Module 7 shows where
# gradients like this come from, and how microgpt gets them automatically.
# The right stride depends on how steep the valley is; 5.0 suits this one.
def train(epochs, lr=5.0):
    for _ in range(epochs):
        g_wte = [[0.0, 0.0] for _ in range(27)]
        g_head = [[0.0, 0.0] for _ in range(27)]
        for a in range(27):
            row_total = sum(pair_counts.get((a, b), 0) for b in range(27))
            if not row_total:
                continue  # skip to the next a at once (break would leave the loop altogether)
            x, p = embed(a), softmax(logits_for(a))
            for b in range(27):
                d = (row_total * p[b] - pair_counts.get((a, b), 0)) / TOTAL
                for j in range(2):
                    g_head[b][j] += d * x[j]
                    g_wte[a][j] += d * lm_head[b][j]
        for i in range(27):
            for j in range(2):
                wte[i][j] -= lr * g_wte[i][j]
                lm_head[i][j] -= lr * g_head[i][j]


print('before training:', round(average_loss(), 4))
train(300)
print('after 300 steps:', round(average_loss(), 4))
for ch in 'aeiou' + 'bkmt':
    print(ch, [round(v, 2) for v in wte[ID[ch]]])