aurora

1
import numpy as np
2
from dataclasses import dataclass
3
from typing import TYPE_CHECKING, Dict, Sequence, Tuple, Union
4

5
from llmtuner.extras.constants import IGNORE_INDEX
6
from llmtuner.extras.packages import (
7
    is_jieba_available, is_nltk_available, is_rouge_available
8
)
9

10
if TYPE_CHECKING:
11
    from transformers.tokenization_utils import PreTrainedTokenizer
12

13
if is_jieba_available():
14
    import jieba
15

16
if is_nltk_available():
17
    from nltk.translate.bleu_score import sentence_bleu, SmoothingFunction
18

19
if is_rouge_available():
20
    from rouge_chinese import Rouge
21

22

23
@dataclass
24
class ComputeMetrics:
25
    r"""
26
    Wraps the tokenizer into metric functions, used in Seq2SeqPeftTrainer.
27
    """
28

29
    tokenizer: "PreTrainedTokenizer"
30

31
    def __call__(self, eval_preds: Sequence[Union[np.ndarray, Tuple[np.ndarray]]]) -> Dict[str, float]:
32
        r"""
33
        Uses the model predictions to compute metrics.
34
        """
35
        preds, labels = eval_preds
36
        score_dict = {"rouge-1": [], "rouge-2": [], "rouge-l": [], "bleu-4": []}
37

38
        preds = np.where(preds != IGNORE_INDEX, preds, self.tokenizer.pad_token_id)
39
        labels = np.where(labels != IGNORE_INDEX, labels, self.tokenizer.pad_token_id)
40

41
        decoded_preds = self.tokenizer.batch_decode(preds, skip_special_tokens=True)
42
        decoded_labels = self.tokenizer.batch_decode(labels, skip_special_tokens=True)
43

44
        for pred, label in zip(decoded_preds, decoded_labels):
45
            hypothesis = list(jieba.cut(pred))
46
            reference = list(jieba.cut(label))
47

48
            if len(" ".join(hypothesis).split()) == 0 or len(" ".join(reference).split()) == 0:
49
                result = {"rouge-1": {"f": 0.0}, "rouge-2": {"f": 0.0}, "rouge-l": {"f": 0.0}}
50
            else:
51
                rouge = Rouge()
52
                scores = rouge.get_scores(" ".join(hypothesis), " ".join(reference))
53
                result = scores[0]
54

55
            for k, v in result.items():
56
                score_dict[k].append(round(v["f"] * 100, 4))
57

58
            bleu_score = sentence_bleu([list(label)], list(pred), smoothing_function=SmoothingFunction().method3)
59
            score_dict["bleu-4"].append(round(bleu_score * 100, 4))
60

61
        return {k: float(np.mean(v)) for k, v in score_dict.items()}
62