"vscode:/vscode.git/clone" did not exist on "0e4ab6d31cbbaa3cc8be2046dfecfc5bde375494"
eval_lm.py 4.26 KB
Newer Older
alexeib's avatar
alexeib committed
1
2
3
4
5
6
7
8
9
10
11
#!/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 numpy as np
import torch

Myle Ott's avatar
Myle Ott committed
12
from fairseq import data, options, progress_bar, tasks, utils
alexeib's avatar
alexeib committed
13
14
15
16
17
18
19
from fairseq.meters import StopwatchMeter, TimeMeter
from fairseq.sequence_scorer import SequenceScorer


def main(args):
    assert args.path is not None, '--path required for evaluation!'

20
    args.tokens_per_sample = getattr(args, 'tokens_per_sample', 1024)
Myle Ott's avatar
Myle Ott committed
21
    print(args)
alexeib's avatar
alexeib committed
22
23

    use_cuda = torch.cuda.is_available() and not args.cpu
Myle Ott's avatar
Myle Ott committed
24
25
26
27
28

    # Load dataset splits
    task = tasks.setup_task(args)
    task.load_dataset(args.gen_subset)
    print('| {} {} {} examples'.format(args.data, args.gen_subset, len(task.dataset(args.gen_subset))))
alexeib's avatar
alexeib committed
29
30

    # Load ensemble
alexeib's avatar
alexeib committed
31
    print('| loading model(s) from {}'.format(args.path))
32
    models, _ = utils.load_ensemble_for_inference(args.path.split(':'), task)
alexeib's avatar
alexeib committed
33
34
35
36

    # Optimize ensemble for generation and set the source and dest dicts on the model (required by scorer)
    for model in models:
        model.make_generation_fast_()
Myle Ott's avatar
Myle Ott committed
37
38
        if args.fp16:
            model.half()
alexeib's avatar
alexeib committed
39

alexeib's avatar
alexeib committed
40
41
    assert len(models) > 0

Myle Ott's avatar
Myle Ott committed
42
43
    itr = data.EpochBatchIterator(
        dataset=task.dataset(args.gen_subset),
Alexei Baevski's avatar
Alexei Baevski committed
44
45
        max_tokens=args.max_tokens or 36000,
        max_sentences=args.max_sentences,
alexeib's avatar
alexeib committed
46
        max_positions=models[0].max_positions(),
Myle Ott's avatar
Myle Ott committed
47
48
        num_shards=args.num_shards,
        shard_id=args.shard_id,
49
        ignore_invalid_inputs=True,
Myle Ott's avatar
Myle Ott committed
50
    ).next_epoch_itr(shuffle=False)
alexeib's avatar
alexeib committed
51
52

    gen_timer = StopwatchMeter()
Myle Ott's avatar
Myle Ott committed
53
    scorer = SequenceScorer(models, task.target_dictionary)
alexeib's avatar
alexeib committed
54
55
56
57
58
    if use_cuda:
        scorer.cuda()

    score_sum = 0.
    count = 0
Alexei Baevski's avatar
Alexei Baevski committed
59
60
61
62

    if args.remove_bpe is not None:
        bpe_cont = args.remove_bpe.rstrip()
        bpe_toks = set(i for i in range(len(task.dictionary)) if task.dictionary[i].endswith(bpe_cont))
63
        bpe_len = len(bpe_cont)
Alexei Baevski's avatar
Alexei Baevski committed
64
65
    else:
        bpe_toks = None
66
        bpe_len = 0
Alexei Baevski's avatar
Alexei Baevski committed
67

alexeib's avatar
alexeib committed
68
69
70
71
72
73
    with progress_bar.build_progress_bar(args, itr) as t:
        results = scorer.score_batched_itr(t, cuda=use_cuda, timer=gen_timer)
        wps_meter = TimeMeter()
        for _, src_tokens, __, hypos in results:
            for hypo in hypos:
                pos_scores = hypo['positional_scores']
Alexei Baevski's avatar
Alexei Baevski committed
74
75
76
77
78
79
80
81
82

                skipped_toks = 0
                if bpe_toks is not None:
                    for i in range(len(hypo['tokens']) - 1):
                        if hypo['tokens'][i].item() in bpe_toks:
                            skipped_toks += 1
                            pos_scores[i + 1] += pos_scores[i]
                            pos_scores[i] = 0

alexeib's avatar
alexeib committed
83
84
85
                inf_scores = pos_scores.eq(float('inf')) | pos_scores.eq(float('-inf'))
                if inf_scores.any():
                    print('| Skipping tokens with inf scores:',
Myle Ott's avatar
Myle Ott committed
86
                          task.target_dictionary.string(hypo['tokens'][inf_scores.nonzero()]))
alexeib's avatar
alexeib committed
87
88
                    pos_scores = pos_scores[(~inf_scores).nonzero()]
                score_sum += pos_scores.sum()
Alexei Baevski's avatar
Alexei Baevski committed
89
                count += pos_scores.numel() - skipped_toks
90
91
92
93
94
95
96
97
98
99
100
101
102
103

                if args.output_word_probs:
                    w = ''
                    word_prob = []
                    for i in range(len(hypo['tokens'])):
                        w_ind = hypo['tokens'][i].item()
                        w += task.dictionary[w_ind]
                        if bpe_toks is not None and w_ind in bpe_toks:
                            w = w[:-bpe_len]
                        else:
                            word_prob.append((w, pos_scores[i].item()))
                            w = ''
                    print('\t'.join('{} [{:2f}]'.format(x[0], x[1]) for x in word_prob))

alexeib's avatar
alexeib committed
104
105
106
107
108
109
110
111
112
113
            wps_meter.update(src_tokens.size(0))
            t.log({'wps': round(wps_meter.avg)})

    avg_nll_loss = -score_sum / count
    print('| Evaluated {} tokens in {:.1f}s ({:.2f} tokens/s)'.format(gen_timer.n, gen_timer.sum, 1. / gen_timer.avg))
    print('| Loss: {:.4f}, Perplexity: {:.2f}'.format(avg_nll_loss, np.exp(avg_nll_loss)))


if __name__ == '__main__':
    parser = options.get_eval_lm_parser()
114
    args = options.parse_args_and_arch(parser)
alexeib's avatar
alexeib committed
115
    main(args)