.. Licensed to the Apache Software Foundation (ASF) under one
   or more contributor license agreements.  See the NOTICE file
   distributed with this work for additional information
   regarding copyright ownership.  The ASF licenses this file
   to you under the Apache License, Version 2.0 (the
   "License"); you may not use this file except in compliance
   with the License.  You may obtain a copy of the License at

     http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing,
   software distributed under the License is distributed on an
   "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
   KIND, either express or implied.  See the License for the
   specific language governing permissions and limitations
   under the License.

Google Neural Machine Translation
=================================

In this notebook, we are going to train Google NMT on IWSLT 2015
English-Vietnamese Dataset. The building process includes four steps: 1)
load and process dataset, 2) create sampler and DataLoader, 3) build
model, and 4) write training epochs.

Load MXNET and Gluon
--------------------

.. code:: python

    import warnings
    warnings.filterwarnings('ignore')

    import argparse
    import time
    import random
    import os
    import io
    import logging
    import numpy as np
    import mxnet as mx
    from mxnet import gluon
    import gluonnlp as nlp
    import nmt

Hyper-parameters
----------------

.. code:: python

    np.random.seed(100)
    random.seed(100)
    mx.random.seed(10000)
    ctx = mx.gpu(0)

    # parameters for dataset
    dataset = 'IWSLT2015'
    src_lang, tgt_lang = 'en', 'vi'
    src_max_len, tgt_max_len = 50, 50

    # parameters for model
    num_hidden = 512
    num_layers = 2
    num_bi_layers = 1
    dropout = 0.2

    # parameters for training
    batch_size, test_batch_size = 128, 32
    num_buckets = 5
    epochs = 1
    clip = 5
    lr = 0.001
    lr_update_factor = 0.5
    log_interval = 10
    save_dir = 'gnmt_en_vi_u512'

    #parameters for testing
    beam_size = 10
    lp_alpha = 1.0
    lp_k = 5

    nmt.utils.logging_config(save_dir)

Load and Preprocess Dataset
---------------------------

The following shows how to process the dataset and cache the processed
dataset for future use. The processing steps include: 1) clip the source
and target sequences, 2) split the string input to a list of tokens, 3)
map the string token into its integer index in the vocabulary, and 4)
append end-of-sentence (EOS) token to source sentence and add BOS and
EOS tokens to target sentence.

