utils.py 5.56 KB
Newer Older
Baber's avatar
Baber committed
1
2
3
4
5
# noqa
import itertools
import json
import os
import re
Baber's avatar
Baber committed
6
from functools import partial, cache
Baber's avatar
Baber committed
7
8
9
10
11
from typing import Literal

import datasets
from transformers import AutoTokenizer

Baber's avatar
Baber committed
12
from lm_eval.tasks.ruler.essays import get_essays, get_all_essays
Baber's avatar
Baber committed
13
14
15
from lm_eval.tasks.ruler.prepare import generate_samples


Baber's avatar
Baber committed
16
@cache
Baber's avatar
Baber committed
17
18
def get_tokenizer(pretrained):
    return AutoTokenizer.from_pretrained(pretrained, trust_remote_code=True)
Baber's avatar
Baber committed
19
20
21


# TOKENIZER = AutoTokenizer.from_pretrained(os.environ.get("TOKENIZER"))
Baber's avatar
Baber committed
22
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
23
24

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

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


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


Baber's avatar
Baber committed
172
def postprocess_pred(predict_str: str) -> str:
Baber's avatar
Baber committed
173
174
175
176
177
178
179
180
181
    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
182
183
184
185
186
187
188
189
190
191
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
192
193
def process_results(doc: dict, results: list[str]) -> dict[str, float]:
    # hacky: set all other lengths to -1
Baber's avatar
Baber committed
194
195
    metrics = {str(length): -1.0 for length in SEQ_LENGTHS}
    input_len = doc["max_length"]
Baber's avatar
Baber committed
196
197
198
    pred = postprocess_pred(results[0])
    score = string_match_all([pred], [doc["outputs"]])
    metrics[str(input_len)] = score
Baber's avatar
Baber committed
199
200
201
    return metrics


Baber's avatar
Baber committed
202
def aggregate_metrics(metrics: list[float]) -> float:
Baber's avatar
Baber committed
203
204
205
206
207
    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)