-
Notifications
You must be signed in to change notification settings - Fork 1
/
nltk_bleu.py
executable file
·59 lines (54 loc) · 2.63 KB
/
nltk_bleu.py
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
#!/usr/bin/env python
# -*- coding:utf-8 -*-
'''
sent bleu using nltk
'''
import nltk
import argparse
def bleu_calculation(reference='', translation='',
sentence_level=False, output='', weights='0.25 0.25 0.25 0.25'):
weight = [float(v) for v in weights.split()]
if sentence_level:
with open(translation, 'r') as trans, open(reference, 'r') as ref:
if output is None:
for tl in trans:
rl = ref.readline()
tran_list = tl.strip().split(' ')
ref_list = rl.strip().split(' ')
print(nltk.translate.bleu_score.sentence_bleu([ref_list], tran_list, weight))
else:
with open(output, 'w') as outputfile:
for tl in trans:
rl = ref.readline()
tran_list = tl.strip().split(' ')
ref_list = rl.strip().split(' ')
outputfile.write('%s\n' % nltk.translate.bleu_score.sentence_bleu([ref_list], tran_list, weight))
else:
with open(translation, 'r') as trans, open(reference, 'r') as ref:
tran_list, ref_list = [], []
for tl in trans:
rl = ref.readline()
tran_list.append(tl.strip().split(' '))
ref_list.append(rl.strip().split(' '))
print(nltk.translate.bleu_score.corpus_bleu(ref_list, tran_list, weight))
return
if __name__ == '__main__':
parser = argparse.ArgumentParser(description='Arguments for calculating BLEU')
parser.add_argument('-r', '--reference', type=str, required=True,
help="reference file")
parser.add_argument('-t', '--translation', type=str, required=True,
help="translation file")
parser.add_argument('-o', '--output', type=str, default=None, required=False,
help="output BLEU score to this file in segment level scenario \
(default: %(default)s)")
parser.add_argument('-sl', '--sentence-level', action='store_true',
help="print segment level BLEU score (default: %(default)s)")
parser.add_argument('-w', '--weights', type=str, required=False,
default='0.25 0.25 0.25 0.25',
help="weights for n-grams (default: %(default)s)")
args = parser.parse_args()
bleu_calculation(reference=args.reference,
translation=args.translation,
sentence_level=args.sentence_level,
output=args.output,
weights=args.weights)