.. code:: python

    def cache_dataset(dataset, prefix):
        """Cache the processed npy dataset  the dataset into a npz

        Parameters
        ----------
        dataset : gluon.data.SimpleDataset
        file_path : str
        """
        if not os.path.exists(nmt._constants.CACHE_PATH):
            os.makedirs(nmt._constants.CACHE_PATH)
        src_data = np.array([ele[0] for ele in dataset])
        tgt_data = np.array([ele[1] for ele in dataset])
        np.savez(os.path.join(nmt._constants.CACHE_PATH, prefix + '.npz'), src_data=src_data, tgt_data=tgt_data)


    def load_cached_dataset(prefix):
        cached_file_path = os.path.join(nmt._constants.CACHE_PATH, prefix + '.npz')
        if os.path.exists(cached_file_path):
            print('Load cached data from {}'.format(cached_file_path))
            dat = np.load(cached_file_path)
            return gluon.data.ArrayDataset(np.array(dat['src_data']), np.array(dat['tgt_data']))
        else:
            return None


    class TrainValDataTransform(object):
        """Transform the machine translation dataset.

        Clip source and the target sentences to the maximum length. For the source sentence, append the
        EOS. For the target sentence, append BOS and EOS.

        Parameters
        ----------
        src_vocab : Vocab
        tgt_vocab : Vocab
        src_max_len : int
        tgt_max_len : int
        """
        def __init__(self, src_vocab, tgt_vocab, src_max_len, tgt_max_len):
            self._src_vocab = src_vocab
            self._tgt_vocab = tgt_vocab
            self._src_max_len = src_max_len
            self._tgt_max_len = tgt_max_len

        def __call__(self, src, tgt):
            if self._src_max_len > 0:
                src_sentence = self._src_vocab[src.split()[:self._src_max_len]]
            else:
                src_sentence = self._src_vocab[src.split()]
            if self._tgt_max_len > 0:
                tgt_sentence = self._tgt_vocab[tgt.split()[:self._tgt_max_len]]
            else:
                tgt_sentence = self._tgt_vocab[tgt.split()]
            src_sentence.append(self._src_vocab[self._src_vocab.eos_token])
            tgt_sentence.insert(0, self._tgt_vocab[self._tgt_vocab.bos_token])
            tgt_sentence.append(self._tgt_vocab[self._tgt_vocab.eos_token])
            src_npy = np.array(src_sentence, dtype=np.int32)
            tgt_npy = np.array(tgt_sentence, dtype=np.int32)
            return src_npy, tgt_npy


    def process_dataset(dataset, src_vocab, tgt_vocab, src_max_len=-1, tgt_max_len=-1):
        start = time.time()
        dataset_processed = dataset.transform(TrainValDataTransform(src_vocab, tgt_vocab,
                                                                    src_max_len,
                                                                    tgt_max_len), lazy=False)
        end = time.time()
        print('Processing time spent: {}'.format(end - start))
        return dataset_processed


    def load_translation_data(dataset, src_lang='en', tgt_lang='vi'):
        """Load translation dataset

        Parameters
        ----------
        dataset : str
        src_lang : str, default 'en'
        tgt_lang : str, default 'vi'

        Returns
        -------
        data_train_processed : Dataset
            The preprocessed training sentence pairs
        data_val_processed : Dataset
            The preprocessed validation sentence pairs
        data_test_processed : Dataset
            The preprocessed test sentence pairs
        val_tgt_sentences : list
            The target sentences in the validation set
        test_tgt_sentences : list
            The target sentences in the test set
        src_vocab : Vocab
            Vocabulary of the source language
        tgt_vocab : Vocab
            Vocabulary of the target language
        """
        common_prefix = 'IWSLT2015_{}_{}_{}_{}'.format(src_lang, tgt_lang,
                                                       src_max_len, tgt_max_len)
        data_train = nlp.data.IWSLT2015('train', src_lang=src_lang, tgt_lang=tgt_lang)
        data_val = nlp.data.IWSLT2015('val', src_lang=src_lang, tgt_lang=tgt_lang)
        data_test = nlp.data.IWSLT2015('test', src_lang=src_lang, tgt_lang=tgt_lang)
        src_vocab, tgt_vocab = data_train.src_vocab, data_train.tgt_vocab
        data_train_processed = load_cached_dataset(common_prefix + '_train')
        if not data_train_processed:
            data_train_processed = process_dataset(data_train, src_vocab, tgt_vocab,
                                                   src_max_len, tgt_max_len)
            cache_dataset(data_train_processed, common_prefix + '_train')
        data_val_processed = load_cached_dataset(common_prefix + '_val')
        if not data_val_processed:
            data_val_processed = process_dataset(data_val, src_vocab, tgt_vocab)
            cache_dataset(data_val_processed, common_prefix + '_val')
        data_test_processed = load_cached_dataset(common_prefix + '_test')
        if not data_test_processed:
            data_test_processed = process_dataset(data_test, src_vocab, tgt_vocab)
            cache_dataset(data_test_processed, common_prefix + '_test')
        fetch_tgt_sentence = lambda src, tgt: tgt.split()
        val_tgt_sentences = list(data_val.transform(fetch_tgt_sentence))
        test_tgt_sentences = list(data_test.transform(fetch_tgt_sentence))
        return data_train_processed, data_val_processed, data_test_processed, \
               val_tgt_sentences, test_tgt_sentences, src_vocab, tgt_vocab


    def get_data_lengths(dataset):
        return list(dataset.transform(lambda srg, tgt: (len(srg), len(tgt))))


    data_train, data_val, data_test, val_tgt_sentences, test_tgt_sentences, src_vocab, tgt_vocab\
        = load_translation_data(dataset=dataset, src_lang=src_lang, tgt_lang=tgt_lang)
    data_train_lengths = get_data_lengths(data_train)
    data_val_lengths = get_data_lengths(data_val)
    data_test_lengths = get_data_lengths(data_test)

    with io.open(os.path.join(save_dir, 'val_gt.txt'), 'w', encoding='utf-8') as of:
        for ele in val_tgt_sentences:
            of.write(' '.join(ele) + '\n')

    with io.open(os.path.join(save_dir, 'test_gt.txt'), 'w', encoding='utf-8') as of:
        for ele in test_tgt_sentences:
            of.write(' '.join(ele) + '\n')


    data_train = data_train.transform(lambda src, tgt: (src, tgt, len(src), len(tgt)), lazy=False)
    data_val = gluon.data.SimpleDataset([(ele[0], ele[1], len(ele[0]), len(ele[1]), i)
                                         for i, ele in enumerate(data_val)])
    data_test = gluon.data.SimpleDataset([(ele[0], ele[1], len(ele[0]), len(ele[1]), i)
                                          for i, ele in enumerate(data_test)])

