utils.py 14.5 KB
Newer Older
1
2
3
import collections
import fnmatch
import functools
4
import hashlib
5
import importlib.util
6
import inspect
7
import json
8
9
10
import logging
import os
import re
11
from dataclasses import asdict, is_dataclass
12
from itertools import islice
13
from typing import Any, Callable, Generator, List, Tuple
14

Lintang Sutawika's avatar
Lintang Sutawika committed
15
import numpy as np
16
import yaml
17
from jinja2 import BaseLoader, Environment, StrictUndefined
sdtblck's avatar
sdtblck committed
18

lintangsutawika's avatar
lintangsutawika committed
19

20
21
22
23
24
logging.basicConfig(
    format="%(asctime)s,%(msecs)03d %(levelname)-8s [%(filename)s:%(lineno)d] %(message)s",
    datefmt="%Y-%m-%d:%H:%M:%S",
    level=logging.INFO,
)
25
eval_logger = logging.getLogger("lm-eval")
sdtblck's avatar
sdtblck committed
26

27
SPACING = " " * 47
sdtblck's avatar
sdtblck committed
28

29
30
31
32
33
HIGHER_IS_BETTER_SYMBOLS = {
    True: "↑",
    False: "↓",
}

sdtblck's avatar
sdtblck committed
34

35
36
37
38
def hash_string(string: str) -> str:
    return hashlib.sha256(string.encode("utf-8")).hexdigest()


39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
def escaped_split(text, sep_char, maxsplit=-1):
    """Split text into a list on occurrences of the given separation
    character `sep_char`. The separation character may be escaped by a
    backslash to avoid splitting at that location.

    The separation character must be a string of size 1.

    If `maxsplit` is given, at most `maxsplit` splits are done (thus,
    the list will have at most `maxsplit + 1` elements). If `maxsplit`
    is not specified or less than 0, then there is no limit on the
    number of splits (all possible splits are made).
    """
    assert (
        len(sep_char) == 1
    ), "separation string must be a single character for escaped splitting"

    if maxsplit == 0:
        return text
    maxsplit = max(0, maxsplit)

    return re.split(r"(?<!\\)" + sep_char, text, maxsplit)


haileyschoelkopf's avatar
haileyschoelkopf committed
62
63
64
65
66
def handle_arg_string(arg):
    if arg.lower() == "true":
        return True
    elif arg.lower() == "false":
        return False
67
68
69
70
71
72
    elif arg.isnumeric():
        return int(arg)
    try:
        return float(arg)
    except ValueError:
        return arg
haileyschoelkopf's avatar
haileyschoelkopf committed
73
74


75
76
77
78
79
80
81
82
83
def handle_non_serializable(o):
    if isinstance(o, np.int64) or isinstance(o, np.int32):
        return int(o)
    elif isinstance(o, set):
        return list(o)
    else:
        return str(o)


84
85
86
87
88
89
90
91
92
93
94
95
def sanitize_list(sub):
    """
    Takes possible nested list and recursively converts all inner component to strings
    """
    if isinstance(sub, list):
        return [sanitize_list(item) for item in sub]
    if isinstance(sub, tuple):
        return tuple(sanitize_list(item) for item in sub)
    else:
        return str(sub)


Jason Phang's avatar
gpt3  
Jason Phang committed
96
97
98
99
100
101
def simple_parse_args_string(args_string):
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Jason Phang's avatar
Jason Phang committed
102
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
103
104
    if not args_string:
        return {}
105
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
106
107
108
    args_dict = {
        k: handle_arg_string(v) for k, v in [arg.split("=") for arg in arg_list]
    }
Jason Phang's avatar
gpt3  
Jason Phang committed
109
    return args_dict
Leo Gao's avatar
Leo Gao committed
110

Fabrizio Milo's avatar
Fabrizio Milo committed
111

Leo Gao's avatar
Leo Gao committed
112
113
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
114
        yield from iter
Leo Gao's avatar
Leo Gao committed
115
116


117
118
119
120
121
def group(arr, fn):
    res = collections.defaultdict(list)

    for ob in arr:
        res[fn(ob)].append(ob)
Fabrizio Milo's avatar
Fabrizio Milo committed
122

123
124
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
125

gakada's avatar
gakada committed
126
127
128
# Returns a list containing all values of the source_list that
# match at least one of the patterns
def pattern_match(patterns, source_list):
129
    if isinstance(patterns, str):
