Now, let's build your very own AI.

These are the small course defaults. width is the size of a token vector; context is the number of visible tokens; learning_rate is the update size. Keep the same settings when resuming.

{
    "width": 8,
    "context": 12,
    "merges": 12,
    "steps": 20,
    "learning_rate": 0.03,
    "seed": 7,
    "tokens": 24,
    "max_tokens": 64,
    "port": 8765
}

Type one Q/A pair per line. These invented English examples keep the first experiment short. Replace them later with your own consistent text.

Q: hello | A: hello friend.
Q: hi | A: hi friend.
Q: name | A: i am xsml.
Q: sun | A: the sun is warm.
Q: cat | A: a cat can nap.
Q: dog | A: a dog can run.
Q: red | A: red is a color.
Q: blue | A: blue is a color.
Q: learn | A: i learn from text.
Q: code | A: i can learn code.
Q: moon | A: the moon is in the sky.
Q: rain | A: rain is water.
Q: book | A: i can read a book.
Q: tea | A: tea is warm.
Q: bye | A: bye friend.

Read UTF-8 text → remove duplicates → split whole records → write train and validation files. Separating records before tokenization avoids sharing the exact same example between both sets.

from pathlib import Path


lines = [line.strip() for line in Path('data/raw.txt').read_text(encoding='utf-8').splitlines() if line.strip()]
lines = list(dict.fromkeys(lines))
assert len(lines) >= 5, 'write at least five different Q/A lines'
assert all(' | A:' in line and line.startswith('Q:') for line in lines)
train = [line for i, line in enumerate(lines) if i % 5 != 0]
val = [line for i, line in enumerate(lines) if i % 5 == 0]
for name, rows in [('train', train), ('val', val)]:
    text = '\n'.join(row.replace(' | A:', '\nA:') for row in rows) + '\n'
    Path('data/' + name + '.txt').write_text(text, encoding='utf-8')
    print(name, len(rows))

merge combines one pair without overlapping. train learns the ordered rules. encode replays them; decode joins vocabulary pieces. The final assertion checks reversibility.

import json
from collections import Counter
from pathlib import Path


def merge(ids, pair, new_id):
    result, i = [], 0
    while i < len(ids):
        if tuple(ids[i:i + 2]) == tuple(pair):
            result.append(new_id)
            i += 2
        else:
            result.append(ids[i])
            i += 1
    return result


class BPE:
    def __init__(self, vocab, merges):
        self.vocab, self.merges = vocab, merges
        self.chars = {c: i for i, c in enumerate(vocab) if len(c) == 1}

    @classmethod
    def train(cls, text, count):
        vocab = ['\ufffd'] + sorted(set(text) - {'\ufffd'})
        ids = [vocab.index(c) for c in text]
        merges = []
        for _ in range(count):
            pairs = Counter((a, b) for a, b in zip(ids, ids[1:])
                            if '\n' not in vocab[a] + vocab[b])
            if not pairs:
                break
            pair, frequency = pairs.most_common(1)[0]
            if frequency < 2:
                break
            new_id = len(vocab)
            vocab.append(vocab[pair[0]] + vocab[pair[1]])
            merges.append([list(pair), new_id])
            ids = merge(ids, pair, new_id)
        return cls(vocab, merges)

    def encode(self, text):
        ids = [self.chars.get(c, 0) for c in text]
        for pair, new_id in self.merges:
            ids = merge(ids, pair, new_id)
        return ids

    def decode(self, ids):
        return ''.join(self.vocab[i] for i in ids)

    def state(self):
        return {'vocab': self.vocab, 'merges': self.merges}


if __name__ == '__main__':
    settings = json.loads(Path('settings.json').read_text())
    text = Path('data/train.txt').read_text(encoding='utf-8')
    tokenizer = BPE.train(text, settings['merges'])
    Path('data/tokenizer.json').write_text(
        json.dumps(tokenizer.state(), ensure_ascii=False), encoding='utf-8')
    assert tokenizer.decode(tokenizer.encode(text)) == text
    print('vocab:', len(tokenizer.vocab), 'merges:', len(tokenizer.merges))
    print('roundtrip: OK', tokenizer.encode('Q: hi\nA:'))

data is a number; grad is its effect on the loss; parents store local derivatives. Addition distributes a gradient; multiplication scales it by the other input. softmax converts scores to probabilities.

import math


