utils.py 5.55 KB
Newer Older
Baber's avatar
Baber committed
1
2
3
# noqa
import itertools
import re
Baber's avatar
Baber committed
4
from functools import cache
Baber's avatar
Baber committed
5
6
7
8
9
from typing import Literal

import datasets
from transformers import AutoTokenizer

Baber's avatar
Baber committed
10
from lm_eval.tasks.ruler.essays import get_all_essays
Baber's avatar
Baber committed
11
12
13
from lm_eval.tasks.ruler.prepare import generate_samples


Baber's avatar
Baber committed
14
@cache
Baber's avatar
Baber committed
15
def get_tokenizer(pretrained):
Baber's avatar
nit  
Baber committed
16
    print("using tokenizer ", pretrained)
Baber's avatar
Baber committed
17
    return AutoTokenizer.from_pretrained(pretrained, trust_remote_code=True)
Baber's avatar
Baber committed
18
19
20


# TOKENIZER = AutoTokenizer.from_pretrained(os.environ.get("TOKENIZER"))
Baber's avatar
Baber committed
21
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? The special magic {type_needle_v} for {query} mentioned in the provided text are"""
Baber's avatar
Baber committed
22
23

SEQ_LENGTHS = (
Baber's avatar
Baber committed
24
25
26
    # 131072,
    # 65536,
    # 32768,
Baber's avatar
Baber committed
27
28
    16384,
    8192,
Baber's avatar
Baber committed
29
30
31
32
33
34
35
36
37
    4096,
)

NUM_SAMPLES = 500
REMOVE_NEWLINE_TAB = ""
STOP_WORDS = ""
RANDOM_SEED = 42


Baber's avatar
Baber committed
38
@cache
Baber's avatar
Baber committed
39
40
41
def get_haystack(type_haystack: Literal["essay", "repeat", "needle"]):
    NEEDLE = "One of the special magic {type_needle_v} for {key} is: {value}."
    if type_haystack == "essay":
Baber's avatar
Baber committed
42
        essay = get_all_essays()["text"]
Baber's avatar
Baber committed
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
        # essay = json.load(open(essay))["text"]
        haystack = re.sub(r"\s+", " ", essay).split(" ")
    elif type_haystack == "repeat":
        haystack = "The grass is green. The sky is blue. The sun is yellow. Here we go. There and back again."
    elif type_haystack == "needle":
        haystack = NEEDLE
    else:
        raise NotImplementedError(f"{type_haystack} is not implemented.")
    return haystack


def flatten(df):
    return {
        "test": datasets.Dataset.from_list(
            list(itertools.chain.from_iterable(df)), split=datasets.Split.TEST
        )
    }


# ruff: noqa
Baber's avatar
Baber committed
63
niah_single_1 = lambda x: flatten(
Baber's avatar
Baber committed
64
65
66
67
68
69
70
    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
Baber committed
71
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
72
73
74
75
    )
    for seq in SEQ_LENGTHS
)
# ruff: noqa
Baber's avatar
Baber committed
76
niah_single_2 = lambda x: flatten(
Baber's avatar
Baber committed
77
78
79
80
81
82
83
    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
Baber committed
84
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
85
86
87
88
    )
    for seq in SEQ_LENGTHS
)
# noqa
Baber's avatar
Baber committed
89
niah_single_3 = lambda x: flatten(
Baber's avatar
Baber committed
90
91
92
93
94
95
96
    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
Baber committed
97
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
98
99
100
101
    )
    for seq in SEQ_LENGTHS
)
# noqa
Baber's avatar
Baber committed
102
niah_multikey_1 = lambda x: flatten(
Baber's avatar
Baber committed
103
104
105
106
107
108
109
110
    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
Baber committed
111
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
112
113
114
115
    )
    for seq in SEQ_LENGTHS
)
# noqa
Baber's avatar
Baber committed
116
niah_multikey_2 = lambda x: flatten(
Baber's avatar
Baber committed
117
118
119
120
121
122
123
    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
Baber committed
124
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
125
126
127
128
    )
    for seq in SEQ_LENGTHS
)
# noqa
Baber's avatar
Baber committed
129
niah_multikey_3 = lambda x: flatten(
Baber's avatar
Baber committed
130
131
132
133
134
135
136
    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
Baber committed
137
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
138
139
140
141
    )
    for seq in SEQ_LENGTHS
)
# noqa
Baber's avatar
Baber committed
142
niah_multivalue = lambda x: flatten(
Baber's avatar
Baber committed
143
144
145
146
147
148
149
150
    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
Baber committed
151
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
152
153
154
155
    )
    for seq in SEQ_LENGTHS
)
# noqa
Baber's avatar
Baber committed
156
niah_multiquery = lambda x: flatten(
Baber's avatar
Baber committed
157
158
159
160
161
162
163
164
    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
Baber committed
165
        TOKENIZER=get_tokenizer(x),
Baber's avatar
Baber committed
166
167
168
169
170
    )
    for seq in SEQ_LENGTHS
)


Baber's avatar
Baber committed
171
def postprocess_pred(predict_str: str) -> str:
Baber's avatar
Baber committed
172
173
174
175
176
177
178
179
180
    predict_str = predict_str.strip()

    # Remove all non-printable characters
    np_pattern = re.compile(r"[\x00-\x1f]")
    predict_str = np_pattern.sub("\n", predict_str).strip()

    return predict_str


Baber's avatar
Baber committed
181
182
183
184
185
186
187
188
189
190
def string_match_all(preds: list[str], refs: list[list[str]]) -> float:
    score = sum(
        [
            sum([1.0 if r.lower() in pred.lower() else 0.0 for r in ref]) / len(ref)
            for pred, ref in zip(preds, refs)
        ]
    ) / len(preds)
    return score


Baber's avatar
Baber committed
191
192
def process_results(doc: dict, results: list[str]) -> dict[str, float]:
    # hacky: set all other lengths to -1
Baber's avatar
Baber committed
193
194
    metrics = {str(length): -1.0 for length in SEQ_LENGTHS}
    input_len = doc["max_length"]
Baber's avatar
Baber committed
195
196
197
    pred = postprocess_pred(results[0])
    score = string_match_all([pred], [doc["outputs"]])
    metrics[str(input_len)] = score
Baber's avatar
Baber committed
198
199
200
    return metrics


Baber's avatar
Baber committed
201
def aggregate_metrics(metrics: list[float]) -> float:
Baber's avatar
Baber committed
202
203
204
205
206
    res = [x for x in metrics if x != -1]
    if not res:
        # we don't have any samples with this length
        return 0.0
    return sum(res) / len(res)