utils.py 12 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, List
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)


Jason Phang's avatar
gpt3  
Jason Phang committed
84
85
86
87
88
89
def simple_parse_args_string(args_string):
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Jason Phang's avatar
Jason Phang committed
90
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
91
92
    if not args_string:
        return {}
93
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
94
95
96
    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
97
    return args_dict
Leo Gao's avatar
Leo Gao committed
98

Fabrizio Milo's avatar
Fabrizio Milo committed
99

Leo Gao's avatar
Leo Gao committed
100
101
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
102
        yield from iter
Leo Gao's avatar
Leo Gao committed
103
104


105
106
107
108
109
def group(arr, fn):
    res = collections.defaultdict(list)

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

111
112
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
113

gakada's avatar
gakada committed
114
115
116
# Returns a list containing all values of the source_list that
# match at least one of the patterns
def pattern_match(patterns, source_list):
117
    if isinstance(patterns, str):
118
119
        patterns = [patterns]

gakada's avatar
gakada committed
120
121
122
123
124
125
126
    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
127
128
129
130
131
132
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
133
134
135
136
def general_detokenize(string):
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
137
138
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
139
    string = re.sub(r" (['.,])", r"\1", string)
140
141
142
    return string


Jason Phang's avatar
Jason Phang committed
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
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
170
    yield ([prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len])
Jason Phang's avatar
Jason Phang committed
171
172
173
174
175
    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
176

Jason Phang's avatar
Jason Phang committed
177
        yield (
lintangsutawika's avatar
lintangsutawika committed
178
179
            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
180
181
182
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
183

Leo Gao's avatar
Leo Gao committed
184
def make_disjoint_window(pair):
Fabrizio Milo's avatar
Fabrizio Milo committed
185
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
186
    a, b = pair
187
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
188

Jason Phang's avatar
Jason Phang committed
189

190
191
192
193
194
195
196
197
198
199
200
201
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)


202
class Reorderer:
baberabb's avatar
baberabb committed
203
204
205
206
207
208
209
    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
        """
210
211
212
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
213
214
215
        # 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]
216
217
218
        arr.sort(key=lambda x: fn(x[1]))

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

220
    def get_reordered(self):
baberabb's avatar
baberabb committed
221
222
223
224
225
        """Gets the reordered array

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

228
    def get_original(self, newarr):
baberabb's avatar
baberabb committed
229
230
231
232
233
234
235
236
        """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
        """
237
238
239
240
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
241
            for ind in inds:
242
243
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
244

245
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
246

247
248
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
249

250
def make_table(result_dict, column: str = "results", sort_results: bool = True):
251
    """Generate table of results."""
252
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
253

lintangsutawika's avatar
lintangsutawika committed
254
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
255
256
257
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
258

lintangsutawika's avatar
lintangsutawika committed
259
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
260
        column_name,
lintangsutawika's avatar
lintangsutawika committed
261
262
        "Version",
        "Filter",
263
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
264
        "Metric",
265
        "",
lintangsutawika's avatar
lintangsutawika committed
266
267
268
269
        "Value",
        "",
        "Stderr",
    ]
270

lintangsutawika's avatar
lintangsutawika committed
271
272
273
274
275
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

276
277
    values = []

278
279
280
281
282
283
    keys = result_dict[column].keys()
    if sort_results:
        # sort entries alphabetically
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
284
        version = result_dict["versions"].get(k, "N/A")
285
        n = str(result_dict["n-shot"][k])
286
        higher_is_better = result_dict.get("higher_is_better", {}).get(k, {})
287
288
289
290

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

291
292
        for (mf), v in dic.items():
            m, _, f = mf.partition(",")
293
294
295
            if m.endswith("_stderr"):
                continue

296
297
            hib = HIGHER_IS_BETTER_SYMBOLS.get(higher_is_better.get(m), "")

298
299
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
300
301
                if se != "N/A":
                    se = "%.4f" % se
302
                values.append([k, version, f, n, m, hib, "%.4f" % v, "±", se])
303
            else:
304
                values.append([k, version, f, n, m, hib, "%.4f" % v, "", ""])
305
306
307
308
309
310
311
312
313
314
315
            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()


316
317
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
318
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
319
320
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
321

322
323
    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
324
325
326
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
327
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
328
329
                "lm-evaluation-harness!"
            )
330
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
331

332
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
333

Fabrizio Milo's avatar
Fabrizio Milo committed
334

335
336
337
338
def ignore_constructor(loader, node):
    return node


lintangsutawika's avatar
lintangsutawika committed
339
340
341
342
def import_function(loader, node):
    function_name = loader.construct_scalar(node)
    yaml_path = os.path.dirname(loader.name)

lintangsutawika's avatar
lintangsutawika committed
343
    *module_name, function_name = function_name.split(".")
344
    if isinstance(module_name, list):
lintangsutawika's avatar
lintangsutawika committed
345
346
        module_name = ".".join(module_name)
    module_path = os.path.normpath(os.path.join(yaml_path, "{}.py".format(module_name)))
lintangsutawika's avatar
lintangsutawika committed
347
348
349
350
351
352
353
354

    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
355

356
357
358
359
360
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
361

362
363
    # Add the import_function constructor to the YAML loader
    yaml.add_constructor("!function", constructor_fn)
364
365
366
    if yaml_config is None:
        with open(yaml_path, "rb") as file:
            yaml_config = yaml.full_load(file)
lintangsutawika's avatar
lintangsutawika committed
367

lintangsutawika's avatar
lintangsutawika committed
368
369
    if yaml_dir is None:
        yaml_dir = os.path.dirname(yaml_path)
370
371
372
373
374
375
376

    assert yaml_dir is not None

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

377
        if isinstance(include_path, str):
378
379
380
381
382
383
384
385
386
387
388
389
390
            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:
391
                included_yaml_config = load_yaml_config(yaml_path=path, mode=mode)
392
393
394
395
396
397
398
399
                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
400
401


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

lintangsutawika's avatar
lintangsutawika committed
406

407
env = Environment(loader=BaseLoader, undefined=StrictUndefined)
408
env.filters["regex_replace"] = regex_replace
409
410


baberabb's avatar
baberabb committed
411
def apply_template(template: str, doc: dict) -> str:
412
413
    rtemplate = env.from_string(template)
    return rtemplate.render(**doc)
414
415


416
def create_iterator(raw_iterator, *, rank=0, world_size=1, limit=None):
417
418
419
    """
    Method for creating a (potentially) sliced and limited
    iterator from a raw document iterator. Used for splitting data
420
421
422
    among ranks in multigpu setting or only pulling a sample of documents
    """
    return islice(raw_iterator, rank, limit, world_size)