Skip to content

cds2.nlp

Educational NLP: BPE tokenizer, scalar autograd, attention, mini-GPT.

Educational NLP toolkit: BPE tokenizer, autograd, attention, mini-GPT.

AttentionWeights dataclass

Attention output plus the weight matrix for inspection.

Source code in src/cds2/nlp/attention.py 
27
28
29
30
31
32
@dataclass(frozen=True)
class AttentionWeights:
    """Attention output plus the weight matrix for inspection."""

    output: FloatArray
    weights: FloatArray

Value

A scalar that records the operations producing it for backprop.

Source code in src/cds2/nlp/autograd.py 
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
class Value:
    """A scalar that records the operations producing it for backprop."""

    def __init__(
        self,
        data: float,
        parents: tuple[Value, ...] = (),
        op: str = "",
    ) -> None:
        self.data = float(data)
        self.grad = 0.0
        self._parents = parents
        self._op = op
        self._backward: Callable[[], None] = lambda: None

    def __add__(self, other: object) -> Value:
        other_value = _coerce(other)
        result = Value(self.data + other_value.data, (self, other_value), "+")

        def propagate() -> None:
            self.grad += result.grad
            other_value.grad += result.grad

        result._backward = propagate
        return result

    def __radd__(self, other: object) -> Value:
        return self.__add__(other)

    def __mul__(self, other: object) -> Value:
        other_value = _coerce(other)
        result = Value(self.data * other_value.data, (self, other_value), "*")

        def propagate() -> None:
            self.grad += other_value.data * result.grad
            other_value.grad += self.data * result.grad

        result._backward = propagate
        return result

    def __rmul__(self, other: object) -> Value:
        return self.__mul__(other)

    def __truediv__(self, other: object) -> Value:
        other_value = _coerce(other)
        return self * other_value**-1.0

    def __neg__(self) -> Value:
        return self * -1.0

    def __sub__(self, other: object) -> Value:
        other_value = _coerce(other)
        return self + (-other_value)

    def __rsub__(self, other: object) -> Value:
        other_value = _coerce(other)
        return other_value + (-self)

    def __pow__(self, exponent: float) -> Value:
        if isinstance(exponent, Value):
            raise TypeError("Value supports only constant exponents")
        result = Value(self.data**exponent, (self,), f"**{exponent}")

        def propagate() -> None:
            self.grad += exponent * self.data ** (exponent - 1.0) * result.grad

        result._backward = propagate
        return result

    def exp(self) -> Value:
        result = Value(math.exp(self.data), (self,), "exp")

        def propagate() -> None:
            self.grad += result.data * result.grad

        result._backward = propagate
        return result

    def tanh(self) -> Value:
        result = Value(math.tanh(self.data), (self,), "tanh")

        def propagate() -> None:
            self.grad += (1.0 - result.data**2) * result.grad

        result._backward = propagate
        return result

    def relu(self) -> Value:
        result = Value(max(0.0, self.data), (self,), "relu")

        def propagate() -> None:
            self.grad += (self.data > 0) * result.grad

        result._backward = propagate
        return result

    def backward(self) -> None:
        """Backpropagate gradients through every ancestor operation."""
        ordered: list[Value] = []
        seen: set[int] = set()

        def visit(node: Value) -> None:
            if id(node) in seen:
                return
            seen.add(id(node))
            for parent in node._parents:
                visit(parent)
            ordered.append(node)

        visit(self)
        self.grad = 1.0
        for node in reversed(ordered):
            node._backward()

    def __repr__(self) -> str:
        return f"Value(data={self.data:g}, grad={self.grad:g})"

backward()

Backpropagate gradients through every ancestor operation.

Source code in src/cds2/nlp/autograd.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
def backward(self) -> None:
    """Backpropagate gradients through every ancestor operation."""
    ordered: list[Value] = []
    seen: set[int] = set()

    def visit(node: Value) -> None:
        if id(node) in seen:
            return
        seen.add(id(node))
        for parent in node._parents:
            visit(parent)
        ordered.append(node)

    visit(self)
    self.grad = 1.0
    for node in reversed(ordered):
        node._backward()

TinyGPT

Untrained-but-deterministic mini GPT for studying the forward pass.

