rouge.py

"""
Adapted from https://github.com/jbdel/vilmedic/blob/main/vilmedic/blocks/scorers/NLG/rouge/rouge.py
"""
import torch.nn as nn
from rouge_score import rouge_scorer
from six.moves import zip_longest
import numpy as np


class Rouge(nn.Module):
    def __init__(self, rouges, **kwargs):
        super().__init__()
        rouges = [r.replace('rougel', 'rougeL') for r in rouges]
        self.scorer = rouge_scorer.RougeScorer(rouges, use_stemmer=True)
        self.rouges = rouges

    def forward(self, refs, hyps):
        scores = []
        for target_rec, prediction_rec in zip_longest(refs, hyps):
            if target_rec is None or prediction_rec is None:
                raise ValueError("Must have equal number of lines across target and "
                                 "prediction.")
            scores.append(self.scorer.score(target_rec, prediction_rec))
        
        f1_rouge = [s[self.rouges[0]].fmeasure for s in scores]
        return np.mean(f1_rouge), f1_rouge


class Rouge1(Rouge):
    def __init__(self, **kwargs):
        super(Rouge1, self).__init__(rouges=['rouge1'])


class Rouge2(Rouge):
    def __init__(self, **kwargs):
        super(Rouge2, self).__init__(rouges=['rouge2'])


class RougeL(Rouge):
    def __init__(self, **kwargs):
        super(RougeL, self).__init__(rouges=['rougeL'])