# Copyright (c) 2017-present, Facebook, Inc.
# All rights reserved.
#
# This source code is licensed under the license found in the LICENSE file in
# the root directory of this source tree. An additional grant of patent rights
# can be found in the PATENTS file in the same directory.
#

import collections
import os
import torch
import math

from fairseq import bleu, data, options, utils
from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter
from fairseq.multiprocessing_trainer import MultiprocessingTrainer
from fairseq.progress_bar import progress_bar
from fairseq.sequence_generator import SequenceGenerator


def main():
    parser = options.get_parser('Trainer')
    dataset_args = options.add_dataset_args(parser)
    dataset_args.add_argument('--max-tokens', default=0, type=int, metavar='N',
                              help='maximum number of tokens in a batch')
    dataset_args.add_argument('--batch-size', default=32, type=int, metavar='N',
                              help='batch size')
    dataset_args.add_argument('--test-batch-size', default=32, type=int, metavar='N',
                              help='batch size for test set')
    dataset_args.add_argument('--valid-batch-size', default=32, type=int, metavar='N',
                              help='batch size for validation set')
    dataset_args.add_argument('--train-subset', default='train', metavar='SPLIT',
                              choices=['train', 'valid', 'test'],
                              help='data subset to use for training (train, valid, test)')
    dataset_args.add_argument('--valid-subset', default='valid', metavar='SPLIT',
                              help='comma separated list ofdata subsets '
                                   ' to use for validation (train, valid, valid1,test, test1)')
    dataset_args.add_argument('--test-subset', default='test', metavar='SPLIT',
                              help='comma separated list ofdata subset '
                                   'to use for testing (train, valid, test)')
    dataset_args.add_argument('--valid-script', nargs='+', metavar='PATH', help='path to external validation script (optional).')

    options.add_optimization_args(parser)
    options.add_checkpoint_args(parser)
    options.add_model_args(parser)

    args = utils.parse_args_and_arch(parser)
    print(args)

    if args.no_progress_bar:
        progress_bar.enabled = False
        progress_bar.print_interval = args.log_interval

    if not os.path.exists(args.save_dir):
        os.makedirs(args.save_dir)
    torch.manual_seed(args.seed)

    # Setting args.max_tokens to infinity(same as setting to None)
    if args.max_tokens == 0:
        args.max_tokens = None

    # Load dataset
    dataset = data.load_with_check(args.data, args.source_lang, args.target_lang)
    if args.source_lang is None or args.target_lang is None:
        # record inferred languages in args, so that it's saved in checkpoints
        args.source_lang, args.target_lang = dataset.src, dataset.dst

    print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
    print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
    for split in dataset.splits:
        print('| {} {} {} examples'.format(args.data, split, len(dataset.splits[split])))

    if not torch.cuda.is_available():
        raise NotImplementedError('Training on CPU is not supported')
    num_gpus = torch.cuda.device_count()

    print('| using {} GPUs (with max tokens per GPU = {})'.format(num_gpus, args.max_tokens))

    # Build model
    print('| model {}'.format(args.arch))
    model = utils.build_model(args, dataset)
    criterion = utils.build_criterion(args, dataset)

    # Start multiprocessing
    trainer = MultiprocessingTrainer(args, model)

    # Load the latest checkpoint if one is available
    epoch, batch_offset = trainer.load_checkpoint(os.path.join(args.save_dir, args.restore_file))

    # Train until the learning rate gets too small
    val_loss = None
    max_epoch = args.max_epoch or math.inf
    lr = trainer.get_lr()
    train_meter = StopwatchMeter()
    train_meter.start()
    while lr > args.min_lr and epoch <= max_epoch:
        # train for one epoch
        train(args, epoch, batch_offset, trainer, criterion, dataset, num_gpus)

        # evaluate on validate set
        for k, subset in enumerate(args.valid_subset.split(',')):
            val_loss = validate(args, epoch, trainer, criterion, dataset, subset, num_gpus)
            if k == 0:
                if not args.no_save:
                    # save checkpoint
                    trainer.save_checkpoint(args, epoch, 0, val_loss, validation_script=args.valid_script)

                # only use first validation loss to update the learning schedule
                lr = trainer.lr_step(val_loss, epoch)


        epoch += 1
        batch_offset = 0
    train_meter.stop()
    print('| done training in {:.1f} seconds'.format(train_meter.sum))

    # Generate on test set and compute BLEU score
    for beam in [1, 5, 10, 20]:
        for subset in args.test_subset.split(','):
            scorer = score_test(args, trainer.get_model(), dataset, subset, beam,
                                cuda_device=(0 if num_gpus > 0 else None))
            print('| Test on {} with beam={}: {}'.format(subset, beam, scorer.result_string()))

    # Stop multiprocessing
    trainer.stop()