130
131
        patterns = [patterns]

gakada's avatar
gakada committed
132
133
134
135
136
137
138
    task_names = set()
    for pattern in patterns:
        for matching in fnmatch.filter(source_list, pattern):
            task_names.add(matching)
    return sorted(list(task_names))


Lintang Sutawika's avatar
Lintang Sutawika committed
139
140
141
142
143
144
def softmax(x):
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


Leo Gao's avatar
Leo Gao committed
145
146
147
148
def general_detokenize(string):
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
149
150
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
151
    string = re.sub(r" (['.,])", r"\1", string)
152
153
154
    return string


155
156
157
158
159
160
161
162
163
164
165
def get_file_task_name(filename: str) -> str:
    """
    Given the sample results filenames, extracts and returns the task name.
    """
    return filename[filename.find("_") + 1 : filename.rfind("_")]


def get_file_datetime(filename: str) -> str:
    """
    Given the results and sample results filenames, extracts and returns the datetime.
    """
166
    return filename[filename.rfind("_") + 1 :].replace(".jsonl", "")
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203


def sanitize_model_name(model_name: str) -> str:
    """
    Given the model name, returns a sanitized version of it.
    """
    return re.sub(r"[\"<>:/\|\\?\*\[\]]+", "__", model_name)


def sanitize_task_name(task_name: str) -> str:
    """
    Given the task name, returns a sanitized version of it.
    """
    return re.sub(r"\W", "_", task_name)


def get_latest_filename(filenames: List[str]) -> str:
    """
    Given a list of filenames, returns the filename with the latest datetime.
    """
    return max(filenames, key=lambda f: get_file_datetime(f))


def get_results_filenames(filenames: List[str]) -> List[str]:
    """
    Extracts filenames that correspond to aggregated results.
    """
    return [f for f in filenames if "/results_" in f and ".json" in f]


def get_sample_results_filenames(filenames: List[str]) -> List[str]:
    """
    Extracts filenames that correspond to sample results.
    """
    return [f for f in filenames if "/samples_" in f and ".json" in f]


204
205
206
def get_rolling_token_windows(
    token_list: List[int], prefix_token: int, max_seq_len: int, context_len: int
) -> Generator[Tuple[List[int], List[int]], None, None]:
Jason Phang's avatar
Jason Phang committed
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
    """
    - context_len allows for a rolling window context, allowing each prediction window to potentially
      condition on some context

    :param token_list: list
        List of tokens to be PREDICTED
    :param max_seq_len: int
        max_seq_len of model (or max_seq_len we want to use)
    :param context_len: int
        Amount of desired token context for prediction. Needs to be at least 1.
    :param prefix_token: token
        Dummy token like <eos> so the first token has something to condition on
    :return: generator
        Generator of tuples
            (input_tokens, pred_tokens)
        Note: Score only the last len(pred_tokens) logits of the LM
    """
    assert 1 <= context_len <= max_seq_len
    if not token_list:
        return
    # +1 offset, going from input->preds
    pred_len = max_seq_len - context_len + 1
    predicted = 0

    # Special handling for first window: predict all tokens
    first_seq_len = min(max_seq_len, len(token_list))
233
    yield [prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len]
Jason Phang's avatar
Jason Phang committed
234
235
236
237
238
    predicted += first_seq_len

    while predicted < len(token_list):
        window_pred_len = min(len(token_list) - predicted, pred_len)
        window_end = predicted + window_pred_len
Leo Gao's avatar
Leo Gao committed
239

