Skip to content

GitLab

Switch branch/tag
Find file Normal viewHistoryPermalink
atc.py 3.61 KB
Newer Older
Mo Li's avatar
[Feature] Upgrade the needle-in-a-haystack experiment to Needlebench (#913)  
Mo Li committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104 
from opencompass.openicl.icl_prompt_template import PromptTemplate
from opencompass.openicl.icl_retriever import ZeroRetriever
from opencompass.openicl.icl_inferencer import GenInferencer
from opencompass.datasets.needlebench.atc import NeedleBenchATCOrderedDataset
from opencompass.datasets.needlebench.atc import NeedleBenchATCDataset
from opencompass.datasets.needlebench.origin import NeedleBenchOriginEvaluator
from opencompass.datasets.needlebench.origin import needlebench_postprocess
from opencompass.datasets.needlebench.origin import needlebench_dataset_postprocess

needlebench_reader_cfg = dict(input_columns=['prompt'], output_column='answer')

needlebench_infer_cfg = dict(
    prompt_template=dict(
        type=PromptTemplate,
        template=dict(
            round=[
                dict(role='HUMAN', prompt='{prompt}'),
                dict(role='BOT', prompt='{answer}\n'),
            ]
        )
        ),
    retriever=dict(type=ZeroRetriever),
    inferencer=dict(type=GenInferencer))

needlebench_eval_cfg = dict(
    evaluator=dict(type=NeedleBenchOriginEvaluator),
    pred_postprocessor=dict(type=needlebench_postprocess),
    dataset_postprocessor=dict(type=needlebench_dataset_postprocess),
    pred_role='BOT')

needle_num_list = list(range(2, 100, 3))
document_depth_percent_intervals = 20
repeats = 30
names_path = './data/needlebench/names.json'

needlebench_atc_datasets_zh = []
needlebench_atc_datasets_en = []
needlebench_atc_datasets_zh_ordered = []
needlebench_atc_datasets_en_ordered = []

for num_needles in needle_num_list:
    # ordered English version
    dataset_dict = {
        'abbr': f'needlebench_atc_challenge'
                f'needle_{num_needles}_en_ordered',
        'type': NeedleBenchATCOrderedDataset,
        'path': names_path,
        'num_needles': num_needles,
        'language': 'English',
        'repeats': repeats,
        'reader_cfg': needlebench_reader_cfg,
        'infer_cfg': needlebench_infer_cfg,
        'eval_cfg': needlebench_eval_cfg
    }
    needlebench_atc_datasets_en_ordered.append(dataset_dict)


for num_needles in needle_num_list:
    # ordered Chinese version
    dataset_dict = {
        'abbr': f'needlebench_atc_challenge'
                f'needle_{num_needles}_zh_ordered',
        'type': NeedleBenchATCOrderedDataset,
        'path': names_path,
        'num_needles': num_needles,
        'language': 'Chinese',
        'repeats': repeats,
        'reader_cfg': needlebench_reader_cfg,
        'infer_cfg': needlebench_infer_cfg,
        'eval_cfg': needlebench_eval_cfg
    }
    needlebench_atc_datasets_zh_ordered.append(dataset_dict)

for num_needles in needle_num_list:
    # standard English version
    dataset_dict = {
        'abbr': f'needlebench_atc_challenge'
                f'needle_{num_needles}_en',
        'type': NeedleBenchATCDataset,
        'path': names_path,
        'num_needles': num_needles,
        'language': 'English',
        'repeats': repeats,
        'reader_cfg': needlebench_reader_cfg,
        'infer_cfg': needlebench_infer_cfg,
        'eval_cfg': needlebench_eval_cfg
    }
    needlebench_atc_datasets_en.append(dataset_dict)

for num_needles in needle_num_list:
    # standard Chinese version
    dataset_dict = {
        'abbr': f'needlebench_atc_challenge'
                f'needle_{num_needles}_zh',
        'type': NeedleBenchATCDataset,
        'path': names_path,
        'num_needles': num_needles,
        'language': 'Chinese',
        'repeats': repeats,
        'reader_cfg': needlebench_reader_cfg,
        'infer_cfg': needlebench_infer_cfg,
        'eval_cfg': needlebench_eval_cfg
    }
    needlebench_atc_datasets_zh.append(dataset_dict)