def train(args, epoch, batch_offset, trainer, criterion, dataset, num_gpus):
    """Train the model for one epoch."""

    itr = dataset.dataloader(args.train_subset, batch_size=args.batch_size,
                             test_batch_size=args.test_batch_size,
                             valid_batch_size=args.valid_batch_size,
                             num_workers=args.workers,
                             max_tokens=args.max_tokens, seed=args.seed, epoch=epoch,
                             max_positions=args.max_positions,
                             sample_without_replacement=args.sample_without_replacement)
    loss_meter = AverageMeter()
    bsz_meter = AverageMeter()    # sentences per batch
    wpb_meter = AverageMeter()    # words per batch
    wps_meter = TimeMeter()       # words per second
    clip_meter = AverageMeter()   # % of updates clipped
    gnorm_meter = AverageMeter()  # gradient norm

    desc = '| epoch {:03d}'.format(epoch)
    lr = trainer.get_lr()
    with progress_bar(itr, desc, leave=False) as t:
        for i, sample in data.skip_group_enumerator(t, num_gpus, batch_offset):
            loss, grad_norm = trainer.train_step(sample, criterion)

            ntokens = sum(s['ntokens'] for s in sample)
            src_size = sum(s['src_tokens'].size(0) for s in sample)
            loss_meter.update(loss, ntokens)
            bsz_meter.update(src_size)
            wpb_meter.update(ntokens)
            wps_meter.update(ntokens)
            clip_meter.update(1 if grad_norm > args.clip_norm else 0)
            gnorm_meter.update(grad_norm)

            t.set_postfix(collections.OrderedDict([
                ('loss', '{:.2f} ({:.2f})'.format(loss, loss_meter.avg)),
                ('wps', '{:5d}'.format(round(wps_meter.avg))),
                ('wpb', '{:5d}'.format(round(wpb_meter.avg))),
                ('bsz', '{:5d}'.format(round(bsz_meter.avg))),
                ('lr', lr),
                ('clip', '{:3.0f}%'.format(clip_meter.avg * 100)),
                ('gnorm', '{:.4f}'.format(gnorm_meter.avg)),
            ]))

            if i == 0:
                # ignore the first mini-batch in words-per-second calculation
                wps_meter.reset()
            if args.save_interval > 0 and (i + 1) % args.save_interval == 0:
                trainer.save_checkpoint(args, epoch, i + 1)

        fmt = desc + ' | train loss {:2.2f} | train ppl {:3.2f}'
        fmt += ' | s/checkpoint {:7d} | words/s {:6d} | words/batch {:6d}'
        fmt += ' | bsz {:5d} | lr {:0.6f} | clip {:3.0f}% | gnorm {:.4f}'
        t.write(fmt.format(loss_meter.avg, math.pow(2, loss_meter.avg),
                           round(wps_meter.elapsed_time),
                           round(wps_meter.avg),
                           round(wpb_meter.avg),
                           round(bsz_meter.avg),
                           lr, clip_meter.avg * 100,
                           gnorm_meter.avg))


