utils.py 11.8 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
def hash_string(string: str) -> str:
    return hashlib.sha256(string.encode("utf-8")).hexdigest()


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)


haileyschoelkopf's avatar
haileyschoelkopf committed
57
58
59
60
61
def handle_arg_string(arg):
    if arg.lower() == "true":
        return True
    elif arg.lower() == "false":
        return False
62
63
64
65
66
67
    elif arg.isnumeric():
        return int(arg)
    try:
        return float(arg)
    except ValueError:
        return arg
haileyschoelkopf's avatar
haileyschoelkopf committed
68
69


70
71
72
73
74
75
76
77
78
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
79
80
81
82
83
84
def simple_parse_args_string(args_string):
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Jason Phang's avatar
Jason Phang committed
85
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
86
87
    if not args_string:
        return {}
88
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
89
90
91
    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
92
    return args_dict
Leo Gao's avatar
Leo Gao committed
93

Fabrizio Milo's avatar
Fabrizio Milo committed
94

Leo Gao's avatar
Leo Gao committed
95
96
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
97
        yield from iter
Leo Gao's avatar
Leo Gao committed
98
99


100
101
102
103
104
def group(arr, fn):
    res = collections.defaultdict(list)

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

106
107
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
108

gakada's avatar
gakada committed
109
110
111
# Returns a list containing all values of the source_list that
# match at least one of the patterns
def pattern_match(patterns, source_list):
112
    if isinstance(patterns, str):
113
114
        patterns = [patterns]

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


Jason Phang's avatar
Jason Phang committed
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
163
164
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
165
    yield ([prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len])
Jason Phang's avatar
Jason Phang committed
166
167
168
169
170
    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
171

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

Fabrizio Milo's avatar
Fabrizio Milo committed
178

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

Jason Phang's avatar
Jason Phang committed
184

185
186
187
188
189
190
191
192
193
194
195
196
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)


197
class Reorderer:
baberabb's avatar
baberabb committed
198
199
200
201
202
203
204
    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
        """
205
206
207
        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
    def get_reordered(self):
baberabb's avatar
baberabb committed
216
217
218
219
220
        """Gets the reordered array

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

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

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
236
            for ind in inds:
237
238
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
239

240
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
241

242
243
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
244

245
def make_table(result_dict, column: str = "results", sort_results: bool = True):
246
    """Generate table of results."""
247
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
248

lintangsutawika's avatar
lintangsutawika committed
249
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
250
251
252
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
253

lintangsutawika's avatar
lintangsutawika committed
254
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
255
        column_name,
lintangsutawika's avatar
lintangsutawika committed
256
257
        "Version",
        "Filter",
258
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
259
260
261
262
263
        "Metric",
        "Value",
        "",
        "Stderr",
    ]
264

lintangsutawika's avatar
lintangsutawika committed
265
266
267
268
269
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

270
271
    values = []

272
273
274
275
276
277
    keys = result_dict[column].keys()
    if sort_results:
        # sort entries alphabetically
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
278
        version = result_dict["versions"].get(k, "N/A")
279
        n = str(result_dict["n-shot"][k])
280
281
282
283

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

284
285
        for (mf), v in dic.items():
            m, _, f = mf.partition(",")
286
287
288
            if m.endswith("_stderr"):
                continue

289
290
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
291
292
293
                if se != "N/A":
                    se = "%.4f" % se
                values.append([k, version, f, n, m, "%.4f" % v, "±", se])
294
            else:
295
                values.append([k, version, f, n, m, "%.4f" % v, "", ""])
296
297
298
299
300
301
302
303
304
305
306
            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()


307
308
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
309
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
310
311
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
312

313
314
    @functools.wraps(fn)
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
315
316
317
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
318
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
319
320
                "lm-evaluation-harness!"
            )
321
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
322

323
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
324

Fabrizio Milo's avatar
Fabrizio Milo committed
325

326
327
328
329
def ignore_constructor(loader, node):
    return node


lintangsutawika's avatar
lintangsutawika committed
330
331
332
333
def import_function(loader, node):
    function_name = loader.construct_scalar(node)
    yaml_path = os.path.dirname(loader.name)

lintangsutawika's avatar
lintangsutawika committed
334
    *module_name, function_name = function_name.split(".")
335
    if isinstance(module_name, list):
lintangsutawika's avatar
lintangsutawika committed
336
337
        module_name = ".".join(module_name)
    module_path = os.path.normpath(os.path.join(yaml_path, "{}.py".format(module_name)))
lintangsutawika's avatar
lintangsutawika committed
338
339
340
341
342
343
344
345

    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
346

347
348
349
350
351
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
352

353
354
    # Add the import_function constructor to the YAML loader
    yaml.add_constructor("!function", constructor_fn)
355
356
357
    if yaml_config is None:
        with open(yaml_path, "rb") as file:
            yaml_config = yaml.full_load(file)
lintangsutawika's avatar
lintangsutawika committed
358

lintangsutawika's avatar
lintangsutawika committed
359
360
    if yaml_dir is None:
        yaml_dir = os.path.dirname(yaml_path)
361
362
363
364
365
366
367

    assert yaml_dir is not None

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

368
        if isinstance(include_path, str):
369
370
371
372
373
374
375
376
377
378
379
380
381
            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:
382
                included_yaml_config = load_yaml_config(yaml_path=path, mode=mode)
383
384
385
386
387
388
389
390
                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
391
392


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

lintangsutawika's avatar
lintangsutawika committed
397

398
env = Environment(loader=BaseLoader, undefined=StrictUndefined)
399
env.filters["regex_replace"] = regex_replace
400
401


baberabb's avatar
baberabb committed
402
def apply_template(template: str, doc: dict) -> str:
403
404
    rtemplate = env.from_string(template)
    return rtemplate.render(**doc)
405
406


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