utils.py 18.7 KB
Newer Older
Baber's avatar
Baber committed
1
2
from __future__ import annotations

3
4
import collections
import fnmatch
5
import hashlib
6
import importlib.util
7
import inspect
8
import json
9
10
11
import logging
import os
import re
Baber's avatar
Baber committed
12
from collections.abc import Generator
13
from dataclasses import asdict, is_dataclass
Baber's avatar
Baber committed
14
from functools import lru_cache, partial, wraps
15
from itertools import islice
16
from pathlib import Path
Baber's avatar
Baber committed
17
from typing import Any, Callable
18

Lintang Sutawika's avatar
Lintang Sutawika committed
19
import numpy as np
20
import yaml
Baber's avatar
Baber committed
21
from jinja2 import BaseLoader, Environment, StrictUndefined, Template
sdtblck's avatar
sdtblck committed
22

lintangsutawika's avatar
lintangsutawika committed
23

24
SPACING = " " * 47
sdtblck's avatar
sdtblck committed
25

26
27
28
29
30
HIGHER_IS_BETTER_SYMBOLS = {
    True: "↑",
    False: "↓",
}

sdtblck's avatar
sdtblck committed
31

Baber's avatar
Baber committed
32
def wrap_text(string: str, width: int = 140, **kwargs) -> str | None:
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
    """
    Wraps the given string to the specified width.
    """
    import textwrap

    return textwrap.fill(
        inspect.cleandoc(string),
        width=width,
        initial_indent="",
        subsequent_indent=" " * 8,
        break_long_words=False,
        break_on_hyphens=False,
        **kwargs,
    )


Lintang Sutawika's avatar
Lintang Sutawika committed
49
50
def setup_logging(verbosity=logging.INFO):
    # Configure the root logger
Baber Abbasi's avatar
Baber Abbasi committed
51
52
    class CustomFormatter(logging.Formatter):
        def format(self, record):
Baber's avatar
Baber committed
53
            record.name = record.name.removeprefix("im_eval.")
Baber Abbasi's avatar
Baber Abbasi committed
54
55
56
57
58
59
60
            return super().format(record)

    formatter = CustomFormatter(
        "%(asctime)s %(levelname)-8s [%(name)s:%(lineno)d] %(message)s",
        datefmt="%Y-%m-%d:%H:%M:%S",
    )

Lintang Sutawika's avatar
Lintang Sutawika committed
61
62
63
64
65
66
67
68
69
70
71
    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)
Baber Abbasi's avatar
Baber Abbasi committed
72

Lintang Sutawika's avatar
Lintang Sutawika committed
73
    if not logging.root.handlers:
Baber Abbasi's avatar
Baber Abbasi committed
74
75
76
77
78
79
80
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)

        root_logger = logging.getLogger()
        root_logger.addHandler(handler)
        root_logger.setLevel(log_level)

Lintang Sutawika's avatar
Lintang Sutawika committed
81
82
83
84
85
86
87
88
        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)


89
90
91
92
def hash_string(string: str) -> str:
    return hashlib.sha256(string.encode("utf-8")).hexdigest()


93
94
95
96
97
98
99
100
101
102
103
104
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
105
106
107
    assert len(sep_char) == 1, (
        "separation string must be a single character for escaped splitting"
    )
108
109
110
111
112

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

Baber's avatar
Baber committed
113
    return re.split(r"(?<!\\)" + sep_char, text, maxsplit=maxsplit)
114
115


haileyschoelkopf's avatar
haileyschoelkopf committed
116
117
118
119
120
def handle_arg_string(arg):
    if arg.lower() == "true":
        return True
    elif arg.lower() == "false":
        return False
121
122
123
124
125
126
    elif arg.isnumeric():
        return int(arg)
    try:
        return float(arg)
    except ValueError:
        return arg
haileyschoelkopf's avatar
haileyschoelkopf committed
127
128


129
def handle_non_serializable(o):
Baber's avatar
Baber committed
130
    if isinstance(o, np.integer):
131
132
133
134
135
136
137
        return int(o)
    elif isinstance(o, set):
        return list(o)
    else:
        return str(o)


138
139
140
141
142
143
144
145
146
147
148
149
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
150
def simple_parse_args_string(args_string: str | None) -> dict:
Jason Phang's avatar
gpt3  
Jason Phang committed
151
152
153
154
155
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Baber Abbasi's avatar
Baber Abbasi committed
156
157
    if args_string is None:
        return {}
