utils.py 15.7 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 pathlib import Path
14
from typing import Any, Callable, Generator, List, Tuple
15

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

lintangsutawika's avatar
lintangsutawika committed
20

21
SPACING = " " * 47
sdtblck's avatar
sdtblck committed
22

23
24
25
26
27
HIGHER_IS_BETTER_SYMBOLS = {
    True: "↑",
    False: "↓",
}

sdtblck's avatar
sdtblck committed
28

Lintang Sutawika's avatar
Lintang Sutawika committed
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
def setup_logging(verbosity=logging.INFO):
    # Configure the root logger
    log_level = os.environ.get("LOGLEVEL", verbosity) or verbosity

    level_map = {
        "DEBUG": logging.DEBUG,
        "INFO": logging.INFO,
        "WARNING": logging.WARNING,
        "ERROR": logging.ERROR,
        "CRITICAL": logging.CRITICAL,
    }

    log_level = level_map.get(str(log_level).upper(), logging.INFO)
    if not logging.root.handlers:
        logging.basicConfig(
            format="%(asctime)s,%(msecs)03d %(levelname)-8s [%(name)s:%(lineno)d] %(message)s",
            datefmt="%Y-%m-%d:%H:%M:%S",
            level=log_level,
        )
        if log_level == logging.DEBUG:
            third_party_loggers = ["urllib3", "filelock", "fsspec"]
            for logger_name in third_party_loggers:
                logging.getLogger(logger_name).setLevel(logging.INFO)
    else:
        logging.getLogger().setLevel(log_level)


56
57
58
59
def hash_string(string: str) -> str:
    return hashlib.sha256(string.encode("utf-8")).hexdigest()


60
61
62
63
64
65
66
67
68
69
70
71
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
72
73
74
    assert len(sep_char) == 1, (
        "separation string must be a single character for escaped splitting"
    )
75
76
77
78
79
80
81
82

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

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


haileyschoelkopf's avatar
haileyschoelkopf committed
83
84
85
86
87
def handle_arg_string(arg):
    if arg.lower() == "true":
        return True
    elif arg.lower() == "false":
        return False
88
89
90
91
92
93
    elif arg.isnumeric():
        return int(arg)
    try:
        return float(arg)
    except ValueError:
        return arg
haileyschoelkopf's avatar
haileyschoelkopf committed
94
95


96
97
98
99
100
101
102
103
104
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)


105
106
107
108
109
110
111
112
113
114
115
116
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
117
118
119
120
121
122
def simple_parse_args_string(args_string):
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Jason Phang's avatar
Jason Phang committed
123
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
124
125
    if not args_string:
        return {}
126
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
127
    args_dict = {
128
129
        kv[0]: handle_arg_string("=".join(kv[1:]))
        for kv in [arg.split("=") for arg in arg_list]
haileyschoelkopf's avatar
haileyschoelkopf committed
130
    }
Jason Phang's avatar
gpt3  
Jason Phang committed
131
    return args_dict
Leo Gao's avatar
Leo Gao committed
132

Fabrizio Milo's avatar
Fabrizio Milo committed
133

Leo Gao's avatar
Leo Gao committed
134
135
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
136
        yield from iter
Leo Gao's avatar
Leo Gao committed
137
138


139
140
141
142
143
def group(arr, fn):
    res = collections.defaultdict(list)

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

145
146
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
147

gakada's avatar
gakada committed
148
149
150
# Returns a list containing all values of the source_list that
# match at least one of the patterns
def pattern_match(patterns, source_list):
151
    if isinstance(patterns, str):
152
153
        patterns = [patterns]

gakada's avatar
gakada committed
154
155
156
157
158
159
160
    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
161
162
163
164
165
166
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
167
168
169
170
def general_detokenize(string):
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
171
172
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
173
    string = re.sub(r" (['.,])", r"\1", string)
174
175
176
    return string


177
178
179
180
181
182
183
184
185
186
187
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.
    """
188
    return filename[filename.rfind("_") + 1 :].replace(".jsonl", "")
189
190
191
192
193
194
195
196
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


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]


226
227
228
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
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
    """
    - 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))
255
    yield [prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len]
Jason Phang's avatar
Jason Phang committed
256
257
258
259
260
    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
261

Jason Phang's avatar
Jason Phang committed
262
        yield (
lintangsutawika's avatar
lintangsutawika committed
263
264
            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
265
266
267
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
268

269
270
271
def make_disjoint_window(
    pair: Tuple[List[int], List[int]],
) -> Tuple[List[int], List[int]]:
Fabrizio Milo's avatar
Fabrizio Milo committed
272
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
273
    a, b = pair
274
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
275

Jason Phang's avatar
Jason Phang committed
276

277
278
279
280
281
282
283
284
285
286
287
288
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)