def validate(args, epoch, trainer, criterion, dataset, subset, ngpus):
    """Evaluate the model on the validation set and return the average loss."""

    itr = dataset.dataloader(subset, batch_size=None,
                             max_tokens=args.max_tokens,
                             max_positions=args.max_positions)
    loss_meter = AverageMeter()

    desc = '| epoch {:03d} | valid on \'{}\' subset'.format(epoch, subset)
    with progress_bar(itr, desc, leave=False) as t:
        for _, sample in data.skip_group_enumerator(t, ngpus):
            ntokens = sum(s['ntokens'] for s in sample)
            loss = trainer.valid_step(sample, criterion)
            loss_meter.update(loss, ntokens)
            t.set_postfix(loss='{:.2f}'.format(loss_meter.avg))

        val_loss = loss_meter.avg
        t.write(desc + ' | valid loss {:2.2f} | valid ppl {:3.2f}'
                .format(val_loss, math.pow(2, val_loss)))

    # update and return the learning rate
    return val_loss


def score_test(args, model, dataset, subset, beam, cuda_device):
    """Evaluate the model on the test set and return the BLEU scorer."""

    translator = SequenceGenerator([model], dataset.dst_dict, beam_size=beam)
    if torch.cuda.is_available():
        translator.cuda()

    scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())
    itr = dataset.dataloader(subset, batch_size=4, max_positions=args.max_positions)
    for _, _, ref, hypos in translator.generate_batched_itr(itr, cuda_device=cuda_device):
        scorer.add(ref.int().cpu(), hypos[0]['tokens'].int().cpu())
    return scorer


if __name__ == '__main__':
    main()

