-
Notifications
You must be signed in to change notification settings - Fork 202
/
SelfBleu.py
87 lines (75 loc) · 2.87 KB
/
SelfBleu.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
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
import os
from multiprocessing import Pool
import nltk
from nltk.translate.bleu_score import SmoothingFunction
from utils.metrics.Metrics import Metrics
class SelfBleu(Metrics):
def __init__(self, test_text='', gram=3):
super().__init__()
self.name = 'Self-Bleu'
self.test_data = test_text
self.gram = gram
self.sample_size = 500
self.reference = None
self.is_first = True
def get_name(self):
return self.name
def get_score(self, is_fast=True, ignore=False):
if ignore:
return 0
if self.is_first:
self.get_reference()
self.is_first = False
if is_fast:
return self.get_bleu_fast()
return self.get_bleu_parallel()
def get_reference(self):
if self.reference is None:
reference = list()
with open(self.test_data) as real_data:
for text in real_data:
text = nltk.word_tokenize(text)
reference.append(text)
self.reference = reference
return reference
else:
return self.reference
def get_bleu(self):
ngram = self.gram
bleu = list()
reference = self.get_reference()
weight = tuple((1. / ngram for _ in range(ngram)))
with open(self.test_data) as test_data:
for hypothesis in test_data:
hypothesis = nltk.word_tokenize(hypothesis)
bleu.append(nltk.translate.bleu_score.sentence_bleu(reference, hypothesis, weight,
smoothing_function=SmoothingFunction().method1))
return sum(bleu) / len(bleu)
def calc_bleu(self, reference, hypothesis, weight):
return nltk.translate.bleu_score.sentence_bleu(reference, hypothesis, weight,
smoothing_function=SmoothingFunction().method1)
def get_bleu_fast(self):
reference = self.get_reference()
# random.shuffle(reference)
reference = reference[0:self.sample_size]
return self.get_bleu_parallel(reference=reference)
def get_bleu_parallel(self, reference=None):
ngram = self.gram
if reference is None:
reference = self.get_reference()
weight = tuple((1. / ngram for _ in range(ngram)))
pool = Pool(os.cpu_count())
result = list()
sentence_num = len(reference)
for index in range(sentence_num):
hypothesis = reference[index]
other = reference[:index] + reference[index+1:]
result.append(pool.apply_async(self.calc_bleu, args=(other, hypothesis, weight)))
score = 0.0
cnt = 0
for i in result:
score += i.get()
cnt += 1
pool.close()
pool.join()
return score / cnt