Create Sampler and DataLoader
-----------------------------

Now, we have obtained ``data_train``, ``data_val``, and ``data_test``.
The next step is to construct sampler and DataLoader. The first step is
to construct batchify function, which pads and stacks sequences to form
mini-batch.

.. code:: python

    train_batchify_fn = nlp.data.batchify.Tuple(nlp.data.batchify.Pad(),
                                                nlp.data.batchify.Pad(),
                                                nlp.data.batchify.Stack(dtype='float32'),
                                                nlp.data.batchify.Stack(dtype='float32'))
    test_batchify_fn = nlp.data.batchify.Tuple(nlp.data.batchify.Pad(),
                                               nlp.data.batchify.Pad(),
                                               nlp.data.batchify.Stack(dtype='float32'),
                                               nlp.data.batchify.Stack(dtype='float32'),
                                               nlp.data.batchify.Stack())

We can then construct bucketing samplers, which generate batches by
grouping sequences with similar lengths. Here, the bucketing scheme is
empirically determined.

.. code:: python

    bucket_scheme = nlp.data.ExpWidthBucket(bucket_len_step=1.2)
    train_batch_sampler = nlp.data.FixedBucketSampler(lengths=data_train_lengths,
                                                      batch_size=batch_size,
                                                      num_buckets=num_buckets,
                                                      shuffle=True,
                                                      bucket_scheme=bucket_scheme)
    logging.info('Train Batch Sampler:\n{}'.format(train_batch_sampler.stats()))
    val_batch_sampler = nlp.data.FixedBucketSampler(lengths=data_val_lengths,
                                                    batch_size=test_batch_size,
                                                    num_buckets=num_buckets,
                                                    shuffle=False)
    logging.info('Valid Batch Sampler:\n{}'.format(val_batch_sampler.stats()))
    test_batch_sampler = nlp.data.FixedBucketSampler(lengths=data_test_lengths,
                                                     batch_size=test_batch_size,
                                                     num_buckets=num_buckets,
                                                     shuffle=False)
    logging.info('Test Batch Sampler:\n{}'.format(test_batch_sampler.stats()))

Given the samplers, we can create DataLoader, which is iterable.

.. code:: python

    train_data_loader = gluon.data.DataLoader(data_train,
                                              batch_sampler=train_batch_sampler,
                                              batchify_fn=train_batchify_fn,
                                              num_workers=4)
    val_data_loader = gluon.data.DataLoader(data_val,
                                            batch_sampler=val_batch_sampler,
                                            batchify_fn=test_batchify_fn,
                                            num_workers=4)
    test_data_loader = gluon.data.DataLoader(data_test,
                                             batch_sampler=test_batch_sampler,
                                             batchify_fn=test_batchify_fn,
                                             num_workers=4)

Build GNMT Model
----------------

After obtaining DataLoader, we can build the model. The GNMT encoder and
decoder can be easily constructed by calling
``get_gnmt_encoder_decoder`` function. Then, we feed the encoder and
decoder to ``NMTModel`` to construct the GNMT model. ``model.hybridize``
allows computation to be done using the symbolic backend.

