generate.py 6.55 KB
Newer Older
Louis Martin's avatar
Louis Martin committed
1
#!/usr/bin/env python3
Sergey Edunov's avatar
Sergey Edunov committed
2
3
4
5
6
7
8
9
10
11
# 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 torch

12
from fairseq import bleu, data, options, tokenizer, utils
Sergey Edunov's avatar
Sergey Edunov committed
13
14
15
16
17
18
19
20
21
22
23
24
25
from fairseq.meters import StopwatchMeter, TimeMeter
from fairseq.sequence_generator import SequenceGenerator


def main():
    parser = options.get_parser('Generation')
    parser.add_argument('--path', metavar='FILE', required=True, action='append',
                        help='path(s) to model file(s)')
    dataset_args = options.add_dataset_args(parser)
    dataset_args.add_argument('--batch-size', default=32, type=int, metavar='N',
                              help='batch size')
    dataset_args.add_argument('--gen-subset', default='test', metavar='SPLIT',
                              help='data subset to generate (train, valid, test)')
Myle Ott's avatar
Myle Ott committed
26
27
28
29
    dataset_args.add_argument('--num-shards', default=1, type=int, metavar='N',
                              help='shard generation over N shards')
    dataset_args.add_argument('--shard-id', default=0, type=int, metavar='ID',
                              help='id of the shard to generate (id < num_shards)')
Sergey Edunov's avatar
Sergey Edunov committed
30
31
32
    options.add_generation_args(parser)

    args = parser.parse_args()
33
    if args.no_progress_bar and args.log_format is None:
34
        args.log_format = 'none'
Sergey Edunov's avatar
Sergey Edunov committed
35
36
37
38
    print(args)

    use_cuda = torch.cuda.is_available() and not args.cpu

39
    # Load dataset
40
41
42
43
    if args.replace_unk is None:
        dataset = data.load_dataset(args.data, [args.gen_subset], args.source_lang, args.target_lang)
    else:
        dataset = data.load_raw_text_dataset(args.data, [args.gen_subset], args.source_lang, args.target_lang)
44
45
46
47
48
    if args.source_lang is None or args.target_lang is None:
        # record inferred languages in args
        args.source_lang, args.target_lang = dataset.src, dataset.dst

    # Load ensemble
Sergey Edunov's avatar
Sergey Edunov committed
49
    print('| loading model(s) from {}'.format(', '.join(args.path)))
Myle Ott's avatar
Myle Ott committed
50
    models, _ = utils.load_ensemble_for_inference(args.path, dataset.src_dict, dataset.dst_dict)
Sergey Edunov's avatar
Sergey Edunov committed
51
52
53

    print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
    print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
Louis Martin's avatar
Louis Martin committed
54
    print('| {} {} {} examples'.format(args.data, args.gen_subset, len(dataset.splits[args.gen_subset])))
Sergey Edunov's avatar
Sergey Edunov committed
55

56
    # Optimize ensemble for generation
Sergey Edunov's avatar
Sergey Edunov committed
57
    for model in models:
Myle Ott's avatar
Myle Ott committed
58
59
        model.make_generation_fast_(
            beamable_mm_beam_size=None if args.no_beamable_mm else args.beam)
Sergey Edunov's avatar
Sergey Edunov committed
60
61

    # Initialize generator
62
    translator = SequenceGenerator(
63
        models, beam_size=args.beam, stop_early=(not args.no_early_stop),
64
65
        normalize_scores=(not args.unnormalized), len_penalty=args.lenpen,
        unk_penalty=args.unkpen)
66
67
68
69
    if use_cuda:
        translator.cuda()

    # Load alignment dictionary for unknown word replacement
Louis Martin's avatar
Louis Martin committed
70
71
72
73
74
75
    # (None if no unknown word replacement, empty if no path to align dictionary)
    align_dict = utils.load_align_dict(args.replace_unk)

    # Generate and compute BLEU score
    scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())
    max_positions = min(model.max_encoder_positions() for model in models)
Myle Ott's avatar
Myle Ott committed
76
    itr = dataset.eval_dataloader(
77
        args.gen_subset, max_sentences=args.batch_size, max_positions=max_positions,
Myle Ott's avatar
Myle Ott committed
78
        skip_invalid_size_inputs_valid_test=args.skip_invalid_size_inputs_valid_test)
Myle Ott's avatar
Myle Ott committed
79
80
81
82
    if args.num_shards > 1:
        if args.shard_id < 0 or args.shard_id >= args.num_shards:
            raise ValueError('--shard-id must be between 0 and num_shards')
        itr = data.sharded_iterator(itr, args.num_shards, args.shard_id)
Louis Martin's avatar
Louis Martin committed
83
    num_sentences = 0
84
    with utils.build_progress_bar(args, itr) as t:
Louis Martin's avatar
Louis Martin committed
85
86
87
88
89
90
91
92
        wps_meter = TimeMeter()
        gen_timer = StopwatchMeter()
        translations = translator.generate_batched_itr(
            t, maxlen_a=args.max_len_a, maxlen_b=args.max_len_b,
            cuda_device=0 if use_cuda else None, timer=gen_timer)
        for sample_id, src_tokens, target_tokens, hypos in translations:
            # Process input and ground truth
            target_tokens = target_tokens.int().cpu()
93
94
95
96
97
98
99
100
            # Either retrieve the original sentences or regenerate them from tokens.
            if align_dict is not None:
                src_str = dataset.splits[args.gen_subset].src.get_original_text(sample_id)
                target_str = dataset.splits[args.gen_subset].dst.get_original_text(sample_id)
            else:
                src_str = dataset.src_dict.string(src_tokens, args.remove_bpe)
                target_str = dataset.dst_dict.string(target_tokens, args.remove_bpe, escape_unk=True)

Louis Martin's avatar
Louis Martin committed
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
            if not args.quiet:
                print('S-{}\t{}'.format(sample_id, src_str))
                print('T-{}\t{}'.format(sample_id, target_str))

            # Process top predictions
            for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
                hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
                    hypo_tokens=hypo['tokens'].int().cpu(),
                    src_str=src_str,
                    alignment=hypo['alignment'].int().cpu(),
                    align_dict=align_dict,
                    dst_dict=dataset.dst_dict,
                    remove_bpe=args.remove_bpe)

                if not args.quiet:
                    print('H-{}\t{}\t{}'.format(sample_id, hypo['score'], hypo_str))
117
                    print('A-{}\t{}'.format(sample_id, ' '.join(map(str, alignment))))
Louis Martin's avatar
Louis Martin committed
118
119
120

                # Score only the top hypothesis
                if i == 0:
121
122
                    if align_dict is not None or args.remove_bpe is not None:
                        # Convert back to tokens for evaluation with unk replacement and/or without BPE
Louis Martin's avatar
Louis Martin committed
123
124
125
126
127
128
                        target_tokens = tokenizer.Tokenizer.tokenize(target_str,
                                                                     dataset.dst_dict,
                                                                     add_if_not_exist=True)
                    scorer.add(target_tokens, hypo_tokens)

            wps_meter.update(src_tokens.size(0))
129
            t.log({'wps': round(wps_meter.avg)})
Louis Martin's avatar
Louis Martin committed
130
131
132
133
134
            num_sentences += 1

    print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} tokens/s)'.format(
        num_sentences, gen_timer.n, gen_timer.sum, 1. / gen_timer.avg))
    print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))
Sergey Edunov's avatar
Sergey Edunov committed
135
136
137
138


if __name__ == '__main__':
    main()