Source code in src/cds2/nlp/gpt.py 
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
class TinyGPT:
    """Untrained-but-deterministic mini GPT for studying the forward pass."""

    def __init__(
        self,
        vocab_size: int,
        embedding_dim: int = 32,
        heads: int = 4,
        layers: int = 2,
        context_length: int = 64,
        seed: int | None = None,
    ) -> None:
        if embedding_dim % heads != 0:
            msg = "embedding_dim must be divisible by heads"
            raise ValueError(msg)
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim
        self.heads = heads
        self.layers = layers
        self.context_length = context_length
        rng = np.random.default_rng(seed)
        scale = 1.0 / np.sqrt(embedding_dim)
        self.token_embeddings = rng.normal(scale=scale, size=(vocab_size, embedding_dim))
        self.position_embeddings = rng.normal(scale=scale, size=(context_length, embedding_dim))
        self.blocks: list[dict[str, FloatArray]] = []
        for _ in range(layers):
            self.blocks.append(
                {
                    "w_q": rng.normal(scale=scale, size=(embedding_dim, embedding_dim)),
                    "w_k": rng.normal(scale=scale, size=(embedding_dim, embedding_dim)),
                    "w_v": rng.normal(scale=scale, size=(embedding_dim, embedding_dim)),
                    "w_o": rng.normal(scale=scale, size=(embedding_dim, embedding_dim)),
                    "w_in": rng.normal(scale=scale, size=(embedding_dim, 4 * embedding_dim)),
                    "w_out": rng.normal(
                        scale=1.0 / (4 * embedding_dim), size=(4 * embedding_dim, embedding_dim)
                    ),
                }
            )
        self.lm_head = rng.normal(scale=scale, size=(embedding_dim, vocab_size))

    def forward(self, token_ids: list[int]) -> FloatArray:
        """Return logits of shape (len(token_ids), vocab_size)."""
        if not token_ids:
            msg = "at least one token is required"
            raise ValueError(msg)
        if any(not 0 <= token < self.vocab_size for token in token_ids):
            msg = "token id outside the vocabulary"
            raise ValueError(msg)
        if len(token_ids) > self.context_length:
            msg = f"sequence longer than context ({self.context_length})"
            raise ValueError(msg)
        positions = np.arange(len(token_ids))
        hidden = self.token_embeddings[np.asarray(token_ids)] + self.position_embeddings[positions]
        causal_mask = np.tril(np.ones((len(token_ids), len(token_ids))))
        for block in self.blocks:
            queries = hidden @ block["w_q"]
            keys = hidden @ block["w_k"]
            values = hidden @ block["w_v"]
            attended = multi_head_attention(
                queries, keys, values, self.heads, mask=causal_mask
            ).output
            hidden = _layer_norm(hidden + attended @ block["w_o"])
            feed_forward = _gelu(hidden @ block["w_in"]) @ block["w_out"]
            hidden = _layer_norm(hidden + feed_forward)
        return np.asarray(hidden @ self.lm_head, dtype=float)

    def sample(
        self,
        prompt_tokens: list[int],
        max_new_tokens: int = 16,
        temperature: float = 1.0,
        seed: int | None = None,
    ) -> list[int]:
        """Autoregressively extend ``prompt_tokens`` by sampling the logits."""
        if temperature <= 0:
            msg = "temperature must be positive"
            raise ValueError(msg)
        rng_values = np.random.default_rng(seed)
        tokens = list(prompt_tokens)
        for _ in range(max_new_tokens):
            window = tokens[-self.context_length :]
            logits = self.forward(window)[-1] / temperature
            probabilities = _softmax(logits)
            next_token = int(rng_values.choice(self.vocab_size, p=probabilities))
            tokens.append(next_token)
        return tokens

forward(token_ids)

Return logits of shape (len(token_ids), vocab_size).

Source code in src/cds2/nlp/gpt.py 
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
def forward(self, token_ids: list[int]) -> FloatArray:
    """Return logits of shape (len(token_ids), vocab_size)."""
    if not token_ids:
        msg = "at least one token is required"
        raise ValueError(msg)
    if any(not 0 <= token < self.vocab_size for token in token_ids):
        msg = "token id outside the vocabulary"
        raise ValueError(msg)
    if len(token_ids) > self.context_length:
        msg = f"sequence longer than context ({self.context_length})"
        raise ValueError(msg)
    positions = np.arange(len(token_ids))
    hidden = self.token_embeddings[np.asarray(token_ids)] + self.position_embeddings[positions]
    causal_mask = np.tril(np.ones((len(token_ids), len(token_ids))))
    for block in self.blocks:
        queries = hidden @ block["w_q"]
        keys = hidden @ block["w_k"]
        values = hidden @ block["w_v"]
        attended = multi_head_attention(
            queries, keys, values, self.heads, mask=causal_mask
        ).output
        hidden = _layer_norm(hidden + attended @ block["w_o"])
        feed_forward = _gelu(hidden @ block["w_in"]) @ block["w_out"]
        hidden = _layer_norm(hidden + feed_forward)
    return np.asarray(hidden @ self.lm_head, dtype=float)

