Register
Login
Resources
Docs Blog Datasets Glossary Case Studies Tutorials & Webinars
Product
Data Engine LLMs Platform Enterprise
Pricing Explore
Connect to our Discord channel

generate.py 7.6 KB

You have to be logged in to leave a comment. Sign In
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
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
  1. #!/usr/bin/env python3 -u
  2. # Copyright (c) 2017-present, Facebook, Inc.
  3. # All rights reserved.
  4. #
  5. # This source code is licensed under the license found in the LICENSE file in
  6. # the root directory of this source tree. An additional grant of patent rights
  7. # can be found in the PATENTS file in the same directory.
  8. """
  9. Translate pre-processed data with a trained model.
  10. """
  11. import torch
  12. from fairseq import bleu, options, progress_bar, tasks, tokenizer, utils
  13. from fairseq.meters import StopwatchMeter, TimeMeter
  14. from fairseq.sequence_generator import SequenceGenerator
  15. from fairseq.sequence_scorer import SequenceScorer
  16. from fairseq.utils import import_user_module
  17. def main(args):
  18. assert args.path is not None, '--path required for generation!'
  19. assert not args.sampling or args.nbest == args.beam, \
  20. '--sampling requires --nbest to be equal to --beam'
  21. assert args.replace_unk is None or args.raw_text, \
  22. '--replace-unk requires a raw text dataset (--raw-text)'
  23. import_user_module(args)
  24. if args.max_tokens is None and args.max_sentences is None:
  25. args.max_tokens = 12000
  26. print(args)
  27. use_cuda = torch.cuda.is_available() and not args.cpu
  28. # Load dataset splits
  29. task = tasks.setup_task(args)
  30. task.load_dataset(args.gen_subset)
  31. print('| {} {} {} examples'.format(args.data, args.gen_subset, len(task.dataset(args.gen_subset))))
  32. # Set dictionaries
  33. src_dict = task.source_dictionary
  34. tgt_dict = task.target_dictionary
  35. # Load ensemble
  36. print('| loading model(s) from {}'.format(args.path))
  37. models, _model_args = utils.load_ensemble_for_inference(
  38. args.path.split(':'), task, model_arg_overrides=eval(args.model_overrides),
  39. )
  40. # Optimize ensemble for generation
  41. for model in models:
  42. model.make_generation_fast_(
  43. beamable_mm_beam_size=None if args.no_beamable_mm else args.beam,
  44. need_attn=args.print_alignment,
  45. )
  46. if args.fp16:
  47. model.half()
  48. # Load alignment dictionary for unknown word replacement
  49. # (None if no unknown word replacement, empty if no path to align dictionary)
  50. align_dict = utils.load_align_dict(args.replace_unk)
  51. # Load dataset (possibly sharded)
  52. itr = task.get_batch_iterator(
  53. dataset=task.dataset(args.gen_subset),
  54. max_tokens=args.max_tokens,
  55. max_sentences=args.max_sentences,
  56. max_positions=utils.resolve_max_positions(
  57. task.max_positions(),
  58. *[model.max_positions() for model in models]
  59. ),
  60. ignore_invalid_inputs=args.skip_invalid_size_inputs_valid_test,
  61. required_batch_size_multiple=8,
  62. num_shards=args.num_shards,
  63. shard_id=args.shard_id,
  64. num_workers=args.num_workers,
  65. ).next_epoch_itr(shuffle=False)
  66. # Initialize generator
  67. gen_timer = StopwatchMeter()
  68. if args.score_reference:
  69. translator = SequenceScorer(models, task.target_dictionary)
  70. else:
  71. translator = SequenceGenerator(
  72. models, task.target_dictionary, beam_size=args.beam, minlen=args.min_len,
  73. stop_early=(not args.no_early_stop), normalize_scores=(not args.unnormalized),
  74. len_penalty=args.lenpen, unk_penalty=args.unkpen,
  75. sampling=args.sampling, sampling_topk=args.sampling_topk, sampling_temperature=args.sampling_temperature,
  76. diverse_beam_groups=args.diverse_beam_groups, diverse_beam_strength=args.diverse_beam_strength,
  77. match_source_len=args.match_source_len, no_repeat_ngram_size=args.no_repeat_ngram_size,
  78. )
  79. if use_cuda:
  80. translator.cuda()
  81. # Generate and compute BLEU score
  82. if args.sacrebleu:
  83. scorer = bleu.SacrebleuScorer()
  84. else:
  85. scorer = bleu.Scorer(tgt_dict.pad(), tgt_dict.eos(), tgt_dict.unk())
  86. num_sentences = 0
  87. has_target = True
  88. with progress_bar.build_progress_bar(args, itr) as t:
  89. if args.score_reference:
  90. translations = translator.score_batched_itr(t, cuda=use_cuda, timer=gen_timer)
  91. else:
  92. translations = translator.generate_batched_itr(
  93. t, maxlen_a=args.max_len_a, maxlen_b=args.max_len_b,
  94. cuda=use_cuda, timer=gen_timer, prefix_size=args.prefix_size,
  95. )
  96. wps_meter = TimeMeter()
  97. for sample_id, src_tokens, target_tokens, hypos in translations:
  98. # Process input and ground truth
  99. has_target = target_tokens is not None
  100. target_tokens = target_tokens.int().cpu() if has_target else None
  101. # Either retrieve the original sentences or regenerate them from tokens.
  102. if align_dict is not None:
  103. src_str = task.dataset(args.gen_subset).src.get_original_text(sample_id)
  104. target_str = task.dataset(args.gen_subset).tgt.get_original_text(sample_id)
  105. else:
  106. src_str = src_dict.string(src_tokens, args.remove_bpe)
  107. if has_target:
  108. target_str = tgt_dict.string(target_tokens, args.remove_bpe, escape_unk=True)
  109. if not args.quiet:
  110. print('S-{}\t{}'.format(sample_id, src_str))
  111. if has_target:
  112. print('T-{}\t{}'.format(sample_id, target_str))
  113. # Process top predictions
  114. for i, hypo in enumerate(hypos[:min(len(hypos), args.nbest)]):
  115. hypo_tokens, hypo_str, alignment = utils.post_process_prediction(
  116. hypo_tokens=hypo['tokens'].int().cpu(),
  117. src_str=src_str,
  118. alignment=hypo['alignment'].int().cpu() if hypo['alignment'] is not None else None,
  119. align_dict=align_dict,
  120. tgt_dict=tgt_dict,
  121. remove_bpe=args.remove_bpe,
  122. )
  123. if not args.quiet:
  124. print('H-{}\t{}\t{}'.format(sample_id, hypo['score'], hypo_str))
  125. print('P-{}\t{}'.format(
  126. sample_id,
  127. ' '.join(map(
  128. lambda x: '{:.4f}'.format(x),
  129. hypo['positional_scores'].tolist(),
  130. ))
  131. ))
  132. if args.print_alignment:
  133. print('A-{}\t{}'.format(
  134. sample_id,
  135. ' '.join(map(lambda x: str(utils.item(x)), alignment))
  136. ))
  137. # Score only the top hypothesis
  138. if has_target and i == 0:
  139. if align_dict is not None or args.remove_bpe is not None:
  140. # Convert back to tokens for evaluation with unk replacement and/or without BPE
  141. target_tokens = tokenizer.Tokenizer.tokenize(
  142. target_str, tgt_dict, add_if_not_exist=True)
  143. if hasattr(scorer, 'add_string'):
  144. scorer.add_string(target_str, hypo_str)
  145. else:
  146. scorer.add(target_tokens, hypo_tokens)
  147. wps_meter.update(src_tokens.size(0))
  148. t.log({'wps': round(wps_meter.avg)})
  149. num_sentences += 1
  150. print('| Translated {} sentences ({} tokens) in {:.1f}s ({:.2f} sentences/s, {:.2f} tokens/s)'.format(
  151. num_sentences, gen_timer.n, gen_timer.sum, num_sentences / gen_timer.sum, 1. / gen_timer.avg))
  152. if has_target:
  153. print('| Generate {} with beam={}: {}'.format(args.gen_subset, args.beam, scorer.result_string()))
  154. def cli_main():
  155. parser = options.get_generation_parser()
  156. args = options.parse_args_and_arch(parser)
  157. main(args)
  158. if __name__ == '__main__':
  159. cli_main()
Tip!

Press p or to see the previous file or, n or to see the next file

Comments

Loading...