289
class Reorderer:
baberabb's avatar
baberabb committed
290
291
292
293
294
295
296
    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
        """
297
298
299
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
300
301
302
        # 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]
303
304
305
        arr.sort(key=lambda x: fn(x[1]))

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

307
    def get_reordered(self):
baberabb's avatar
baberabb committed
308
309
310
311
312
        """Gets the reordered array

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

315
    def get_original(self, newarr):
baberabb's avatar
baberabb committed
316
317
318
319
320
321
322
323
        """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
        """
324
325
326
327
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
328
            for ind in inds:
329
330
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
331

332
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
333

334
335
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
336

Lintang Sutawika's avatar
Lintang Sutawika committed
337
def make_table(result_dict, column: str = "results", sort_results: bool = False):
338
    """Generate table of results."""
339
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
340

lintangsutawika's avatar
lintangsutawika committed
341
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
342
343
344
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
345

lintangsutawika's avatar
lintangsutawika committed
346
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
347
        column_name,
lintangsutawika's avatar
lintangsutawika committed
348
349
        "Version",
        "Filter",
350
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
351
        "Metric",
352
        "",
lintangsutawika's avatar
lintangsutawika committed
353
354
355
356
        "Value",
        "",
        "Stderr",
    ]
357

lintangsutawika's avatar
lintangsutawika committed
358
359
360
361
362
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

363
364
    values = []

365
366
    keys = result_dict[column].keys()
    if sort_results:
Lintang Sutawika's avatar
Lintang Sutawika committed
367
368
369
        # 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
370
371
372
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
Lintang Sutawika's avatar
Lintang Sutawika committed
373
374
        version = result_dict["versions"].get(k, "    N/A")
        n = str(result_dict.get("n-shot", " ").get(k, " "))
375
        higher_is_better = result_dict.get("higher_is_better", {}).get(k, {})
376
377
378
379

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

380
        metric_items = dic.items()
Lintang Sutawika's avatar
Lintang Sutawika committed
381
        metric_items = sorted(metric_items)
382
383

        for (mf), v in metric_items:
384
            m, _, f = mf.partition(",")
385
386
387
            if m.endswith("_stderr"):
                continue

388
389
            hib = HIGHER_IS_BETTER_SYMBOLS.get(higher_is_better.get(m), "")

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

392
393
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
Lintang Sutawika's avatar
Lintang Sutawika committed
394
                se = "   N/A" if se == "N/A" else "%.4f" % se
Lintang Sutawika's avatar
Lintang Sutawika committed
395
                values.append([k, version, f, n, m, hib, v, "±", se])
396
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
397
                values.append([k, version, f, n, m, hib, v, "", ""])
398
399
400
401
402
403
404
405
406
407
408
            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()


409
410
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
411
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
412
413
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
414

415
416
    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
417
418
419
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
420
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
421
422
                "lm-evaluation-harness!"
            )
423
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
424

425
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
426

Fabrizio Milo's avatar
Fabrizio Milo committed
427

428
429
430
431
def ignore_constructor(loader, node):
    return node


432
def import_function(loader: yaml.Loader, node, yaml_path: Path):
lintangsutawika's avatar
lintangsutawika committed
433
434
    function_name = loader.construct_scalar(node)

lintangsutawika's avatar
lintangsutawika committed
435
    *module_name, function_name = function_name.split(".")
436
    if isinstance(module_name, list):
lintangsutawika's avatar
lintangsutawika committed
437
        module_name = ".".join(module_name)
438
    module_path = yaml_path.parent / f"{module_name}.py"
lintangsutawika's avatar
lintangsutawika committed
439

440
441
442
443
    spec = importlib.util.spec_from_file_location(module_name, module_path.as_posix())

    if spec is None:
        raise ImportError(f"Could not import module {module_name} from {module_path}.")
lintangsutawika's avatar
lintangsutawika committed
444
    module = importlib.util.module_from_spec(spec)
445
446
447

    if spec.loader is None:
        raise ImportError(f"Module loader is None, {module_name} from {module_path}.")
lintangsutawika's avatar
lintangsutawika committed
448
449
450
451
452
    spec.loader.exec_module(module)

    function = getattr(module, function_name)
    return function

lintangsutawika's avatar
lintangsutawika committed
453

454
455
456
457
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":
458
459
460
461
        if yaml_path is None:
            raise ValueError("yaml_path must be provided if mode is 'full'.")
        # Attach yaml_path to the import function so that it can be used later
        constructor_fn = functools.partial(import_function, yaml_path=Path(yaml_path))
lintangsutawika's avatar
lintangsutawika committed
462

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

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

    assert yaml_dir is not None

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

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


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

lintangsutawika's avatar
lintangsutawika committed
508

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


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


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


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