sample(prompt_tokens, max_new_tokens=16, temperature=1.0, seed=None)

Autoregressively extend prompt_tokens by sampling the logits.

Source code in src/cds2/nlp/gpt.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
def sample(
    self,
    prompt_tokens: list[int],
    max_new_tokens: int = 16,
    temperature: float = 1.0,
    seed: int | None = None,
) -> list[int]:
    """Autoregressively extend ``prompt_tokens`` by sampling the logits."""
    if temperature <= 0:
        msg = "temperature must be positive"
        raise ValueError(msg)
    rng_values = np.random.default_rng(seed)
    tokens = list(prompt_tokens)
    for _ in range(max_new_tokens):
        window = tokens[-self.context_length :]
        logits = self.forward(window)[-1] / temperature
        probabilities = _softmax(logits)
        next_token = int(rng_values.choice(self.vocab_size, p=probabilities))
        tokens.append(next_token)
    return tokens

BPETokenizer

Trainable sub-word tokenizer using greedy byte-pair merges.

Source code in src/cds2/nlp/tokenizer.py 
 8
 9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
class BPETokenizer:
    """Trainable sub-word tokenizer using greedy byte-pair merges."""

    END_OF_WORD = "</w>"

    def __init__(self, merges: int = 64) -> None:
        if merges < 1:
            msg = "merges must be at least 1"
            raise ValueError(msg)
        self.merges = merges
        self.merge_rules: list[tuple[str, str]] = []
        self.vocabulary: set[str] = set()

    def train(self, corpus: str) -> BPETokenizer:
        """Learn ``merges`` pair rules from whitespace-split words."""
        words: list[tuple[str, ...]] = []
        for raw_word in corpus.lower().split():
            cleaned = "".join(character for character in raw_word if character.isalnum())
            if cleaned:
                symbols = tuple(cleaned) + (self.END_OF_WORD,)
                words.append(symbols)

        pair_counts: dict[tuple[str, str], int] = {}
        for word in words:
            for index in range(len(word) - 1):
                pair_counts[word[index], word[index + 1]] = (
                    pair_counts.get((word[index], word[index + 1]), 0) + 1
                )

        for _ in range(self.merges):
            if not pair_counts:
                break
            best_pair, best_count = max(pair_counts.items(), key=lambda item: item[1])
            if best_count < 2:
                break
            self.merge_rules.append(best_pair)
            merged = best_pair[0] + best_pair[1]
            self.vocabulary.add(merged)

            words = [self._merge_in_word(word, best_pair) for word in words]
            pair_counts = {}
            for word in words:
                for index in range(len(word) - 1):
                    pair_counts[word[index], word[index + 1]] = (
                        pair_counts.get((word[index], word[index + 1]), 0) + 1
                    )

        for word in words:
            for symbol in word:
                self.vocabulary.add(symbol)
        return self

    @staticmethod
    def _merge_in_word(word: tuple[str, ...], pair: tuple[str, str]) -> tuple[str, ...]:
        merged_symbol = pair[0] + pair[1]
        output: list[str] = []
        index = 0
        while index < len(word):
            if index < len(word) - 1 and word[index] == pair[0] and word[index + 1] == pair[1]:
                output.append(merged_symbol)
                index += 2
            else:
                output.append(word[index])
                index += 1
        return tuple(output)

    def encode(self, word: str) -> list[str]:
        """Split ``word`` into learned sub-word symbols."""
        symbols = tuple("".join(character for character in word.lower() if character.isalnum()))
        if not symbols:
            return []
        symbols = symbols + (self.END_OF_WORD,)
        for pair in self.merge_rules:
            symbols = self._merge_in_word(symbols, pair)
        return list(symbols)

    def decode(self, tokens: list[str]) -> str:
        """Reconstruct text by joining symbols and dropping end markers."""
        text = "".join(tokens)
        return text.replace(self.END_OF_WORD, " ").strip()

train(corpus)

Learn merges pair rules from whitespace-split words.

