run_gpt.py

import argparse

import os
from itertools import islice
from typing import Iterable, List, Optional

from keras import optimizers, losses
from keras.models import load_model
# noinspection PyPep8Naming
from keras import backend as K
from keras import callbacks
import numpy as np

from . import wikitext
from .bpe import BPEEncoder, ID_FOR_PADDING
from .utils import (
    load_optimizer_weights, contain_tf_gpu_mem_usage, CosineLRSchedule)
from .models import (
    universal_transformer_gpt_model, vanilla_transformer_gpt_model)


def pad_lm_samples(samples: Iterable[List[int]],
                   required_sequence_length: int):
    tail_padding = [ID_FOR_PADDING]
    for sample in samples:
        assert len(sample) > 0
        sample.extend(tail_padding * (required_sequence_length - len(sample)))


def training_data_to_samples(training_set_name: str,
                             encoder: BPEEncoder,
                             max_sequence_length: int) -> np.ndarray:
    """
    Reads WikiText dataset, interpreting each line as an independent sequence,
    then splits those lines with BPE tokenizer and turns them into word ids
    based on previously constructed BPE vocabulary (both the tokenizer
    and the vocabulary are parts of the BPEEncoder instance).

    Those word id's then packed into a matrix the size of
    (number of lines x max_sequence_length + 1), which can be later sliced
    to get X and Y matrices of sequences for training).
    """
    training_set = wikitext.read_wikitext_file(training_set_name)
    useful_sequences = []
    for line in training_set.splitlines():
        clean_line = line.strip()
        is_header = clean_line.startswith('=') and clean_line.endswith('=')
        if is_header or not clean_line:
            continue
        # the encoder is supposed to add <SEQ> and </SEQ>
        id_word_pairs = list(encoder(clean_line))
        useful_sequences.append(
            [word_id for word_id, _ in id_word_pairs[:max_sequence_length]])

    pad_lm_samples(useful_sequences, max_sequence_length + 1)
    result = np.empty(
        (len(useful_sequences), max_sequence_length + 1),
        dtype='int32')
    for i, sequence in enumerate(useful_sequences):
        result[i, :] = sequence
    return result


def training_data_to_dense_samples(training_set_name: str,
                                   encoder: BPEEncoder,
                                   max_sequence_length: int) -> np.ndarray:
    """
    Reads WikiText dataset, interpreting each line as an independent sequence,
    then splits those lines with BPE tokenizer and turns them into word ids
    based on previously constructed BPE vocabulary (both the tokenizer
    and the vocabulary are parts of the BPEEncoder instance).

    Those word id's then packed into a matrix the size of
    (number of lines x max_sequence_length + 1), which can be later sliced
    to get X and Y matrices of sequences for training).
    """
    training_set = wikitext.read_wikitext_file(training_set_name)
    useful_sequences = []

    def stream_bpe_tokens():
        for line in training_set.splitlines():
            clean_line = line.strip()
            if not clean_line:
                continue
            # the encoder is supposed to add <SEQ> and </SEQ>
            id_word_pairs = encoder(clean_line)
            yield from id_word_pairs

    id_word_stream = stream_bpe_tokens()
    while True:
        chunk = list(islice(id_word_stream, max_sequence_length))
        if len(chunk) == 0:
            break
        sample_sequence = [word_id for word_id, _ in chunk]
        useful_sequences.append(sample_sequence)

    pad_lm_samples(useful_sequences, max_sequence_length + 1)
    result = np.empty(
        (len(useful_sequences), max_sequence_length + 1),
        dtype='int32')
    for i, sequence in enumerate(useful_sequences):
        result[i, :] = sequence
    return result


def perplexity(y_true, y_pred):
    """
    Popular metric for evaluating language modelling architectures.
    More info: http://cs224d.stanford.edu/lecture_notes/LectureNotes4.pdf
    """
    cross_entropy = K.sparse_categorical_crossentropy(y_true, y_pred)
    return K.mean(K.exp(K.mean(cross_entropy, axis=-1)))