Jason Phang's avatar
Jason Phang committed
158
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
159
160
    if not args_string:
        return {}
161
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
162
    args_dict = {
163
164
        kv[0]: handle_arg_string("=".join(kv[1:]))
        for kv in [arg.split("=") for arg in arg_list]
haileyschoelkopf's avatar
haileyschoelkopf committed
165
    }
Jason Phang's avatar
gpt3  
Jason Phang committed
166
    return args_dict
Leo Gao's avatar
Leo Gao committed
167

Fabrizio Milo's avatar
Fabrizio Milo committed
168

Leo Gao's avatar
Leo Gao committed
169
170
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
171
        yield from iter
Leo Gao's avatar
Leo Gao committed
172
173


174
175
176
177
178
def group(arr, fn):
    res = collections.defaultdict(list)

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

180
181
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
182

gakada's avatar
gakada committed
183
184
# Returns a list containing all values of the source_list that
# match at least one of the patterns
Baber's avatar
Baber committed
185
def pattern_match(patterns: list[str], source_list: list[str]) -> list[str]:
186
    if isinstance(patterns, str):
187
188
        patterns = [patterns]

gakada's avatar
gakada committed
189
190
191
192
193
194
195
    task_names = set()
    for pattern in patterns:
        for matching in fnmatch.filter(source_list, pattern):
            task_names.add(matching)
    return sorted(list(task_names))


Baber Abbasi's avatar
Baber Abbasi committed
196
def softmax(x) -> np.ndarray:
Lintang Sutawika's avatar
Lintang Sutawika committed
197
198
199
200
201
    """Compute softmax values for each sets of scores in x."""
    e_x = np.exp(x - np.max(x))
    return e_x / e_x.sum()


Baber's avatar
Baber committed
202
def general_detokenize(string: str) -> str:
Leo Gao's avatar
Leo Gao committed
203
204
205
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
206
207
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
208
    string = re.sub(r" (['.,])", r"\1", string)
209
210
211
    return string


212
213
214
215
216
217
218
219
220
221
222
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.
    """
223
    return filename[filename.rfind("_") + 1 :].replace(".jsonl", "")
224
225
226
227
228
229


def sanitize_model_name(model_name: str) -> str:
    """
    Given the model name, returns a sanitized version of it.
    """
Baber's avatar
Baber committed
230
    return re.sub(r"[\"<>:/|\\?*\[\]]+", "__", model_name)
231
232
233
234
235
236
237
238
239


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)


Baber's avatar
Baber committed
240
def get_latest_filename(filenames: list[str]) -> str:
241
242
243
244
245
246
    """
    Given a list of filenames, returns the filename with the latest datetime.
    """
    return max(filenames, key=lambda f: get_file_datetime(f))


Baber's avatar
Baber committed
247
def get_results_filenames(filenames: list[str]) -> list[str]:
248
249
250
251
252
253
    """
    Extracts filenames that correspond to aggregated results.
    """
    return [f for f in filenames if "/results_" in f and ".json" in f]


Baber's avatar
Baber committed
254
def get_sample_results_filenames(filenames: list[str]) -> list[str]:
255
256
257
258
259
260
    """
    Extracts filenames that correspond to sample results.
    """
    return [f for f in filenames if "/samples_" in f and ".json" in f]


261
def get_rolling_token_windows(
Baber's avatar
Baber committed
262
263
    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
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
    """
    - 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))
290
    yield [prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len]
Jason Phang's avatar
Jason Phang committed
291
292
293
294
295
    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
296

Jason Phang's avatar
Jason Phang committed
297
        yield (
lintangsutawika's avatar
lintangsutawika committed
298
299
            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
300
301
302
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
303

304
def make_disjoint_window(
Baber's avatar
Baber committed
305
306
    pair: tuple[list[int], list[int]],
) -> tuple[list[int], list[int]]:
Fabrizio Milo's avatar
Fabrizio Milo committed
307
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
308
    a, b = pair
309
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
310

Jason Phang's avatar
Jason Phang committed
311

312
313
314
315
316
317
318
319
320
321
322
323
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)


324
class Reorderer:
Baber's avatar
Baber committed
325
    def __init__(self, arr: list[Any], fn: Callable) -> None:
baberabb's avatar
baberabb committed
326
327
328
329
330
331
        """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
        """
332
333
334
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
335
336
337
        # 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]
