1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
|
- #!/usr/bin/env python3 -u
- # 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
- from fairseq import bleu, data, options, progress_bar, tokenizer, utils
- from fairseq.meters import StopwatchMeter, TimeMeter
- from fairseq.sequence_generator import SequenceGenerator
- from fairseq.sequence_scorer import SequenceScorer
- def main(args):
- print(args)
- use_cuda = torch.cuda.is_available() and not args.cpu
- # Load dataset
- 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,
- )
- 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
- print('| loading model(s) from {}'.format(', '.join(args.path)))
- models, _ = utils.load_ensemble_for_inference(args.path, dataset.src_dict, dataset.dst_dict)
- print('| [{}] dictionary: {} types'.format(dataset.src, len(dataset.src_dict)))
- print('| [{}] dictionary: {} types'.format(dataset.dst, len(dataset.dst_dict)))
- print('| {} {} {} examples'.format(args.data, args.gen_subset, len(dataset.splits[args.gen_subset])))
- # Optimize ensemble for generation
- for model in models:
- model.make_generation_fast_(
- beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
- )
- # Load alignment dictionary for unknown word replacement
- # (None if no unknown word replacement, empty if no path to align dictionary)
- align_dict = utils.load_align_dict(args.replace_unk)
- # Load dataset (possibly sharded)
- max_positions = min(model.max_encoder_positions() for model in models)
- itr = dataset.eval_dataloader(
- args.gen_subset,
- max_sentences=args.max_sentences,
- max_positions=max_positions,
- skip_invalid_size_inputs_valid_test=args.skip_invalid_size_inputs_valid_test,
- )
- 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)
- # Initialize generator
- gen_timer = StopwatchMeter()
- if args.score_reference:
- translator = SequenceScorer(models)
- else:
- translator = SequenceGenerator(
- models, beam_size=args.beam, stop_early=(not args.no_early_stop),
- normalize_scores=(not args.unnormalized), len_penalty=args.lenpen,
- unk_penalty=args.unkpen)
- if use_cuda:
- translator.cuda()
- # Generate and compute BLEU score
- scorer = bleu.Scorer(dataset.dst_dict.pad(), dataset.dst_dict.eos(), dataset.dst_dict.unk())
- num_sentences = 0
- has_target = True
- with progress_bar.build_progress_bar(args, itr) as t:
- if args.score_reference:
- translations = translator.score_batched_itr(t, cuda=use_cuda, timer=gen_timer)
- else:
- translations = translator.generate_batched_itr(
- t, maxlen_a=args.max_len_a, maxlen_b=args.max_len_b,
- cuda=use_cuda, timer=gen_timer, prefix_size=args.prefix_size)
- wps_meter = TimeMeter()
- for sample_id, src_tokens, target_tokens, hypos in translations:
- # Process input and ground truth
- has_target = target_tokens is not None
- target_tokens = target_tokens.int().cpu() if has_target else None
- # 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) if has_target else ''
- if not args.quiet:
- print('S-{}\t{}'.format(sample_id, src_str))
- if has_target:
- 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))
- print('P-{}\t{}'.format(
- sample_id,
- ' '.join(map(
- lambda x: '{:.4f}'.format(x),
- hypo['positional_scores'].tolist(),
- ))
- ))
- print('A-{}\t{}'.format(
- sample_id,
- ' '.join(map(lambda x: str(utils.item(x)), alignment))
- ))
- # Score only the top hypothesis
- if has_target and i == 0:
- 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
- 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))
- t.log({'wps': round(wps_meter.avg)})
- 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))
- if has_target:
- print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))
- if __name__ == '__main__':
- parser = options.get_generation_parser()
- args = parser.parse_args()
- main(args)
|