utils.py 15.3 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
Baber's avatar
Baber committed
13
from typing import Any, Callable, Generator, List, Optional, 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
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).
    """
Baber Abbasi's avatar
Baber Abbasi committed
51
52
53
    assert len(sep_char) == 1, (
        "separation string must be a single character for escaped splitting"
    )
54
55
56
57
58
59
60
61

    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)


Baber's avatar
Baber committed
96
def simple_parse_args_string(args_string: Optional[str]) -> dict:
Jason Phang's avatar
gpt3  
Jason Phang committed
97
98
99
100
101
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Baber's avatar
Baber committed
102
103
    if args_string is None:
        return {}
Jason Phang's avatar
Jason Phang committed
104
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
105
106
    if not args_string:
        return {}
107
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
108
    args_dict = {
109
110
        kv[0]: handle_arg_string("=".join(kv[1:]))
        for kv in [arg.split("=") for arg in arg_list]
haileyschoelkopf's avatar
haileyschoelkopf committed
111
    }
Jason Phang's avatar
gpt3  
Jason Phang committed
112
    return args_dict
Leo Gao's avatar
Leo Gao committed
113

Fabrizio Milo's avatar
Fabrizio Milo committed
114

Baber's avatar
Baber committed
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
def parse_keyed_list_string(s: str) -> dict[str, list]:
    """Parse a string of key-value pairs into a dictionary where all values are lists."""
    result = {}
    current_key = None
    values = []

    parts = s.split(",")

    for part in parts:
        if "=" in part:
            # Save previous key's values
            if current_key is not None:
                result[current_key] = values

            # Start new key-value pair
            current_key, value = part.split("=")
            values = [handle_arg_string(value)]
        else:
            values.append(handle_arg_string(part))

    # Add the last key-value pair
    if current_key is not None:
        result[current_key] = values

    return result


Leo Gao's avatar
Leo Gao committed
142
143
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
144
        yield from iter
Leo Gao's avatar
Leo Gao committed
145
146


147
148
149
150
151
def group(arr, fn):
    res = collections.defaultdict(list)

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

153
154
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
155

gakada's avatar
gakada committed
156
157
158
# Returns a list containing all values of the source_list that
# match at least one of the patterns
def pattern_match(patterns, source_list):
159
    if isinstance(patterns, str):
160
161
        patterns = [patterns]

gakada's avatar
gakada committed
162
163
164
165
166
167
168
    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
169
170
171
172
173
174
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
175
176
177
178
def general_detokenize(string):
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
179
180
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
181
    string = re.sub(r" (['.,])", r"\1", string)
182
183
184
    return string


185
186
187
188
189
190
191
192
193
194
195
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.
    """
196
    return filename[filename.rfind("_") + 1 :].replace(".jsonl", "")
197
198
199
200
201
202
203
204
205
206
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
233


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]


234
235
236
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
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
    """
    - 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))
263
    yield [prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len]
Jason Phang's avatar
Jason Phang committed
264
265
266
267
268
    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
269

Jason Phang's avatar
Jason Phang committed
270
        yield (
lintangsutawika's avatar
lintangsutawika committed
271
272
            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
273
274
275
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
276

277
278
279
def make_disjoint_window(
    pair: Tuple[List[int], List[int]],
) -> Tuple[List[int], List[int]]:
Fabrizio Milo's avatar
Fabrizio Milo committed
280
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
281
    a, b = pair
282
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
283

Jason Phang's avatar
Jason Phang committed
284

285
286
287
288
289
290
291
292
293
294
295
296
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)


297
class Reorderer:
baberabb's avatar
baberabb committed
298
299
300
301
302
303
304
    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
        """
305
306
307
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
308
309
310
        # 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]
311
312
313
        arr.sort(key=lambda x: fn(x[1]))

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

315
    def get_reordered(self):