338
339
340
        arr.sort(key=lambda x: fn(x[1]))

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

342
    def get_reordered(self):
baberabb's avatar
baberabb committed
343
344
345
346
347
        """Gets the reordered array

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

350
    def get_original(self, newarr):
baberabb's avatar
baberabb committed
351
352
353
354
355
356
357
358
        """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
        """
359
360
361
362
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
363
            for ind in inds:
364
365
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
366

367
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
368

369
370
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
371

Lintang Sutawika's avatar
Lintang Sutawika committed
372
def make_table(result_dict, column: str = "results", sort_results: bool = False):
373
    """Generate table of results."""
374
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
375

lintangsutawika's avatar
lintangsutawika committed
376
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
377
378
379
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
380

lintangsutawika's avatar
lintangsutawika committed
381
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
382
        column_name,
lintangsutawika's avatar
lintangsutawika committed
383
384
        "Version",
        "Filter",
385
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
386
        "Metric",
387
        "",
lintangsutawika's avatar
lintangsutawika committed
388
389
390
391
        "Value",
        "",
        "Stderr",
    ]
392

lintangsutawika's avatar
lintangsutawika committed
393
394
395
396
397
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

398
399
    values = []

400
401
    keys = result_dict[column].keys()
    if sort_results:
Lintang Sutawika's avatar
Lintang Sutawika committed
402
403
404
        # 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
405
406
407
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
Lintang Sutawika's avatar
Lintang Sutawika committed
408
409
        version = result_dict["versions"].get(k, "    N/A")
        n = str(result_dict.get("n-shot", " ").get(k, " "))
Baber's avatar
Baber committed
410
411
        # TODO: fix this
        # higher_is_better = result_dict.get("higher_is_better", {}).get(k, {})
412
413
414
415

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

416
        metric_items = dic.items()
Lintang Sutawika's avatar
Lintang Sutawika committed
417
        metric_items = sorted(metric_items)
418
419

        for (mf), v in metric_items:
420
            m, _, f = mf.partition(",")
421
422
423
            if m.endswith("_stderr"):
                continue

Baber's avatar
Baber committed
424
425
426
            # hib = HIGHER_IS_BETTER_SYMBOLS.get(higher_is_better.get(m), "")
            # TODO: fix
            hib = "↑"
427

Baber's avatar
Baber committed
428
            v = f"{v:.4f}" if isinstance(v, float) else v
Lintang Sutawika's avatar
Lintang Sutawika committed
429

430
431
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
Baber's avatar
Baber committed
432
                se = "   N/A" if se == "N/A" else f"{se:.4f}"
Lintang Sutawika's avatar
Lintang Sutawika committed
433
                values.append([k, version, f, n, m, hib, v, "±", se])
434
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
435
                values.append([k, version, f, n, m, hib, v, "", ""])
436
437
438
439
440
441
442
443
444
445
446
            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()


447
448
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
449
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
450
451
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
452

Baber's avatar
Baber committed
453
454
    wraps(fn)

455
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
456
457
458
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
459
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
460
461
                "lm-evaluation-harness!"
            )
462
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
463

464
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
465

Fabrizio Milo's avatar
Fabrizio Milo committed
466

467
468
469
470
def ignore_constructor(loader, node):
    return node


471
def import_function(loader: yaml.Loader, node, yaml_path: Path):
lintangsutawika's avatar
lintangsutawika committed
472
473
    function_name = loader.construct_scalar(node)

lintangsutawika's avatar
lintangsutawika committed
474
    *module_name, function_name = function_name.split(".")
475
    if isinstance(module_name, list):
lintangsutawika's avatar
lintangsutawika committed
476
        module_name = ".".join(module_name)
477
    module_path = yaml_path.parent / f"{module_name}.py"
lintangsutawika's avatar
lintangsutawika committed
478

479
480
481
482
    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
483
    module = importlib.util.module_from_spec(spec)
484
485
486

    if spec.loader is None:
        raise ImportError(f"Module loader is None, {module_name} from {module_path}.")
lintangsutawika's avatar
lintangsutawika committed
487
488
489
490
491
    spec.loader.exec_module(module)

    function = getattr(module, function_name)
    return function

lintangsutawika's avatar
lintangsutawika committed
492

Baber's avatar
Baber committed
493
494
495
def load_yaml_config(
    yaml_path: str | None = None, yaml_config=None, yaml_dir=None, mode="full"
):
496
497
498
    if mode == "simple":
        constructor_fn = ignore_constructor
    elif mode == "full":
499
500
501
        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
Baber's avatar
Baber committed
502
        constructor_fn = partial(import_function, yaml_path=Path(yaml_path))
lintangsutawika's avatar
lintangsutawika committed
503

504
    loader = yaml.CLoader if yaml.__with_libyaml__ else yaml.FullLoader
505
    # Add the import_function constructor to the YAML loader
506
    yaml.add_constructor("!function", constructor_fn, Loader=loader)
507
508
    if yaml_config is None:
        with open(yaml_path, "rb") as file:
509
            yaml_config = yaml.load(file, Loader=loader)
lintangsutawika's avatar
lintangsutawika committed
510

lintangsutawika's avatar
lintangsutawika committed
511
512
    if yaml_dir is None:
        yaml_dir = os.path.dirname(yaml_path)
513
514
515
516
517
518
519

    assert yaml_dir is not None

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

520
        if isinstance(include_path, str):
521
522
523
524
525
526
527
528
529
530
531
532
533
            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:
534
                included_yaml_config = load_yaml_config(yaml_path=path, mode=mode)
535
536
537
538
539
540
541
542
                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
543
544


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

lintangsutawika's avatar
lintangsutawika committed
549

550
env = Environment(
Baber's avatar
Baber committed
551
    loader=BaseLoader(), undefined=StrictUndefined, keep_trailing_newline=True
552
)
553
env.filters["regex_replace"] = regex_replace
554
555


Baber's avatar
Baber committed
556
@lru_cache(maxsize=128)
Baber's avatar
Baber committed
557
def _compile(raw: str) -> Template:
Baber's avatar
Baber committed
558
559
560
    return env.from_string(raw)


baberabb's avatar
baberabb committed
561
def apply_template(template: str, doc: dict) -> str:
Baber's avatar
Baber committed
562
    rtemplate = _compile(template)
563
    return rtemplate.render(**doc)
564
565


Baber's avatar
Baber committed
566
567
568
569
570
571
572
def create_iterator(
    raw_iterator: collections.Iterator,
    *,
    rank: int = 0,
    world_size: int = 1,
    limit: int | None = None,
) -> islice:
573
574
575
    """
    Method for creating a (potentially) sliced and limited
    iterator from a raw document iterator. Used for splitting data
