-
Notifications
You must be signed in to change notification settings - Fork 0
/
rouge_score.py
63 lines (52 loc) · 1.98 KB
/
rouge_score.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
import argparse
import pandas as pd
from tqdm import tqdm
from rouge_metric import Rouge
class RougeScorer():
def __init__(self):
self.rouge_evaluator = Rouge(
metrics=["rouge-n", "rouge-l"],
max_n=2,
limit_length=True,
length_limit=1000,
length_limit_type="words",
use_tokenizer=True,
apply_avg=True,
apply_best=False,
alpha=0.5, # Default F1_score
weight_factor=1.2,
)
def compute_rouge(self, ref_df, hyp_df):
#ref_df = pd.read_csv(ref_path)
#hyp_df = pd.read_csv(hyp_path)
hyp_df.iloc[:,1] = hyp_df.iloc[:,1].fillna(' ')
ids = ref_df['id']
hyp_df = hyp_df[hyp_df['id'].isin(ids)]
hyp_df.index = ref_df.index
ref_df = ref_df.sort_values(by=["id"])
hyp_df = hyp_df.sort_values(by=["id"])
ref_df["id"] = ref_df["id"].astype(int)
hyp_df["id"] = hyp_df["id"].astype(int)
hyps = [tuple(row) for row in hyp_df.values]
refs = [tuple(row) for row in ref_df.values]
reference_summaries = []
generated_summaries = []
for ref_tp, hyp_tp in zip(refs, hyps):
ref_id, ref = ref_tp
hyp_id, hyp = hyp_tp
assert ref_id == hyp_id
reference_summaries.append(ref)
generated_summaries.append(hyp)
scores = self.rouge_evaluator.get_scores(generated_summaries, reference_summaries)
str_scores = self.format_rouge_scores(scores)
#self.save_rouge_scores(str_scores)
return str_scores
def save_rouge_scores(self, str_scores):
with open("rouge_scores.txt", "w") as output:
output.write(str_scores)
def format_rouge_scores(self, scores):
return "{:.3f},{:.3f},{:.3f}".format(
scores["rouge-1"]["f"],
scores["rouge-2"]["f"],
scores["rouge-l"]["f"],
)