baberabb's avatar
baberabb committed
316
317
318
319
320
        """Gets the reordered array

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

323
    def get_original(self, newarr):
baberabb's avatar
baberabb committed
324
325
326
327
328
329
330
331
        """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
        """
332
333
334
335
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
336
            for ind in inds:
337
338
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
339

340
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
341

342
343
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
344

Lintang Sutawika's avatar
Lintang Sutawika committed
345
def make_table(result_dict, column: str = "results", sort_results: bool = False):
346
    """Generate table of results."""
347
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
348

lintangsutawika's avatar
lintangsutawika committed
349
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
350
351
352
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
353

lintangsutawika's avatar
lintangsutawika committed
354
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
355
        column_name,
lintangsutawika's avatar
lintangsutawika committed
356
357
        "Version",
        "Filter",
358
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
359
        "Metric",
360
        "",
lintangsutawika's avatar
lintangsutawika committed
361
362
363
364
        "Value",
        "",
        "Stderr",
    ]
365

lintangsutawika's avatar
lintangsutawika committed
366
367
368
369
370
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

371
372
    values = []

373
374
    keys = result_dict[column].keys()
    if sort_results:
Lintang Sutawika's avatar
Lintang Sutawika committed
375
376
377
        # 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
378
379
380
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
Lintang Sutawika's avatar
Lintang Sutawika committed
381
382
        version = result_dict["versions"].get(k, "    N/A")
        n = str(result_dict.get("n-shot", " ").get(k, " "))
383
        higher_is_better = result_dict.get("higher_is_better", {}).get(k, {})
384
385
386
387

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

388
        metric_items = dic.items()
Lintang Sutawika's avatar
Lintang Sutawika committed
389
        metric_items = sorted(metric_items)
390
391

        for (mf), v in metric_items:
392
            m, _, f = mf.partition(",")
393
394
395
            if m.endswith("_stderr"):
                continue

396
397
            hib = HIGHER_IS_BETTER_SYMBOLS.get(higher_is_better.get(m), "")

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

400
401
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
Lintang Sutawika's avatar
Lintang Sutawika committed
402
                se = "   N/A" if se == "N/A" else "%.4f" % se
Lintang Sutawika's avatar
Lintang Sutawika committed
403
                values.append([k, version, f, n, m, hib, v, "±", se])
404
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
405
                values.append([k, version, f, n, m, hib, v, "", ""])
406
407
408
409
410
411
412
413
414
415
416
            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()


417
418
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
419
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
420
421
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
422

423
424
    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
425
426
427
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
428
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
429
430
                "lm-evaluation-harness!"
            )
431
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
432

433
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
434

Fabrizio Milo's avatar
Fabrizio Milo committed
435

436
437
438
439
def ignore_constructor(loader, node):
    return node


lintangsutawika's avatar
lintangsutawika committed
440
441
442
443
def import_function(loader, node):
    function_name = loader.construct_scalar(node)
    yaml_path = os.path.dirname(loader.name)

lintangsutawika's avatar
lintangsutawika committed
444
    *module_name, function_name = function_name.split(".")
445
    if isinstance(module_name, list):
lintangsutawika's avatar
lintangsutawika committed
446
447
        module_name = ".".join(module_name)
    module_path = os.path.normpath(os.path.join(yaml_path, "{}.py".format(module_name)))
lintangsutawika's avatar
lintangsutawika committed
448
449
450
451
452
453
454
455

    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
456

457
458
459
460
461
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
462

463
464
    # Add the import_function constructor to the YAML loader
    yaml.add_constructor("!function", constructor_fn)
465
466
467
    if yaml_config is None:
        with open(yaml_path, "rb") as file:
            yaml_config = yaml.full_load(file)
lintangsutawika's avatar
lintangsutawika committed
468

lintangsutawika's avatar
lintangsutawika committed
469
470
    if yaml_dir is None:
        yaml_dir = os.path.dirname(yaml_path)
471
472
473
474
475
476
477

    assert yaml_dir is not None

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

478
        if isinstance(include_path, str):
479
480
481
482
483
484
485
486
487
488
489
490
491
            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:
492
                included_yaml_config = load_yaml_config(yaml_path=path, mode=mode)
493
494
495
496
497
498
499
500
                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
501
502


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

lintangsutawika's avatar
lintangsutawika committed
507

508
509
510
env = Environment(
    loader=BaseLoader, undefined=StrictUndefined, keep_trailing_newline=True
)
511
env.filters["regex_replace"] = regex_replace
512
513


baberabb's avatar
baberabb committed
514
def apply_template(template: str, doc: dict) -> str:
515
516
    rtemplate = env.from_string(template)
    return rtemplate.render(**doc)
517
518


519
def create_iterator(raw_iterator, *, rank=0, world_size=1, limit=None):
520
521
522
    """
    Method for creating a (potentially) sliced and limited
    iterator from a raw document iterator. Used for splitting data
523
524
525
    among ranks in multigpu setting or only pulling a sample of documents
    """
    return islice(raw_iterator, rank, limit, world_size)
526
527
528
529
530
531
532
533
534
535


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