.. code:: python

    encoder, decoder = nmt.gnmt.get_gnmt_encoder_decoder(hidden_size=num_hidden,
                                                         dropout=dropout,
                                                         num_layers=num_layers,
                                                         num_bi_layers=num_bi_layers)
    model = nmt.translation.NMTModel(src_vocab=src_vocab, tgt_vocab=tgt_vocab, encoder=encoder, decoder=decoder,
                                     embed_size=num_hidden, prefix='gnmt_')
    model.initialize(init=mx.init.Uniform(0.1), ctx=ctx)
    static_alloc = True
    model.hybridize(static_alloc=static_alloc)
    logging.info(model)

    # Due to the paddings, we need to mask out the losses corresponding to padding tokens.
    loss_function = nmt.loss.SoftmaxCEMaskedLoss()
    loss_function.hybridize(static_alloc=static_alloc)

We also build the beam search translator.

.. code:: python

    translator = nmt.translation.BeamSearchTranslator(model=model, beam_size=beam_size,
                                                      scorer=nlp.model.BeamSearchScorer(alpha=lp_alpha,
                                                                                        K=lp_k),
                                                      max_length=tgt_max_len + 100)
    logging.info('Use beam_size={}, alpha={}, K={}'.format(beam_size, lp_alpha, lp_k))

We define evaluation function as follows. The ``evaluate`` function use
beam search translator to generate outputs for the validation and
testing datasets.

.. code:: python

    def evaluate(data_loader):
        """Evaluate given the data loader

        Parameters
        ----------
        data_loader : gluon.data.DataLoader

        Returns
        -------
        avg_loss : float
            Average loss
        real_translation_out : list of list of str
            The translation output
        """
        translation_out = []
        all_inst_ids = []
        avg_loss_denom = 0
        avg_loss = 0.0
        for _, (src_seq, tgt_seq, src_valid_length, tgt_valid_length, inst_ids) \
                in enumerate(data_loader):
            src_seq = src_seq.as_in_context(ctx)
            tgt_seq = tgt_seq.as_in_context(ctx)
            src_valid_length = src_valid_length.as_in_context(ctx)
            tgt_valid_length = tgt_valid_length.as_in_context(ctx)
            # Calculating Loss
            out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1)
            loss = loss_function(out, tgt_seq[:, 1:], tgt_valid_length - 1).mean().asscalar()
            all_inst_ids.extend(inst_ids.asnumpy().astype(np.int32).tolist())
            avg_loss += loss * (tgt_seq.shape[1] - 1)
            avg_loss_denom += (tgt_seq.shape[1] - 1)
            # Translate
            samples, _, sample_valid_length =\
                translator.translate(src_seq=src_seq, src_valid_length=src_valid_length)
            max_score_sample = samples[:, 0, :].asnumpy()
            sample_valid_length = sample_valid_length[:, 0].asnumpy()
            for i in range(max_score_sample.shape[0]):
                translation_out.append(
                    [tgt_vocab.idx_to_token[ele] for ele in
                     max_score_sample[i][1:(sample_valid_length[i] - 1)]])
        avg_loss = avg_loss / avg_loss_denom
        real_translation_out = [None for _ in range(len(all_inst_ids))]
        for ind, sentence in zip(all_inst_ids, translation_out):
            real_translation_out[ind] = sentence
        return avg_loss, real_translation_out


    def write_sentences(sentences, file_path):
        with io.open(file_path, 'w', encoding='utf-8') as of:
            for sent in sentences:
                of.write(' '.join(sent) + '\n')

Training Epochs
---------------

Before entering the training stage, we need to create trainer for
updating the parameters. In the following example, we create a trainer
that uses ADAM optimzier.

.. code:: python

    trainer = gluon.Trainer(model.collect_params(), 'adam', {'learning_rate': lr})

We can then write the training loop. During the training, we evaluate on
the validation and testing datasets every epoch, and record the
parameters that give the hightest BLEU score on the validation dataset.
Before performing forward and backward, we first use ``as_in_context``
function to copy the mini-batch to GPU. The statement
``with mx.autograd.record()`` tells Gluon backend to compute the
gradients for the part inside the block.