class Value:
    def __init__(self, data, parents=()):
        self.data = float(data)
        self.grad = 0.0
        self.parents = parents

    def __add__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data + other.data, ((self, 1), (other, 1)))

    def __mul__(self, other):
        other = other if isinstance(other, Value) else Value(other)
        return Value(self.data * other.data,
                     ((self, other.data), (other, self.data)))

    def __pow__(self, power):
        return Value(self.data ** power,
                     ((self, power * self.data ** (power - 1)),))

    def __neg__(self):
        return self * -1

    def __sub__(self, other):
        return self + (-other)

    def __truediv__(self, other):
        return self * (other ** -1)

    __radd__ = __add__
    __rmul__ = __mul__

    def exp(self):
        result = math.exp(self.data)
        return Value(result, ((self, result),))

    def log(self):
        return Value(math.log(self.data), ((self, 1 / self.data),))

    def relu(self):
        return Value(max(0, self.data), ((self, float(self.data > 0)),))

    def backward(self):
        order, seen, stack = [], set(), [(self, False)]
        while stack:
            node, expanded = stack.pop()
            if expanded:
                order.append(node)
            elif node not in seen:
                seen.add(node)
                stack.append((node, True))
                stack.extend((parent, False) for parent, _ in node.parents)
        for node in order:
            node.grad = 0.0
        self.grad = 1.0
        for node in reversed(order):
            for parent, derivative in node.parents:
                parent.grad += derivative * node.grad


def softmax(values):
    peak = max(v.data for v in values)
    exps = [(v - peak).exp() for v in values]
    total = sum(exps)
    return [v / total for v in exps]


def linear(x, weights):
    return [sum(a * b for a, b in zip(x, row)) for row in weights]


def norm(x):
    mean = sum(x) / len(x)
    variance = sum((v - mean) ** 2 for v in x) / len(x)
    return [(v - mean) * (variance + 1e-5) ** -0.5 for v in x]


if __name__ == '__main__':
    x = Value(3)
    y = x * x + 2 * x
    y.backward()
    print('y =', y.data, 'dy/dx =', x.grad)
    assert y.data == 15 and x.grad == 8

Read shapes first, then forward. Every Value in weights is trainable. Lists store vectors and matrices. loss compares each position with its next-token target. state and load serialize numbers, not Python objects.

import random
from engine import Value, linear, norm, softmax


class GPT:
    def __init__(self, vocab_size, width, context, seed=7):
        self.width, self.context = width, context
        rng = random.Random(seed)
        shapes = {
            'token': (vocab_size, width), 'position': (context, width),
            'query': (width, width), 'key': (width, width),
            'value': (width, width), 'project': (width, width),
            'expand': (2 * width, width), 'contract': (width, 2 * width),
            'head': (vocab_size, width),
        }
        self.weights = {
            name: [[Value(rng.gauss(0, 0.1)) for _ in range(cols)]
                   for _ in range(rows)] for name, (rows, cols) in shapes.items()
        }
        self.parameters = [v for matrix in self.weights.values()
                           for row in matrix for v in row]

    def forward(self, ids):
        assert 0 < len(ids) <= self.context
        w, keys, values, output = self.weights, [], [], []
        for position, token in enumerate(ids):
            x = [a + b for a, b in zip(w['token'][token], w['position'][position])]
            h = norm(x)
            q = linear(h, w['query'])
            keys.append(linear(h, w['key']))
            values.append(linear(h, w['value']))
            scores = [sum(a * b for a, b in zip(q, k)) / self.width ** 0.5
                      for k in keys]
            attention = softmax(scores)
            mixed = [sum(p * v[i] for p, v in zip(attention, values))
                     for i in range(self.width)]
            x = [a + b for a, b in zip(x, linear(mixed, w['project']))]
            hidden = [v.relu() for v in linear(norm(x), w['expand'])]
            x = [a + b for a, b in zip(x, linear(hidden, w['contract']))]
            output.append(linear(norm(x), w['head']))
        return output

    def loss(self, inputs, targets):
        logits = self.forward(inputs)
        losses = [-softmax(row)[target].log()
                  for row, target in zip(logits, targets)]
        return sum(losses) / len(losses)

    def state(self):
        return {name: [[v.data for v in row] for row in matrix]
                for name, matrix in self.weights.items()}

    def load(self, state):
        assert set(state) == set(self.weights), 'checkpoint keys'
        for name, matrix in self.weights.items():
            assert len(state[name]) == len(matrix), 'checkpoint rows'
            for row, saved in zip(matrix, state[name]):
                assert len(saved) == len(row), 'checkpoint columns'
                for value, number in zip(row, saved):
                    value.data = float(number)

