generate.py 6.63 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
    print(args)

    use_cuda = torch.cuda.is_available() and not args.cpu
38
39
    if hasattr(torch, 'set_grad_enabled'):
        torch.set_grad_enabled(False)
Sergey Edunov's avatar
Sergey Edunov committed
40

41
    # Load dataset
42
43
44
45
    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)
46
47
48
49
50
    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
51
    print('| loading model(s) from {}'.format(', '.join(args.path)))
Myle Ott's avatar
Myle Ott committed
52
    models, _ = utils.load_ensemble_for_inference(args.path, dataset.src_dict, dataset.dst_dict)
Sergey Edunov's avatar
Sergey Edunov committed
53
54
55

    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
56
    print('| {} {} {} examples'.format(args.data, args.gen_subset, len(dataset.splits[args.gen_subset])))
Sergey Edunov's avatar
Sergey Edunov committed
57

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

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

    # Load alignment dictionary for unknown word replacement
Louis Martin's avatar
Louis Martin committed
72
73
74
75
76
77
    # (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
78
    itr = dataset.eval_dataloader(
79
        args.gen_subset, max_sentences=args.batch_size, max_positions=max_positions,
Myle Ott's avatar
Myle Ott committed
80
        skip_invalid_size_inputs_valid_test=args.skip_invalid_size_inputs_valid_test)
Myle Ott's avatar
Myle Ott committed
81
82
83
84
    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
85
    num_sentences = 0
86
    with utils.build_progress_bar(args, itr) as t:
Louis Martin's avatar
Louis Martin committed
87
88
89
90
91
92
93
94
        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()
95
96
97
98
99
100
101
102
            # 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
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
            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))
119
                    print('A-{}\t{}'.format(sample_id, ' '.join(map(str, alignment))))
Louis Martin's avatar
Louis Martin committed
120
121
122

                # Score only the top hypothesis
                if i == 0:
123
124
                    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
125
126
127
128
129
130
                        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))
131
            t.log({'wps': round(wps_meter.avg)})
Louis Martin's avatar
Louis Martin committed
132
133
134
135
136
            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
137
138
139
140


if __name__ == '__main__':
    main()