__init__.py 4.34 KB
Newer Older
lintangsutawika's avatar
lintangsutawika committed
1
import ast
2
import os
3
from typing import Dict
4

5
from lm_eval import utils
6
from lm_eval.utils import eval_logger
7

8

lintangsutawika's avatar
lintangsutawika committed
9
# Prompt library.
10
11
12
# Stores prompts in a dictionary indexed by 2 levels:
# prompt category name, and prompt name.
# This allows us to access prompts
13
PROMPT_REGISTRY: Dict[str, Dict[str, str]] = {
14
15
    "qa-basic": {
        "question-newline-answer": "Question: {{question}}\nAnswer:",
lintangsutawika's avatar
lintangsutawika committed
16
        "q-newline-a": "Q: {{question}}\nA:",
17
18
19
    },
}

lintangsutawika's avatar
lintangsutawika committed
20

Ethan Smith's avatar
Ethan Smith committed
21
def get_prompt(prompt_id: str, dataset_name: str = None, subset_name: str = None):
lintangsutawika's avatar
lintangsutawika committed
22
    # unpack prompt name
23
    category_name, prompt_name = prompt_id.split(":")
lintangsutawika's avatar
update  
lintangsutawika committed
24
25
26
27
28
    if subset_name is None:
        dataset_full_name = dataset_name
    else:
        dataset_full_name = f"{dataset_name}-{subset_name}"
    eval_logger.info(f"Loading prompt from {category_name} for {dataset_full_name}")
29
    if category_name == "promptsource":
30
        try:
31
            from promptsource.templates import DatasetTemplates
32
        except ModuleNotFoundError:
33
34
35
36
            raise Exception(
                "Tried to load a Promptsource template, but promptsource is not installed ",
                "please install promptsource via pip install lm-eval[promptsource] or pip install -e .[promptsource]",
            )
37
        try:
lintangsutawika's avatar
lintangsutawika committed
38
            if subset_name is None:
39
40
                prompts = DatasetTemplates(dataset_name=dataset_name)
            else:
lintangsutawika's avatar
lintangsutawika committed
41
42
                prompts = DatasetTemplates(
                    dataset_name=dataset_name, subset_name=subset_name
43
                )
lintangsutawika's avatar
lintangsutawika committed
44
45
        except Exception:
            raise ValueError(f"{dataset_name} and {subset_name} not found")
46
47
        if prompt_name in prompts.all_template_names:
            return prompts[prompt_name]
48
        else:
49
50
            raise ValueError(
                f"{prompt_name} not in prompt list {prompts.all_template_names}"
lintangsutawika's avatar
lintangsutawika committed
51
            )
52
53
54
55
56
57
58
59
    elif ".yaml" in category_name:
        import yaml

        with open(category_name, "rb") as file:
            prompt_yaml_file = yaml.full_load(file)

        prompt_string = prompt_yaml_file["prompts"][prompt_name]
        return PromptString(prompt_string)
60
61
62
    else:
        try:
            return PROMPT_REGISTRY[category_name][prompt_name]
lintangsutawika's avatar
lintangsutawika committed
63
        except Exception:
64
65
66
            raise ValueError(
                f"expected only a single `:` as separator between \
                prompt category and name, but got `{prompt_id}` instead"
lintangsutawika's avatar
lintangsutawika committed
67
            )
68
69


lintangsutawika's avatar
lintangsutawika committed
70
71
72
def load_prompt_list(
    use_prompt: str, dataset_name=None, subset_name=None, yaml_path=None, **kwargs
):
73
    category_name, prompt_name = use_prompt.split(":")
74

75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
    if category_name == "promptsource":
        from promptsource.templates import DatasetTemplates

        if subset_name is None:
            prompts = DatasetTemplates(dataset_name=dataset_name)
        else:
            prompts = DatasetTemplates(
                dataset_name=dataset_name, subset_name=subset_name
            )

        prompt_list = utils.pattern_match(prompt_name, prompts.all_template_names)

    elif ".yaml" in category_name:
        import yaml

lintangsutawika's avatar
lintangsutawika committed
90
91
        if yaml_path is not None:
            category_name = os.path.realpath(os.path.join(yaml_path, category_name))
lintangsutawika's avatar
lintangsutawika committed
92

93
94
95
96
97
98
        with open(category_name, "rb") as file:
            prompt_yaml_file = yaml.full_load(file)

        prompt_list = utils.pattern_match(
            prompt_name, prompt_yaml_file["prompts"].keys()
        )
99

lintangsutawika's avatar
lintangsutawika committed
100
    # category_name, *prompt_name = use_prompt.split(":")
lintangsutawika's avatar
lintangsutawika committed
101
102
103
104
105
106
    # TODO allow to multiple prompt naming
    # if len(prompt_name) > 1:
    #     prompt_list = []
    #     for prompt in prompt_name:
    #         prompt_list.append(utils.pattern_match(prompt_name, prompts.all_template_names))
    # else:
lintangsutawika's avatar
lintangsutawika committed
107
    #     prompt_list = utils.pattern_match(prompt_name, prompts.all_template_names)
108
    return [":".join([category_name, prompt]) for prompt in prompt_list]
109
110
111


class PromptString:
lintangsutawika's avatar
lintangsutawika committed
112
    def __init__(self, prompt_string):
113
114
115
116
117
        self.prompt_string = prompt_string

    def apply(self, doc):
        doc_to_text = self.prompt_string["doc_to_text"]
        doc_to_target = self.prompt_string["doc_to_target"]
lintangsutawika's avatar
lintangsutawika committed
118
119
120

        # TODO need a way to process doc_to_choice
        if "doc_to_choice" in self.prompt_string:
121
            raise Exception("Not yet implemented to accept doc_to_choice")
lintangsutawika's avatar
lintangsutawika committed
122

123
124
125
126
        text_string = utils.apply_template(doc_to_text, doc)
        target_string = utils.apply_template(doc_to_target, doc)

        return [text_string, target_string]