The data hash prevents resuming on silently changed data. Saving through a temporary file and os.replace preserves the previous complete checkpoint if writing is interrupted. SGD has no momentum state to save.

import argparse
import hashlib
import json
import math
import os
import random
from pathlib import Path
from model import GPT
from tokenizer import BPE


def main():
    cfg = json.loads(Path('settings.json').read_text())
    parser = argparse.ArgumentParser()
    parser.add_argument('--steps', type=int, default=cfg['steps'])
    args = parser.parse_args()
    assert args.steps > 0, '--steps must be positive'
    token_state = json.loads(Path('data/tokenizer.json').read_text())
    tokenizer = BPE(**token_state)
    train_text = Path('data/train.txt').read_text(encoding='utf-8')
    val_text = Path('data/val.txt').read_text(encoding='utf-8')
    fingerprint = hashlib.sha256((train_text + '\0' + val_text).encode()).hexdigest()
    train, val = tokenizer.encode(train_text), tokenizer.encode(val_text)
    model = GPT(len(tokenizer.vocab), cfg['width'], cfg['context'], cfg['seed'])
    rng, step = random.Random(cfg['seed']), 0
    checkpoint = Path('checkpoints/model.json')
    if checkpoint.exists():
        saved = json.loads(checkpoint.read_text())
        assert saved['settings'] == cfg and saved['tokenizer'] == token_state, 'changed settings/tokenizer: start a new checkpoint'
        assert saved['data_hash'] == fingerprint, 'changed data: start a new checkpoint'
        model.load(saved['weights'])
        step = saved['step']
        rng.setstate((saved['random'][0], tuple(saved['random'][1]), saved['random'][2]))
    n = model.context
    assert len(train) > n and len(val) > n, 'add more text'
    print('parameters:', len(model.parameters), 'resume:', step, flush=True)
    print('validation:', round(model.loss(val[:n], val[1:n + 1]).data, 4), flush=True)
    checkpoint.parent.mkdir(exist_ok=True)
    for _ in range(args.steps):
        start = rng.randrange(len(train) - n)
        window = train[start:start + n + 1]
        for p in model.parameters:
            p.grad = 0
        loss = model.loss(window[:-1], window[1:])
        loss.backward()
        size = math.sqrt(sum(p.grad * p.grad for p in model.parameters))
        scale = min(1.0, 1.0 / (size + 1e-8))
        for p in model.parameters:
            p.data -= cfg['learning_rate'] * p.grad * scale
        step += 1
        state = {'settings': cfg, 'tokenizer': token_state, 'data_hash': fingerprint,
                 'weights': model.state(), 'step': step, 'random': rng.getstate()}
        temporary = checkpoint.with_suffix('.tmp')
        temporary.write_text(json.dumps(state, ensure_ascii=False), encoding='utf-8')
        os.replace(temporary, checkpoint)
        print('step:', step, 'loss:', round(loss.data, 4), flush=True)
    print('validation:', round(model.loss(val[:n], val[1:n + 1]).data, 4), flush=True)
    print('saved:', checkpoint, flush=True)


if __name__ == '__main__':
    main()

load rebuilds the architecture and restores numbers. generate crops context, scores candidates and samples repeatedly. A fixed seed makes comparisons reproducible.

import argparse
import json
import math
import random
from pathlib import Path
from model import GPT
from tokenizer import BPE


def load():
    saved = json.loads(Path('checkpoints/model.json').read_text())
    cfg = saved['settings']
    tokenizer = BPE(**saved['tokenizer'])
    model = GPT(len(tokenizer.vocab), cfg['width'], cfg['context'], cfg['seed'])
    model.load(saved['weights'])
    return model, tokenizer, cfg