[((22, 13, 22, 42), 'fairseq.options.get_parser', 'options.get_parser', ({(22, 32, 22, 41): '"""Trainer"""'}, {}), "('Trainer')", False, 'from fairseq import bleu, data, options, utils\n'), ((23, 19, 23, 51), 'fairseq.options.add_dataset_args', 'options.add_dataset_args', ({(23, 44, 23, 50): 'parser'}, {}), '(parser)', False, 'from fairseq import bleu, data, options, utils\n'), ((43, 4, 43, 41), 'fairseq.options.add_optimization_args', 'options.add_optimization_args', ({(43, 34, 43, 40): 'parser'}, {}), '(parser)', False, 'from fairseq import bleu, data, options, utils\n'), ((44, 4, 44, 39), 'fairseq.options.add_checkpoint_args', 'options.add_checkpoint_args', ({(44, 32, 44, 38): 'parser'}, {}), '(parser)', False, 'from fairseq import bleu, data, options, utils\n'), ((45, 4, 45, 34), 'fairseq.options.add_model_args', 'options.add_model_args', ({(45, 27, 45, 33): 'parser'}, {}), '(parser)', False, 'from fairseq import bleu, data, options, utils\n'), ((47, 11, 47, 44), 'fairseq.utils.parse_args_and_arch', 'utils.parse_args_and_arch', ({(47, 37, 47, 43): 'parser'}, {}), '(parser)', False, 'from fairseq import bleu, data, options, utils\n'), ((56, 4, 56, 32), 'torch.manual_seed', 'torch.manual_seed', ({(56, 22, 56, 31): 'args.seed'}, {}), '(args.seed)', False, 'import torch\n'), ((63, 14, 63, 81), 'fairseq.data.load_with_check', 'data.load_with_check', ({(63, 35, 63, 44): 'args.data', (63, 46, 63, 62): 'args.source_lang', (63, 64, 63, 80): 'args.target_lang'}, {}), '(args.data, args.source_lang, args.target_lang)', False, 'from fairseq import bleu, data, options, utils\n'), ((75, 15, 75, 40), 'torch.cuda.device_count', 'torch.cuda.device_count', ({}, {}), '()', False, 'import torch\n'), ((81, 12, 81, 44), 'fairseq.utils.build_model', 'utils.build_model', ({(81, 30, 81, 34): 'args', (81, 36, 81, 43): 'dataset'}, {}), '(args, dataset)', False, 'from fairseq import bleu, data, options, utils\n'), ((82, 16, 82, 52), 'fairseq.utils.build_criterion', 'utils.build_criterion', ({(82, 38, 82, 42): 'args', (82, 44, 82, 51): 'dataset'}, {}), '(args, dataset)', False, 'from fairseq import bleu, data, options, utils\n'), ((85, 14, 85, 49), 'fairseq.multiprocessing_trainer.MultiprocessingTrainer', 'MultiprocessingTrainer', ({(85, 37, 85, 41): 'args', (85, 43, 85, 48): 'model'}, {}), '(args, model)', False, 'from fairseq.multiprocessing_trainer import MultiprocessingTrainer\n'), ((94, 18, 94, 34), 'fairseq.meters.StopwatchMeter', 'StopwatchMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((138, 17, 138, 31), 'fairseq.meters.AverageMeter', 'AverageMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((139, 16, 139, 30), 'fairseq.meters.AverageMeter', 'AverageMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((140, 16, 140, 30), 'fairseq.meters.AverageMeter', 'AverageMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((141, 16, 141, 27), 'fairseq.meters.TimeMeter', 'TimeMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((142, 17, 142, 31), 'fairseq.meters.AverageMeter', 'AverageMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((143, 18, 143, 32), 'fairseq.meters.AverageMeter', 'AverageMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((194, 17, 194, 31), 'fairseq.meters.AverageMeter', 'AverageMeter', ({}, {}), '()', False, 'from fairseq.meters import AverageMeter, StopwatchMeter, TimeMeter\n'), ((215, 17, 215, 77), 'fairseq.sequence_generator.SequenceGenerator', 'SequenceGenerator', (), '', False, 'from fairseq.sequence_generator import SequenceGenerator\n'), ((216, 7, 216, 32), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((54, 11, 54, 40), 'os.path.exists', 'os.path.exists', ({(54, 26, 54, 39): 'args.save_dir'}, {}), '(args.save_dir)', False, 'import os\n'), ((55, 8, 55, 34), 'os.makedirs', 'os.makedirs', ({(55, 20, 55, 33): 'args.save_dir'}, {}), '(args.save_dir)', False, 'import os\n'), ((73, 11, 73, 36), 'torch.cuda.is_available', 'torch.cuda.is_available', ({}, {}), '()', False, 'import torch\n'), ((88, 50, 88, 96), 'os.path.join', 'os.path.join', ({(88, 63, 88, 76): 'args.save_dir', (88, 78, 88, 95): 'args.restore_file'}, {}), '(args.save_dir, args.restore_file)', False, 'import os\n'), ((147, 9, 147, 45), 'fairseq.progress_bar.progress_bar', 'progress_bar', (), '', False, 'from fairseq.progress_bar import progress_bar\n'), ((148, 25, 148, 78), 'fairseq.data.skip_group_enumerator', 'data.skip_group_enumerator', ({(148, 52, 148, 53): 't', (148, 55, 148, 63): 'num_gpus', (148, 65, 148, 77): 'batch_offset'}, {}), '(t, num_gpus, batch_offset)', False, 'from fairseq import bleu, data, options, utils\n'), ((197, 9, 197, 45), 'fairseq.progress_bar.progress_bar', 'progress_bar', (), '', False, 'from fairseq.progress_bar import progress_bar\n'), ((198, 25, 198, 61), 'fairseq.data.skip_group_enumerator', 'data.skip_group_enumerator', ({(198, 52, 198, 53): 't', (198, 55, 198, 60): 'ngpus'}, {}), '(t, ngpus)', False, 'from fairseq import bleu, data, options, utils\n'), ((179, 43, 179, 70), 'math.pow', 'math.pow', ({(179, 52, 179, 53): '(2)', (179, 55, 179, 69): 'loss_meter.avg'}, {}), '(2, loss_meter.avg)', False, 'import math\n'), ((206, 34, 206, 55), 'math.pow', 'math.pow', ({(206, 43, 206, 44): '(2)', (206, 46, 206, 54): 'val_loss'}, {}), '(2, val_loss)', False, 'import math\n')]