utils.py 12.3 KB
Newer Older
sdtblck's avatar
sdtblck committed
1
import os
Leo Gao's avatar
Leo Gao committed
2
import re
Stephen Hogg's avatar
Stephen Hogg committed
3
import sys
4
5
6
7
8
9
import yaml
import inspect
import pathlib
import functools
import subprocess
import collections
lintangsutawika's avatar
lintangsutawika committed
10
import importlib.util
gakada's avatar
gakada committed
11
import fnmatch
12

13
14
from typing import List, Union

15
import gc
16
import torch
sdtblck's avatar
sdtblck committed
17

Xingjian Shi's avatar
Xingjian Shi committed
18
from omegaconf import OmegaConf
19
from jinja2 import BaseLoader, Environment, StrictUndefined
20
from itertools import islice
sdtblck's avatar
sdtblck committed
21

22
from lm_eval.logger import eval_logger
sdtblck's avatar
sdtblck committed
23
24
25
26
27
28
29
30
31
32
33


class ExitCodeError(Exception):
    pass


def sh(x):
    if os.system(x):
        raise ExitCodeError()


34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
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)


Jason Phang's avatar
gpt3  
Jason Phang committed
57
58
59
60
61
62
def simple_parse_args_string(args_string):
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Jason Phang's avatar
Jason Phang committed
63
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
64
65
66
    if not args_string:
        return {}
    arg_list = args_string.split(",")
Xingjian Shi's avatar
Xingjian Shi committed
67
    args_dict = OmegaConf.to_object(OmegaConf.from_dotlist(arg_list))
Jason Phang's avatar
gpt3  
Jason Phang committed
68
    return args_dict
Leo Gao's avatar
Leo Gao committed
69

Fabrizio Milo's avatar
Fabrizio Milo committed
70

Leo Gao's avatar
Leo Gao committed
71
72
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
73
        yield from iter
Leo Gao's avatar
Leo Gao committed
74
75


76
def chunks(iter, n=0, fn=None):
Leo Gao's avatar
Leo Gao committed
77
    arr = []
78
    for i, x in enumerate(iter):
Leo Gao's avatar
Leo Gao committed
79
        arr.append(x)
80
        if len(arr) == (fn(i) if fn else n):
Leo Gao's avatar
Leo Gao committed
81
82
            yield arr
            arr = []
Fabrizio Milo's avatar
Fabrizio Milo committed
83
84
85
86

    if arr:
        yield arr

Leo Gao's avatar
Leo Gao committed
87

88
89
90
91
92
def group(arr, fn):
    res = collections.defaultdict(list)

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

94
95
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
96

gakada's avatar
gakada committed
97
98
99
100
101
102
103
class MultiChoice:
    def __init__(self, choices):
        self.choices = choices

    # Simple wildcard support (linux filename patterns)
    def __contains__(self, values):
        for value in values.split(","):
104
105
106
107
108
            if len(fnmatch.filter(self.choices, value)) == 0:
                eval_logger.warning("{} is not in task list.".format(value))
                eval_logger.info(f"Available tasks to choose:")
                for choice in self.choices:
                    eval_logger.info(f"  - {choice}")
gakada's avatar
gakada committed
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
        return True

    def __iter__(self):
        for choice in self.choices:
            yield choice


# Returns a list containing all values of the source_list that
# match at least one of the patterns
def pattern_match(patterns, source_list):
    task_names = set()
    for pattern in patterns:
        for matching in fnmatch.filter(source_list, pattern):
            task_names.add(matching)
    return sorted(list(task_names))


Leo Gao's avatar
Leo Gao committed
126
127
128
129
def general_detokenize(string):
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
130
131
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
132
    string = re.sub(r" (['.,])", r"\1", string)
133
134
135
    return string