Jason Phang's avatar
Jason Phang committed
240
        yield (
lintangsutawika's avatar
lintangsutawika committed
241
242
            token_list[window_end - max_seq_len - 1 : window_end - 1],
            token_list[window_end - window_pred_len : window_end],
Jason Phang's avatar
Jason Phang committed
243
244
245
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
246

247
248
249
def make_disjoint_window(
    pair: Tuple[List[int], List[int]],
) -> Tuple[List[int], List[int]]:
Fabrizio Milo's avatar
Fabrizio Milo committed
250
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
251
    a, b = pair
252
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
253

Jason Phang's avatar
Jason Phang committed
254

255
256
257
258
259
260
261
262
263
264
265
266
class EnhancedJSONEncoder(json.JSONEncoder):
    """
    Provides a proper json encoding for the loggers and trackers json dumps.
    Notably manages the json encoding of dataclasses.
    """

    def default(self, o):
        if is_dataclass(o):
            return asdict(o)
        return super().default(o)


267
class Reorderer:
baberabb's avatar
baberabb committed
268
269
270
271
272
273
274
    def __init__(self, arr: List[Any], fn: Callable) -> None:
        """Reorder an array according to some function

        Args:
            arr (List[Any]): The initial array
            fn (Callable[[Any], Any]): A function to determine the priority of elements
        """
275
276
277
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
278
279
280
        # arr = [([y[0] for y in x], x[0][1]) for x in arr]
        # TODO: overhaul reorderer. It currently grouped requests by content but we don't want this
        arr = [([y[0]], x[0][1]) for x in arr for y in x]
281
282
283
        arr.sort(key=lambda x: fn(x[1]))

        self.arr = arr
Fabrizio Milo's avatar
Fabrizio Milo committed
284

285
    def get_reordered(self):
baberabb's avatar
baberabb committed
286
287
288
289
290
        """Gets the reordered array

        Returns:
            List[Any]: The reordered array
        """
291
        return [x[1] for x in self.arr]
Fabrizio Milo's avatar
Fabrizio Milo committed
292

293
    def get_original(self, newarr):
baberabb's avatar
baberabb committed
294
295
296
297
298
299
300
301
        """Restores the original order of a new array based on the old array's order

        Args:
            newarr (List[Any]): The array to be restored

        Returns:
            List[Any]: The array restored to the original order
        """
302
303
304
305
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
306
            for ind in inds:
307
308
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
309

310
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
311

312
313
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
314

Lintang Sutawika's avatar
Lintang Sutawika committed
315
def make_table(result_dict, column: str = "results", sort_results: bool = False):
316
    """Generate table of results."""
317
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
318

lintangsutawika's avatar
lintangsutawika committed
319
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
320
321
322
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
323

lintangsutawika's avatar
lintangsutawika committed
324
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
325
        column_name,
lintangsutawika's avatar
lintangsutawika committed
326
327
        "Version",
        "Filter",
328
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
329
        "Metric",
330
        "",
lintangsutawika's avatar
lintangsutawika committed
331
332
333
334
        "Value",
        "",
        "Stderr",
    ]
335

lintangsutawika's avatar
lintangsutawika committed
336
337
338
339
340
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

341
342
    values = []

343
344
    keys = result_dict[column].keys()
    if sort_results:
Lintang Sutawika's avatar
Lintang Sutawika committed
345
346
347
        # sort entries alphabetically by task or group name.
        # NOTE: we default here to false, because order matters for multi-level table printing a la mmlu.
        # sorting here would mess that up
348
349
350
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
Lintang Sutawika's avatar
Lintang Sutawika committed
351
352
        version = result_dict["versions"].get(k, "    N/A")
        n = str(result_dict.get("n-shot", " ").get(k, " "))
353
        higher_is_better = result_dict.get("higher_is_better", {}).get(k, {})
354
355
356
357

        if "alias" in dic:
            k = dic.pop("alias")

358
        metric_items = dic.items()
Lintang Sutawika's avatar
Lintang Sutawika committed
359
        metric_items = sorted(metric_items)
360
361

        for (mf), v in metric_items:
362
            m, _, f = mf.partition(",")
363
364
365
            if m.endswith("_stderr"):
                continue

366
367
            hib = HIGHER_IS_BETTER_SYMBOLS.get(higher_is_better.get(m), "")

Lintang Sutawika's avatar
Lintang Sutawika committed
368
369
            v = "%.4f" % v if isinstance(v, float) else v

370
371
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
Lintang Sutawika's avatar
Lintang Sutawika committed
372
                se = "   N/A" if se == "N/A" else "%.4f" % se
Lintang Sutawika's avatar
Lintang Sutawika committed
373
                values.append([k, version, f, n, m, hib, v, "±", se])
374
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
375
                values.append([k, version, f, n, m, hib, v, "", ""])
376
377
378
379
380
381
382
383
384
385
386
            k = ""
            version = ""
    md_writer.value_matrix = values
    latex_writer.value_matrix = values

    # todo: make latex table look good
    # print(latex_writer.dumps())

    return md_writer.dumps()


387
388
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
389
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
390
391
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
392

393
394
    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
