さあ、あなただけの人工知能を作ってみましょう。

widthはベクトル幅、contextは文脈長、learning_rateは更新量です。

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

一行に質問と回答を一組。短い自作の英語例です。

Q: hello | A: hello friend.
Q: hi | A: hi friend.
Q: name | A: i am xslm.
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.

UTF-8読込、重複除去、文単位分割、学習・検証の保存を行います。

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で合併し、trainで順序を学び、encodeで再適用、decodeで復元します。

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は値、gradは勾配、parentsは局所微分です。softmaxは確率化します。

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

shapesで寸法を確認。すべての重みが学習対象で、stateとloadは数値を扱います。

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)

資料ハッシュで変更を検出。仮ファイルとos.replaceで完全な保存を保ちます。SGDにモメンタムはありません。

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()

構造と重みを復元し、文脈を切って反復抽選します。

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))

モデルを共有し、入力を検証してUTF-8 JSONを返します。

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': 'XSLM'})
            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': 'XSLM'})
            except (ValueError, KeyError, TypeError):
                self.reply(400, {'error': 'invalid request'})

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


if __name__ == '__main__':
    main()

inputで読んでHTTP POSTを送り、実際の回答や接続エラーを表示します。

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 -> XSLM | /exit')
    while True:
        try:
            question = input('You> ').strip()
            if question == '/exit':
                break
            if question:
                print('XSLM>', ask(question))
        except (EOFError, KeyboardInterrupt):
            break
        except (URLError, TimeoutError) as error:
            print('Runner:', error)