from mmengine.evaluator import BaseMetric
import evaluate
import torch

class rouge(BaseMetric):

    default_prefix = 'rouge'  # set default_prefix
    rouge = evaluate.load('rouge')

    def process(self, data_batch, data_samples):
        outputs = [''.join(data_samples)]
        labels = data_batch['data_samples']['references']
        self.results.append({
            'outputs': outputs,
            'labels': labels,
        })

    def compute_metrics(self, results):
        outputs_list = []
        labels_list = []
        for utt_result in results:
            outputs_list = outputs_list + utt_result['outputs']
            labels_list = labels_list + utt_result['labels']

        return self.rouge.compute(predictions=outputs_list, references=labels_list)