.. code:: python

    best_valid_bleu = 0.0
    for epoch_id in range(epochs):
        log_avg_loss = 0
        log_avg_gnorm = 0
        log_wc = 0
        log_start_time = time.time()
        for batch_id, (src_seq, tgt_seq, src_valid_length, tgt_valid_length)\
                in enumerate(train_data_loader):
            # logging.info(src_seq.context) Context suddenly becomes GPU.
            src_seq = src_seq.as_in_context(ctx)
            tgt_seq = tgt_seq.as_in_context(ctx)
            src_valid_length = src_valid_length.as_in_context(ctx)
            tgt_valid_length = tgt_valid_length.as_in_context(ctx)
            with mx.autograd.record():
                out, _ = model(src_seq, tgt_seq[:, :-1], src_valid_length, tgt_valid_length - 1)
                loss = loss_function(out, tgt_seq[:, 1:], tgt_valid_length - 1).mean()
                loss = loss * (tgt_seq.shape[1] - 1) / (tgt_valid_length - 1).mean()
                loss.backward()
            grads = [p.grad(ctx) for p in model.collect_params().values()]
            gnorm = gluon.utils.clip_global_norm(grads, clip)
            trainer.step(1)
            src_wc = src_valid_length.sum().asscalar()
            tgt_wc = (tgt_valid_length - 1).sum().asscalar()
            step_loss = loss.asscalar()
            log_avg_loss += step_loss
            log_avg_gnorm += gnorm
            log_wc += src_wc + tgt_wc
            if (batch_id + 1) % log_interval == 0:
                wps = log_wc / (time.time() - log_start_time)
                logging.info('[Epoch {} Batch {}/{}] loss={:.4f}, ppl={:.4f}, gnorm={:.4f}, '
                             'throughput={:.2f}K wps, wc={:.2f}K'
                             .format(epoch_id, batch_id + 1, len(train_data_loader),
                                     log_avg_loss / log_interval,
                                     np.exp(log_avg_loss / log_interval),
                                     log_avg_gnorm / log_interval,
                                     wps / 1000, log_wc / 1000))
                log_start_time = time.time()
                log_avg_loss = 0
                log_avg_gnorm = 0
                log_wc = 0
        valid_loss, valid_translation_out = evaluate(val_data_loader)
        valid_bleu_score, _, _, _, _ = nmt.bleu.compute_bleu([val_tgt_sentences], valid_translation_out)
        logging.info('[Epoch {}] valid Loss={:.4f}, valid ppl={:.4f}, valid bleu={:.2f}'
                     .format(epoch_id, valid_loss, np.exp(valid_loss), valid_bleu_score * 100))
        test_loss, test_translation_out = evaluate(test_data_loader)
        test_bleu_score, _, _, _, _ = nmt.bleu.compute_bleu([test_tgt_sentences], test_translation_out)
        logging.info('[Epoch {}] test Loss={:.4f}, test ppl={:.4f}, test bleu={:.2f}'
                     .format(epoch_id, test_loss, np.exp(test_loss), test_bleu_score * 100))
        write_sentences(valid_translation_out,
                        os.path.join(save_dir, 'epoch{:d}_valid_out.txt').format(epoch_id))
        write_sentences(test_translation_out,
                        os.path.join(save_dir, 'epoch{:d}_test_out.txt').format(epoch_id))
        if valid_bleu_score > best_valid_bleu:
            best_valid_bleu = valid_bleu_score
            save_path = os.path.join(save_dir, 'valid_best.params')
            logging.info('Save best parameters to {}'.format(save_path))
            model.save_parameters(save_path)
        if epoch_id + 1 >= (epochs * 2) // 3:
            new_lr = trainer.learning_rate * lr_update_factor
            logging.info('Learning rate change to {}'.format(new_lr))
            trainer.set_learning_rate(new_lr)

Summary
-------

In this notebook, we have shown how to train a GNMT model on IWSLT 2015
English-Vietnamese using Gluon NLP toolkit. The complete training script
can be found
`here <https://github.com/dmlc/gluon-nlp/blob/master/scripts/nmt/train_gnmt.py>`__.
The command to reproduce the result can be seen in the `nmt scripts
page <http://gluon-nlp.mxnet.io/scripts/index.html#machine-translation>`__.