from torchmetrics.text import BLEUScore

from torchmetrics.text import WordErrorRate

from torchmetrics.functional.text.rouge import rouge_score

def compute_metrics(preds, labels):

    result={}

    wer = WordErrorRate()

    result['wer']=wer(preds,labels).item()

    rouge_result=compute_rouge(preds,labels)

    for k,v in rouge_result.items():

        result[k]=v.item()

    labels=[[label] for i,label in enumerate(labels)]

    for i in range(1,5):

        bleu=BLEUScore(n_gram=i)

        result[f'bleu-{i}']=bleu(preds,labels).item()

    return result

def compute_rouge(preds, labels):

    metrics={}

    for decoded_label, decoded_pred in zip(labels, preds):

        metric=rouge_score(decoded_pred,decoded_label)

        for key in metric.keys():

            metrics[key]=metrics.get(key, 0) + metric[key]

    for key in metrics.keys():

        metrics[key]=metrics[key]/len(labels)*100

    return metrics