Source code in src/cds2/nlp/tokenizer.py 
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
def train(self, corpus: str) -> BPETokenizer:
    """Learn ``merges`` pair rules from whitespace-split words."""
    words: list[tuple[str, ...]] = []
    for raw_word in corpus.lower().split():
        cleaned = "".join(character for character in raw_word if character.isalnum())
        if cleaned:
            symbols = tuple(cleaned) + (self.END_OF_WORD,)
            words.append(symbols)

    pair_counts: dict[tuple[str, str], int] = {}
    for word in words:
        for index in range(len(word) - 1):
            pair_counts[word[index], word[index + 1]] = (
                pair_counts.get((word[index], word[index + 1]), 0) + 1
            )

    for _ in range(self.merges):
        if not pair_counts:
            break
        best_pair, best_count = max(pair_counts.items(), key=lambda item: item[1])
        if best_count < 2:
            break
        self.merge_rules.append(best_pair)
        merged = best_pair[0] + best_pair[1]
        self.vocabulary.add(merged)

        words = [self._merge_in_word(word, best_pair) for word in words]
        pair_counts = {}
        for word in words:
            for index in range(len(word) - 1):
                pair_counts[word[index], word[index + 1]] = (
                    pair_counts.get((word[index], word[index + 1]), 0) + 1
                )

    for word in words:
        for symbol in word:
            self.vocabulary.add(symbol)
    return self

encode(word)

Split word into learned sub-word symbols.

Source code in src/cds2/nlp/tokenizer.py 
74
75
76
77
78
79
80
81
82
def encode(self, word: str) -> list[str]:
    """Split ``word`` into learned sub-word symbols."""
    symbols = tuple("".join(character for character in word.lower() if character.isalnum()))
    if not symbols:
        return []
    symbols = symbols + (self.END_OF_WORD,)
    for pair in self.merge_rules:
        symbols = self._merge_in_word(symbols, pair)
    return list(symbols)

decode(tokens)

Reconstruct text by joining symbols and dropping end markers.

Source code in src/cds2/nlp/tokenizer.py 
84
85
86
87
def decode(self, tokens: list[str]) -> str:
    """Reconstruct text by joining symbols and dropping end markers."""
    text = "".join(tokens)
    return text.replace(self.END_OF_WORD, " ").strip()

multi_head_attention(queries, keys, values, heads, mask=None)

Split the last dimension into heads and attend independently.

Source code in src/cds2/nlp/attention.py 
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
def multi_head_attention(
    queries: FloatArray,
    keys: FloatArray,
    values: FloatArray,
    heads: int,
    mask: FloatArray | None = None,
) -> AttentionWeights:
    """Split the last dimension into ``heads`` and attend independently."""
    if queries.shape[-1] % heads != 0:
        msg = "embedding dimension must be divisible by the head count"
        raise ValueError(msg)
    batch_shape = queries.shape[:-1]
    head_dim = queries.shape[-1] // heads

    def split(array: FloatArray) -> FloatArray:
        reshaped = array.reshape(*batch_shape, heads, head_dim)
        return np.swapaxes(reshaped, -2, -3)

    per_head = scaled_dot_product_attention(split(queries), split(keys), split(values), mask)
    merged = np.swapaxes(per_head.output, -2, -3).reshape(*batch_shape, heads * head_dim)
    return AttentionWeights(output=merged, weights=per_head.weights)

scaled_dot_product_attention(queries, keys, values, mask=None)

Classic scaled dot-product attention over a batch or 2-D input.

Source code in src/cds2/nlp/attention.py 
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
def scaled_dot_product_attention(
    queries: FloatArray,
    keys: FloatArray,
    values: FloatArray,
    mask: FloatArray | None = None,
) -> AttentionWeights:
    """Classic scaled dot-product attention over a batch or 2-D input."""
    dimension = queries.shape[-1]
    scores = queries @ np.swapaxes(keys, -1, -2) / np.sqrt(dimension)
    if mask is not None:
        scores = np.where(mask.astype(bool), scores, -np.inf)
    weights = softmax(scores, axis=-1)
    output = weights @ values
    return AttentionWeights(
        output=np.asarray(output, dtype=float), weights=np.asarray(weights, dtype=float)
    )

softmax(values, axis=-1)

Numerically stable softmax along axis.

Source code in src/cds2/nlp/attention.py 
20
21
22
23
24
def softmax(values: NDArray[np.float64], axis: int = -1) -> FloatArray:
    """Numerically stable softmax along ``axis``."""
    shifted = values - values.max(axis=axis, keepdims=True)
    exponentials = np.exp(shifted)
    return np.asarray(exponentials / exponentials.sum(axis=axis, keepdims=True))