576
577
578
    among ranks in multigpu setting or only pulling a sample of documents
    """
    return islice(raw_iterator, rank, limit, world_size)
579
580
581
582
583
584
585
586
587
588


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
589
590
591
592
593
594
595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617


def convert_pil_to_hash(value):
    from io import BytesIO

    img_bytes = BytesIO()
    value.save(img_bytes, format="PNG")
    return hashlib.sha256(str(img_bytes).encode()).hexdigest()


def convert_bytes_to_hash(value):
    return hashlib.sha256(str(value).encode()).hexdigest()


def hash_dict_images(data_dict):
    """
    Create a deep copy of `data_dict` where all bytes and PIL.Image.Image values
    are replaced by their respective hashes using the provided converter functions.

    Parameters:
        data_dict (dict): The input dictionary with arbitrary nesting of dicts and lists.

    Returns:
        dict: A new dictionary with the same structure as `data_dict`, but with all
              bytes and PIL.Image.Image objects replaced by their hashes.
    """

    def _process_value(value):
        # Bytes -> hash
Baber Abbasi's avatar
Baber Abbasi committed
618
619
        from PIL import Image

620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
        if isinstance(value, (bytes, bytearray)):
            return convert_bytes_to_hash(value)
        # PIL Image -> hash
        if isinstance(value, Image.Image):
            return convert_pil_to_hash(value)
        # Nested dictionary -> recurse
        if isinstance(value, dict):
            return {k: _process_value(v) for k, v in value.items()}
        # List or tuple -> recurse, preserving type
        if isinstance(value, list):
            return [_process_value(v) for v in value]
        if isinstance(value, tuple):
            return tuple(_process_value(v) for v in value)
        # Other types remain unchanged
        return value

    # Ensure the top-level is a dict
    if not isinstance(data_dict, dict):
        raise TypeError("Input must be a dictionary")

Baber Abbasi's avatar
Baber Abbasi committed
640
641
642
643
644
    return (
        {key: _process_value(val) for key, val in data_dict.items()}
        if importlib.util.find_spec("PIL")
        else data_dict
    )