File size: 2,691 Bytes
158b61b |
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 |
# -*- encoding: utf-8 -*-
import argparse
import os
import time
import pyrouge
import shutil
import sys
import codecs
from onmt.utils.logging import init_logger, logger
def test_rouge(cand, ref):
"""Calculate ROUGE scores of sequences passed as an iterator
e.g. a list of str, an open file, StringIO or even sys.stdin
"""
current_time = time.strftime('%Y-%m-%d-%H-%M-%S', time.localtime())
tmp_dir = ".rouge-tmp-{}".format(current_time)
try:
if not os.path.isdir(tmp_dir):
os.mkdir(tmp_dir)
os.mkdir(tmp_dir + "/candidate")
os.mkdir(tmp_dir + "/reference")
candidates = [line.strip() for line in cand]
references = [line.strip() for line in ref]
assert len(candidates) == len(references)
cnt = len(candidates)
for i in range(cnt):
if len(references[i]) < 1:
continue
with open(tmp_dir + "/candidate/cand.{}.txt".format(i), "w",
encoding="utf-8") as f:
f.write(candidates[i])
with open(tmp_dir + "/reference/ref.{}.txt".format(i), "w",
encoding="utf-8") as f:
f.write(references[i])
r = pyrouge.Rouge155()
r.model_dir = tmp_dir + "/reference/"
r.system_dir = tmp_dir + "/candidate/"
r.model_filename_pattern = 'ref.#ID#.txt'
r.system_filename_pattern = r'cand.(\d+).txt'
rouge_results = r.convert_and_evaluate()
results_dict = r.output_to_dict(rouge_results)
return results_dict
finally:
pass
if os.path.isdir(tmp_dir):
shutil.rmtree(tmp_dir)
def rouge_results_to_str(results_dict):
return ">> ROUGE(1/2/3/L/SU4): {:.2f}/{:.2f}/{:.2f}/{:.2f}/{:.2f}".format(
results_dict["rouge_1_f_score"] * 100,
results_dict["rouge_2_f_score"] * 100,
results_dict["rouge_3_f_score"] * 100,
results_dict["rouge_l_f_score"] * 100,
results_dict["rouge_su*_f_score"] * 100)
if __name__ == "__main__":
init_logger('test_rouge.log')
parser = argparse.ArgumentParser()
parser.add_argument('-c', type=str, default="candidate.txt",
help='candidate file')
parser.add_argument('-r', type=str, default="reference.txt",
help='reference file')
args = parser.parse_args()
if args.c.upper() == "STDIN":
candidates = sys.stdin
else:
candidates = codecs.open(args.c, encoding="utf-8")
references = codecs.open(args.r, encoding="utf-8")
results_dict = test_rouge(candidates, references)
logger.info(rouge_results_to_str(results_dict))
|