def main(model_save_path: str,
         model_name: str,
         tensorboard_log_path: Optional[str],
         num_epochs: int,
         learning_rate: float,
         batch_size: int,
         max_seq_length: int,
         word_embedding_size: int,
         load_weights_only: bool,
         show_model_summary: bool):
    contain_tf_gpu_mem_usage()
    encoder = wikitext.build_wikitext_bpe_encoder()

    def x_y_for_dataset(dataset_name):
        fat_sample = training_data_to_dense_samples(
            dataset_name, encoder, max_seq_length)
        _x = fat_sample[:, :max_seq_length]
        _y = np.expand_dims(fat_sample[:, 1:], axis=-1)
        return _x, _y

    x, y = x_y_for_dataset(wikitext.TRAINING_SET_NAME)

    def compile_new_model():
        if model_name == 'universal':
            optimizer = optimizers.Adam(
                lr=learning_rate, beta_1=0.6, beta_2=0.999)
            _model = universal_transformer_gpt_model(
                max_seq_length=max_seq_length,
                vocabulary_size=encoder.vocabulary_size(),
                word_embedding_size=word_embedding_size,
                transformer_depth=5,
                num_heads=8)
            _model.compile(
                optimizer,
                loss=losses.sparse_categorical_crossentropy,
                metrics=[perplexity])
        elif model_name == 'vanilla':
            optimizer = optimizers.Adam(
                lr=learning_rate, beta_1=0.9, beta_2=0.999, clipvalue=5.0)
            _model = vanilla_transformer_gpt_model(
                max_seq_length=max_seq_length,
                vocabulary_size=encoder.vocabulary_size(),
                word_embedding_size=word_embedding_size,
                transformer_depth=5,
                num_heads=8)
            _model.compile(
                optimizer,
                loss=losses.sparse_categorical_crossentropy,
                metrics=[perplexity])
        else:
            raise RuntimeError(f'Unknown model {model_name}')
        return _model

    if os.path.exists(model_save_path):
        if load_weights_only:
            print('Loading weights from', model_save_path)
            model = compile_new_model()
            model.load_weights(
                model_save_path, skip_mismatch=True, by_name=True)
            load_optimizer_weights(model, model_save_path)
        else:
            print('Loading the whole model from', model_save_path)
            model = load_model(
                model_save_path,
                custom_objects={
                    'perplexity': perplexity,
                })
    else:
        model = compile_new_model()

    if show_model_summary:
        model.summary(120)

    lr_scheduler = callbacks.LearningRateScheduler(
        CosineLRSchedule(lr_high=learning_rate,
                         lr_low=learning_rate / 32,
                         initial_period=num_epochs),
        verbose=1)
    model_callbacks = [
        callbacks.ModelCheckpoint(
            model_save_path,
            monitor='val_loss', save_best_only=True, verbose=True),
        lr_scheduler,
    ]
    if tensorboard_log_path:
        model_callbacks.append(callbacks.TensorBoard(tensorboard_log_path))
    model.fit(
        x, y,
        validation_data=x_y_for_dataset(wikitext.VALIDATION_SET_NAME),
        batch_size=batch_size, epochs=num_epochs,
        callbacks=model_callbacks)
    # Evaluation using test set
    print('-' * 80)
    test_x, test_y = x_y_for_dataset(wikitext.TEST_SET_NAME)
    test_metrics = model.evaluate(test_x, test_y, batch_size=batch_size)
    for metric_name, metric_value in zip(model.metrics_names, test_metrics):
        print(f'Test {metric_name}:', metric_value)


if __name__ == '__main__':
    _argparser = argparse.ArgumentParser(
        description='A simple example of the Transformer model in work',
        formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    _argparser.add_argument(
        '--save', type=str, required=True, metavar='PATH',
        help='A path where the best model should be saved / restored from')
    _argparser.add_argument(
        '--tensorboard-log', type=str, metavar='PATH', default=None,
        help='Path to a directory for Tensorboard logs')
    _argparser.add_argument(
        '--epochs', type=int, default=200, metavar='INTEGER',
        help='The number of epochs to train')
    _argparser.add_argument(
        '--lr', type=float, default=2e-4, metavar='FLOAT',
        help='Learning rate')
    _argparser.add_argument(
        '--batch-size', type=int, default=32, metavar='INTEGER',
        help='Training batch size')
    _argparser.add_argument(
        '--seq-len', type=int, default=256, metavar='INTEGER',
        help='Max sequence length')
    _argparser.add_argument(
        '--we-size', type=int, default=512, metavar='INTEGER',
        help='Word embedding size')
    _argparser.add_argument(
        '--model', type=str, default='universal', metavar='NAME',
        choices=['universal', 'vanilla'],
        help='The type of the model to train: "vanilla" or "universal"')
    _argparser.add_argument(
        '--load-weights-only', action='store_true',
        help='Use the save file only to initialize weights '
             '(do not load the whole model)')
    _argparser.add_argument(
        '--model-summary', action='store_true',
        help='Display the summary of the model before the training begins')
    _args = _argparser.parse_args()

    main(model_save_path=_args.save,
         model_name=_args.model,
         tensorboard_log_path=_args.tensorboard_log,
         num_epochs=_args.epochs,
         learning_rate=_args.lr,
         batch_size=_args.batch_size,
         max_seq_length=_args.seq_len,
         word_embedding_size=_args.we_size,
         load_weights_only=_args.load_weights_only,
         show_model_summary=_args.model_summary)