자, 이제 여러분들만의 인공지능을 만들어 봅시다.

작게 시작하는 수업 기본값입니다. 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는 조각을 이어 붙입니다. 마지막 assert는 원문 복원을 검사합니다.

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에서 크기를 본 뒤 forward를 읽으세요. weights의 Value는 모두 학습 대상입니다. 리스트로 벡터와 행렬을 표현합니다. loss는 각 위치의 다음 토큰과 비교하며 state와 load는 Python 객체 대신 숫자를 저장하고 읽습니다.

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

load는 모델 구조를 다시 만들고 숫자를 복원합니다. generate는 문맥을 자르고 후보 점수를 계산해 반복해서 뽑습니다. 같은 seed로 비교 실험을 재현할 수 있습니다.

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으로 응답합니다. 가상 리눅스의 loopback에 연결하여 이 리눅스 세션 안에서만 씁니다.

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은 한 줄을 읽고, Request는 JSON POST를 만들고, urlopen은 실제 응답을 읽습니다. 연결이 실패하면 가짜 답을 넣지 않고 오류를 알려 줍니다.

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)