Jason Phang's avatar
Jason Phang committed
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
def get_rolling_token_windows(token_list, prefix_token, max_seq_len, context_len):
    """
    - 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))
Fabrizio Milo's avatar
Fabrizio Milo committed
163
    yield ([prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len])
Jason Phang's avatar
Jason Phang committed
164
165
166
167
168
    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
169

Jason Phang's avatar
Jason Phang committed
170
        yield (
171
172
            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
173
174
175
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
176

Leo Gao's avatar
Leo Gao committed
177
def make_disjoint_window(pair):
Fabrizio Milo's avatar
Fabrizio Milo committed
178
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
179
    a, b = pair
180
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
181

Jason Phang's avatar
Jason Phang committed
182

183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
def select_continuation_from_batch_left_padding(
    generations: Union[List[List[int]], torch.Tensor], max_context_size: int
):
    """Select the continuation from the batch, removing prompts of different lengths.
    Args:
        generations (Union[List[List[int]], torch.Tensor]):
            A tensor or list-of-lists of shape [batch_size, sequence length].
        max_context_size (int):
            The size of the biggest context; generations will proceed from that
            index.
    Example:
        PAD     PAD Continue : The dog chased the cat  [every       day of the week]
        Riddle  me    this   : The  dog chased the  cat [yesterday] PAD PAD PAD PAD
    Output:
        [every day of the week]
        [yesterday]  PAD PAD PAD PAD
    """
    return generations[:, max_context_size:]


203
204
205
206
207
class Reorderer:
    def __init__(self, arr, fn):
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
208
209
210
        # 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]
211
212
213
        arr.sort(key=lambda x: fn(x[1]))

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

215
216
    def get_reordered(self):
        return [x[1] for x in self.arr]
Fabrizio Milo's avatar
Fabrizio Milo committed
217

218
219
220
221
222
    def get_original(self, newarr):
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
223
            for ind in inds:
224
225
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
226

227
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
228

229
230
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
231

232
233
234
235
236
237
def make_table(result_dict):
    """Generate table of results."""
    from pytablewriter import MarkdownTableWriter, LatexTableWriter

    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
238
    md_writer.headers = ["Task", "Version", "Filter", "Metric", "Value", "", "Stderr"]
lintangsutawika's avatar
lintangsutawika committed
239
240
241
242
243
244
245
246
247
    latex_writer.headers = [
        "Task",
        "Version",
        "Filter",
        "Metric",
        "Value",
        "",
        "Stderr",
    ]
248
249
250
251
252

    values = []

    for k, dic in result_dict["results"].items():
        version = result_dict["versions"][k]
253
254
        for (mf), v in dic.items():
            m, _, f = mf.partition(",")
255
256
257
            if m.endswith("_stderr"):
                continue

258
259
260
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
                values.append([k, version, f, m, "%.4f" % v, "±", "%.4f" % se])
261
            else:
262
                values.append([k, version, f, m, "%.4f" % v, "", ""])
263
264
265
266
267
268
269
270
271
272
273
            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()


274
275
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
276
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
277
278
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
279

280
281
    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
282
283
284
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
285
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
286
287
                "lm-evaluation-harness!"
            )
288
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
289

290
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
291

Fabrizio Milo's avatar
Fabrizio Milo committed
292

Stephen Hogg's avatar
Stephen Hogg committed
293
294
295
296
297
298
299
300
301
@positional_deprecated
def find_test_root(start_path: pathlib.Path) -> pathlib.Path:
    """
    Search upward in the directory tree to a maximum of three layers
    to find and return the package root (containing the 'tests' folder)
    """
    cur_path = start_path.resolve()
    max_layers = 3
    for _ in range(max_layers):
Fabrizio Milo's avatar
Fabrizio Milo committed
302
        if (cur_path / "tests" / "test_version_stable.py").exists():
Stephen Hogg's avatar
Stephen Hogg committed
303
304
305
            return cur_path
        else:
            cur_path = cur_path.parent.resolve()
Fabrizio Milo's avatar
Fabrizio Milo committed
306
307
308
309
    raise FileNotFoundError(
        f"Unable to find package root within {max_layers} upwards" + f"of {start_path}"
    )

Stephen Hogg's avatar
Stephen Hogg committed
310
311

@positional_deprecated
312
def run_task_tests(task_list: List[str]):
Stephen Hogg's avatar
Stephen Hogg committed
313
314
315
    """
    Find the package root and run the tests for the given tasks
    """
jon-tow's avatar
jon-tow committed
316
317
    import pytest

318
    package_root = find_test_root(start_path=pathlib.Path(__file__))
Fabrizio Milo's avatar
Fabrizio Milo committed
319
320
321
322
323
324
325
    task_string = " or ".join(task_list)
    args = [
        f"{package_root}/tests/test_version_stable.py",
        f"--rootdir={package_root}",
        "-k",
        f"{task_string}",
    ]
Stephen Hogg's avatar
Stephen Hogg committed
326
327
328
    sys.path.append(str(package_root))
    pytest_return_val = pytest.main(args)
    if pytest_return_val:
Fabrizio Milo's avatar
Fabrizio Milo committed
329
330
331
        raise ValueError(
            f"Not all tests for the specified tasks ({task_list}) ran successfully! Error code: {pytest_return_val}"
        )
332
333


334
335
336
337
338
339
340
341
342
343
344
345
346
def get_git_commit_hash():
    """
    Gets the git commit hash of your current repo (if it exists).
    Source: https://github.com/EleutherAI/gpt-neox/blob/b608043be541602170bfcfb8ec9bf85e8a0799e0/megatron/neox_arguments/neox_args.py#L42
    """
    try:
        git_hash = subprocess.check_output(["git", "describe", "--always"]).strip()
        git_hash = git_hash.decode()
    except subprocess.CalledProcessError:
        git_hash = None
    return git_hash


lintangsutawika's avatar
lintangsutawika committed
347
348
349
350
351
def import_function(loader, node):

    function_name = loader.construct_scalar(node)
    yaml_path = os.path.dirname(loader.name)

lintangsutawika's avatar
lintangsutawika committed
352
    module_name, function_name = function_name.split(".")
lintangsutawika's avatar
lintangsutawika committed
353
354
355
356
357
358
359
360
361
    module_path = os.path.join(yaml_path, "{}.py".format(module_name))

    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
362

lintangsutawika's avatar
lintangsutawika committed
363
# Add the import_function constructor to the YAML loader
lintangsutawika's avatar
lintangsutawika committed
364
yaml.add_constructor("!function", import_function)
lintangsutawika's avatar
lintangsutawika committed
365
366
367


def load_yaml_config(yaml_path):
lintangsutawika's avatar
lintangsutawika committed
368
    with open(yaml_path, "rb") as file:
lintangsutawika's avatar
lintangsutawika committed
369
370
        yaml_config = yaml.full_load(file)
        yaml_dir = os.path.dirname(yaml_path)
lintangsutawika's avatar
lintangsutawika committed
371
372
373
374

        if "include" in yaml_config:
            include_path = yaml_config["include"]
            del yaml_config["include"]
lintangsutawika's avatar
lintangsutawika committed
375
376
377

            if type(include_path) == str:
                include_path = [include_path]
lintangsutawika's avatar
lintangsutawika committed
378

lintangsutawika's avatar
lintangsutawika committed
379
380
381
382
383
384
            # Load from the last one first
            include_path.reverse()
            final_yaml_config = {}
            for path in include_path:

                # Assumes that path is a full path.
lintangsutawika's avatar
lintangsutawika committed
385
                # If not found, assume the included yaml
lintangsutawika's avatar
lintangsutawika committed
386
387
388
389
390
391
392
                # is in the same dir as the original yaml
                if not os.path.isfile(path):
                    path = os.path.join(yaml_dir, path)

                try:
                    included_yaml_config = load_yaml_config(path)
                    final_yaml_config.update(included_yaml_config)
lintangsutawika's avatar
lintangsutawika committed
393
                except Exception as ex:
lintangsutawika's avatar
lintangsutawika committed
394
                    # If failed to load, ignore
lintangsutawika's avatar
lintangsutawika committed
395
                    raise ex
lintangsutawika's avatar
lintangsutawika committed
396
397
398
399
400
401

            final_yaml_config.update(yaml_config)
            return final_yaml_config
        return yaml_config


402
env = Environment(loader=BaseLoader, undefined=StrictUndefined)
403
404
405
406
407


def apply_template(template, doc):
    rtemplate = env.from_string(template)
    return rtemplate.render(**doc)
408
409


410
411
412
413
def create_iterator(raw_iterator, rank, world_size, limit=None):
    """
    Method for creating a (potentially) sliced and limited
    iterator from a raw document iterator. Used for splitting data
414
415
416
    among ranks in multigpu setting or only pulling a sample of documents
    """
    return islice(raw_iterator, rank, limit, world_size)
417
418


419
420
421
def clear_torch_cache():
    gc.collect()
    torch.cuda.empty_cache()