niah_utils.py 4.73 KB
Newer Older
Baber's avatar
cleanup  
Baber committed
1
import itertools
Baber's avatar
nit  
Baber committed
2
import logging
Baber's avatar
cleanup  
Baber committed
3
from typing import Generator
Baber's avatar
cleanup  
Baber committed
4
5
6

import datasets

Baber's avatar
Baber committed
7
from lm_eval.tasks.ruler.common_utils import DEFAULT_SEQ_LENGTHS, get_tokenizer
Baber's avatar
cleanup  
Baber committed
8
from lm_eval.tasks.ruler.prepare_niah import generate_samples, get_haystack
Baber's avatar
Baber committed
9

Baber's avatar
cleanup  
Baber committed
10
11

TEMPLATE = """Some special magic {type_needle_v} are hidden within the following text. Make sure to memorize it. I will quiz you about the {type_needle_v} afterwards.\n{context}\nWhat are all the special magic {type_needle_v} for {query} mentioned in the provided text?"""
Baber's avatar
nit  
Baber committed
12
eval_logger = logging.getLogger(__name__)
Baber's avatar
cleanup  
Baber committed
13
14
15
16
17
18
19
20
21
22


def download_dataset(df: Generator) -> dict[str, datasets.Dataset]:
    return {
        "test": datasets.Dataset.from_list(
            list(itertools.chain.from_iterable(df)), split=datasets.Split.TEST
        )
    }


Baber's avatar
Baber committed
23
24
25
26
27
28
29
30
31
32
def niah_single_1(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="repeat"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="repeat",
            type_needle_k="words",
            type_needle_v="numbers",
Baber's avatar
nit  
Baber committed
33
            num_samples=500,
Baber's avatar
Baber committed
34
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
35
36
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
37
    )
Baber's avatar
Baber committed
38
39
40
41
42
43
44
45
46
47
48
49


def niah_single_2(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="essay"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="essay",
            type_needle_k="words",
            type_needle_v="numbers",
Baber's avatar
nit  
Baber committed
50
            num_samples=500,
Baber's avatar
Baber committed
51
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
52
53
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
54
    )
Baber's avatar
Baber committed
55
56
57
58
59
60
61
62
63
64
65
66


def niah_single_3(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="essay"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="essay",
            type_needle_k="words",
            type_needle_v="uuids",
Baber's avatar
nit  
Baber committed
67
            num_samples=500,
Baber's avatar
Baber committed
68
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
69
70
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
71
    )
Baber's avatar
Baber committed
72
73
74
75
76
77
78
79
80
81
82
83
84


def niah_multikey_1(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="essay"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="essay",
            type_needle_k="words",
            type_needle_v="numbers",
            num_needle_k=4,
Baber's avatar
nit  
Baber committed
85
            num_samples=500,
Baber's avatar
Baber committed
86
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
87
88
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
89
    )
Baber's avatar
Baber committed
90
91
92
93
94
95
96
97
98
99
100
101


def niah_multikey_2(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="needle"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="needle",
            type_needle_k="words",
            type_needle_v="numbers",
Baber's avatar
nit  
Baber committed
102
            num_samples=500,
Baber's avatar
Baber committed
103
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
104
105
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
106
    )
Baber's avatar
Baber committed
107
108
109
110
111
112
113
114
115
116
117
118


def niah_multikey_3(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="needle"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="needle",
            type_needle_k="uuids",
            type_needle_v="uuids",
Baber's avatar
nit  
Baber committed
119
            num_samples=500,
Baber's avatar
Baber committed
120
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
121
122
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
123
    )
Baber's avatar
Baber committed
124
125
126
127
128
129
130
131
132
133
134
135
136


def niah_multivalue(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="essay"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="essay",
            type_needle_k="words",
            type_needle_v="numbers",
            num_needle_v=4,
Baber's avatar
nit  
Baber committed
137
            num_samples=500,
Baber's avatar
Baber committed
138
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
139
140
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
141
    )
Baber's avatar
Baber committed
142
143
144
145
146
147
148
149
150
151
152
153
154


def niah_multiquery(**kwargs):
    seq_lengths = kwargs.pop("max_seq_lengths", DEFAULT_SEQ_LENGTHS)
    return download_dataset(
        generate_samples(
            get_haystack(type_haystack="essay"),
            max_seq_length=seq,
            template=TEMPLATE,
            type_haystack="essay",
            type_needle_k="words",
            type_needle_v="numbers",
            num_needle_q=4,
Baber's avatar
nit  
Baber committed
155
            num_samples=500,
Baber's avatar
Baber committed
156
            TOKENIZER=get_tokenizer(**kwargs),
Baber's avatar
Baber committed
157
158
        )
        for seq in seq_lengths
Baber's avatar
cleanup  
Baber committed
159
    )