395
396
397
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
398
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
399
400
                "lm-evaluation-harness!"
            )
401
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
402

403
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
404

Fabrizio Milo's avatar
Fabrizio Milo committed
405

406
407
408
409
def ignore_constructor(loader, node):
    return node


lintangsutawika's avatar
lintangsutawika committed
410
411
412
413
def import_function(loader, node):
    function_name = loader.construct_scalar(node)
    yaml_path = os.path.dirname(loader.name)

lintangsutawika's avatar
lintangsutawika committed
414
    *module_name, function_name = function_name.split(".")
415
    if isinstance(module_name, list):
lintangsutawika's avatar
lintangsutawika committed
416
417
        module_name = ".".join(module_name)
    module_path = os.path.normpath(os.path.join(yaml_path, "{}.py".format(module_name)))
lintangsutawika's avatar
lintangsutawika committed
418
419
420
421
422
423
424
425

    spec = importlib.util.spec_from_file_location(module_name, module_path)
    module = importlib.util.module_from_spec(spec)
    spec.loader.exec_module(module)

    function = getattr(module, function_name)
    return function

lintangsutawika's avatar
lintangsutawika committed
426

427
428
429
430
431
def load_yaml_config(yaml_path=None, yaml_config=None, yaml_dir=None, mode="full"):
    if mode == "simple":
        constructor_fn = ignore_constructor
    elif mode == "full":
        constructor_fn = import_function
lintangsutawika's avatar
lintangsutawika committed
432

433
434
    # Add the import_function constructor to the YAML loader
    yaml.add_constructor("!function", constructor_fn)
435
436
437
    if yaml_config is None:
        with open(yaml_path, "rb") as file:
            yaml_config = yaml.full_load(file)
lintangsutawika's avatar
lintangsutawika committed
438

lintangsutawika's avatar
lintangsutawika committed
439
440
    if yaml_dir is None:
        yaml_dir = os.path.dirname(yaml_path)
441
442
443
444
445
446
447

    assert yaml_dir is not None

    if "include" in yaml_config:
        include_path = yaml_config["include"]
        del yaml_config["include"]

448
        if isinstance(include_path, str):
449
450
451
452
453
454
455
456
457
458
459
460
461
            include_path = [include_path]

        # Load from the last one first
        include_path.reverse()
        final_yaml_config = {}
        for path in include_path:
            # Assumes that path is a full path.
            # If not found, assume the included yaml
            # is in the same dir as the original yaml
            if not os.path.isfile(path):
                path = os.path.join(yaml_dir, path)

            try:
462
                included_yaml_config = load_yaml_config(yaml_path=path, mode=mode)
463
464
465
466
467
468
469
470
                final_yaml_config.update(included_yaml_config)
            except Exception as ex:
                # If failed to load, ignore
                raise ex

        final_yaml_config.update(yaml_config)
        return final_yaml_config
    return yaml_config
lintangsutawika's avatar
lintangsutawika committed
471
472


Ethan Smith's avatar
Ethan Smith committed
473
def regex_replace(string, pattern, repl, count: int = 0):
474
475
    """Implements the `re.sub` function as a custom Jinja filter."""
    return re.sub(pattern, repl, string, count=count)
lintangsutawika's avatar
lintangsutawika committed
476

lintangsutawika's avatar
lintangsutawika committed
477

478
479
480
env = Environment(
    loader=BaseLoader, undefined=StrictUndefined, keep_trailing_newline=True
)
481
env.filters["regex_replace"] = regex_replace
482
483


baberabb's avatar
baberabb committed
484
def apply_template(template: str, doc: dict) -> str:
485
486
    rtemplate = env.from_string(template)
    return rtemplate.render(**doc)
487
488


489
def create_iterator(raw_iterator, *, rank=0, world_size=1, limit=None):
490
491
492
    """
    Method for creating a (potentially) sliced and limited
    iterator from a raw document iterator. Used for splitting data
493
494
495
    among ranks in multigpu setting or only pulling a sample of documents
    """
    return islice(raw_iterator, rank, limit, world_size)
496
497
498
499
500
501
502
503
504
505


def weighted_f1_score(items):
    from sklearn.metrics import f1_score

    unzipped_list = list(zip(*items))
    golds = unzipped_list[0]
    preds = unzipped_list[1]
    fscore = f1_score(golds, preds, average="weighted")
    return fscore