aurora

1
from dataclasses import dataclass
2
from typing import TYPE_CHECKING, Dict, List, Tuple
3

4
from llmtuner.extras.constants import CHOICES
5

6
if TYPE_CHECKING:
7
    from datasets import Dataset
8

9

10
@dataclass
11
class EvalTemplate:
12

13
    system: str
14
    choice: str
15
    answer: str
16
    prefix: str
17

18
    def parse_example(
19
        self,
20
        example: Dict[str, str]
21
    ) -> Tuple[str, str]:
22
        candidates = [self.choice.format(choice=ch, content=example[ch]) for ch in CHOICES if ch in example]
23
        return "".join([example["question"]] + candidates + [self.answer]), example["answer"]
24

25
    def format_example(
26
        self,
27
        target_data: Dict[str, str],
28
        support_set: "Dataset",
29
        subject_name: str,
30
        use_history: bool
31
    ) -> Tuple[str, str, List[Tuple[str, str]]]:
32
        query, resp = self.parse_example(target_data)
33
        history = [self.parse_example(support_set[k]) for k in range(len(support_set))]
34

35
        if len(history):
36
            temp = history.pop(0)
37
            history.insert(0, (self.system.format(subject=subject_name) + temp[0], temp[1]))
38
        else:
39
            query = self.system.format(subject=subject_name) + query
40

41
        if not use_history:
42
            query = "\n\n".join(["".join(item) for item in history] + [query])
43
            history = []
44
        return query.strip(), resp, history
45

46

47
eval_templates: Dict[str, EvalTemplate] = {}
48

49

50
def register_eval_template(
51
    name: str,
52
    system: str,
53
    choice: str,
54
    answer: str,
55
    prefix: str
56
) -> None:
57
    eval_templates[name] = EvalTemplate(
58
        system=system,
59
        choice=choice,
60
        answer=answer,
61
        prefix=prefix
62
    )
63

64

65
def get_eval_template(name: str) -> EvalTemplate:
66
    eval_template = eval_templates.get(name, None)
67
    assert eval_template is not None, "Template {} does not exist.".format(name)
68
    return eval_template
69

70

71
register_eval_template(
72
    name="en",
73
    system="The following are multiple choice questions (with answers) about {subject}.\n\n",
74
    choice="\n{choice}. {content}",
75
    answer="\nAnswer: ",
76
    prefix=" "
77
)
78

79

80
register_eval_template(
81
    name="zh",
82
    system="以下是中国关于{subject}考试的单项选择题，请选出其中的正确答案。\n\n",
83
    choice="\n{choice}. {content}",
84
    answer="\n答案：",
85
    prefix="\n"
86
)
87