__init__.py 1.77 KB
Newer Older
1
2
3
from lm_eval.logger import eval_logger
from promptsource.templates import DatasetTemplates

4
# TODO: decide whether we want jinja2 or f-string prompts. would it be cursed to support both?
lintangsutawika's avatar
lintangsutawika committed
5
# Prompt library.
6
7
8
9
10
11
# Stores prompts in a dictionary indexed by 2 levels:
# prompt category name, and prompt name.
# This allows us to access prompts
PROMPT_REGISTRY = {
    "qa-basic": {
        "question-newline-answer": "Question: {{question}}\nAnswer:",
lintangsutawika's avatar
lintangsutawika committed
12
        "q-newline-a": "Q: {{question}}\nA:",
13
14
15
    },
}

lintangsutawika's avatar
lintangsutawika committed
16

17
def get_prompt(prompt_id: str, dataset_name=None, subset_name=None):
lintangsutawika's avatar
lintangsutawika committed
18
    # unpack prompt name
19
    category_name, prompt_name = prompt_id.split(":")
lintangsutawika's avatar
update  
lintangsutawika committed
20
21
22
23
24
    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}")
25
26
    if category_name == "promptsource":
        try:
lintangsutawika's avatar
lintangsutawika committed
27
            if subset_name is None:
28
29
                prompts = DatasetTemplates(dataset_name=dataset_name)
            else:
lintangsutawika's avatar
lintangsutawika committed
30
31
                prompts = DatasetTemplates(
                    dataset_name=dataset_name, subset_name=subset_name
32
                )
lintangsutawika's avatar
lintangsutawika committed
33
34
        except Exception:
            raise ValueError(f"{dataset_name} and {subset_name} not found")
35
36
        if prompt_name in prompts.all_template_names:
            return prompts[prompt_name]
37
        else:
38
39
            raise ValueError(
                f"{prompt_name} not in prompt list {prompts.all_template_names}"
lintangsutawika's avatar
lintangsutawika committed
40
            )
41
42
43
    else:
        try:
            return PROMPT_REGISTRY[category_name][prompt_name]
lintangsutawika's avatar
lintangsutawika committed
44
        except Exception:
45
46
47
            raise ValueError(
                f"expected only a single `:` as separator between \
                prompt category and name, but got `{prompt_id}` instead"
lintangsutawika's avatar
lintangsutawika committed
48
            )