def generate(model, tokenizer, prompt, tokens=24, temperature=0.8, top_k=8, seed=7):
    if not prompt or tokens < 1 or temperature < 0 or not math.isfinite(temperature) or top_k < 0:
        raise ValueError('invalid sampling parameters')
    ids, result, rng = tokenizer.encode(prompt), [], random.Random(seed)
    for _ in range(tokens):
        logits = model.forward(ids[-model.context:])[-1]
        ranked = sorted(range(len(logits)), key=lambda i: logits[i].data, reverse=True)
        if temperature == 0:
            token = ranked[0]
        else:
            candidates = ranked[:top_k] if top_k else ranked
            peak = logits[candidates[0]].data
            weights = [math.exp((logits[i].data - peak) / temperature) for i in candidates]
            token = rng.choices(candidates, weights=weights)[0]
        ids.append(token)
        if tokenizer.vocab[token] == '\n':
            break
        result.append(token)
    return tokenizer.decode(result)


if __name__ == '__main__':
    model, tokenizer, cfg = load()
    parser = argparse.ArgumentParser()
    parser.add_argument('--prompt', default='Q: hi\nA:')
    parser.add_argument('--tokens', type=int, default=cfg['tokens'])
    parser.add_argument('--temperature', type=float, default=0.8)
    parser.add_argument('--top-k', type=int, default=8)
    parser.add_argument('--seed', type=int, default=cfg['seed'])
    args = parser.parse_args()
    assert 1 <= args.tokens <= cfg['max_tokens']
    print(generate(model, tokenizer, args.prompt, args.tokens,
                   args.temperature, args.top_k, args.seed))

One model instance is shared by requests. Validate the body, generate, then return UTF-8 JSON. Binding to the guest loopback keeps this exercise local to your Linux session.

import json
from http.server import BaseHTTPRequestHandler, HTTPServer
from generate import load, generate


def main():
    model, tokenizer, cfg = load()

    class Handler(BaseHTTPRequestHandler):
        def reply(self, status, body):
            data = json.dumps(body, ensure_ascii=False).encode('utf-8')
            self.send_response(status)
            self.send_header('Content-Type', 'application/json; charset=utf-8')
            self.send_header('Content-Length', str(len(data)))
            self.end_headers()
            self.wfile.write(data)

        def do_GET(self):
            if self.path == '/health':
                self.reply(200, {'status': 'ok', 'model': 'XSML'})
            else:
                self.reply(404, {'error': 'not found'})

        def do_POST(self):
            if self.path != '/generate':
                return self.reply(404, {'error': 'not found'})
            try:
                length = int(self.headers.get('Content-Length', '0'))
                if not 0 < length <= 4096:
                    raise ValueError('request size')
                data = json.loads(self.rfile.read(length))
                prompt = data['prompt']
                tokens = data.get('tokens', cfg['tokens'])
                temperature = data.get('temperature', 0.8)
                top_k = data.get('top_k', 8)
                seed = data.get('seed', cfg['seed'])
                if not isinstance(prompt, str) or not 0 < len(prompt) <= 512:
                    raise ValueError('prompt')
                if type(tokens) is not int or not 1 <= tokens <= cfg['max_tokens']:
                    raise ValueError('tokens')
                if type(top_k) is not int or top_k < 0 or type(seed) is not int:
                    raise ValueError('top_k/seed')
                if type(temperature) not in (int, float):
                    raise ValueError('temperature')
                result = generate(model, tokenizer, prompt, tokens, temperature, top_k, seed)
                self.reply(200, {'text': result, 'model': 'XSML'})
            except (ValueError, KeyError, TypeError):
                self.reply(400, {'error': 'invalid request'})

    address = ('127.0.0.1', cfg['port'])
    print('XSML Runner:', address, flush=True)
    HTTPServer(address, Handler).serve_forever()


if __name__ == '__main__':
    main()

input reads a line; Request creates the JSON POST; urlopen reads the real response. Connection errors are reported instead of replaced with a fake answer.

import json
from pathlib import Path
from urllib.request import Request, urlopen
from urllib.error import URLError


def ask(question):
    cfg = json.loads(Path('settings.json').read_text())
    body = {'prompt': 'Q: ' + question + '\nA:', 'tokens': cfg['tokens']}
    request = Request('http://127.0.0.1:' + str(cfg['port']) + '/generate',
                      data=json.dumps(body).encode(),
                      headers={'Content-Type': 'application/json'})
    with urlopen(request, timeout=120) as response:
        return json.load(response)['text']


if __name__ == '__main__':
    print('XChat -> Runner -> XSML | /exit')
    while True:
        try:
            question = input('You> ').strip()
            if question == '/exit':
                break
            if question:
                print('XSML>', ask(question))
        except (EOFError, KeyboardInterrupt):
            break
        except (URLError, TimeoutError) as error:
            print('Runner:', error)