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

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

Lintang Sutawika's avatar
Lintang Sutawika committed
20
import numpy as np
Baber's avatar
Baber committed
21
from jinja2 import BaseLoader, Environment, StrictUndefined
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
HIGHER_IS_BETTER_SYMBOLS = {
    True: "↑",
    False: "↓",
}
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def wrap_text(string: str, width: int = 140, **kwargs) -> Optional[str]:
    """
    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,
    )


Baber's avatar
Baber committed
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

def get_logger(level: Optional[str] = None) -> logging.Logger:
    """
    Get a logger with a stream handler that captures all lm_eval logs.

    Args:
        level (Optional[str]): The logging level.
    Example:
        >>> logger = get_logger("INFO")
        >>> logger.info("Log this")
        INFO:lm_eval:Log this!

    Returns:
        logging.Logger: The logger.
    """
    logger = logging.getLogger("lm_eval")
    if not logger.hasHandlers():
        logger.addHandler(logging.StreamHandler())
        logger.setLevel(logging.INFO)
    if level is not None:
        level = getattr(logging, level.upper())
        logger.setLevel(level)
    return logger


def setup_logging(verbosity=logging.INFO, suppress_third_party=True):
    """
    Configure logging for the lm_eval CLI application.

    WARNING: This function is intended for CLI use only. Library users should
    use get_logger() instead to avoid interfering with their application's
    logging configuration.

    Args:
        verbosity: Log level (int) or string name. Can be overridden by LOGLEVEL env var.
        suppress_third_party: Whether to suppress verbose third-party library logs.

    Returns:
        logging.Logger: The configured lm_eval logger instance.
    """
    # Validate verbosity parameter
    if isinstance(verbosity, str):
        level_map = {
            "DEBUG": logging.DEBUG,
            "INFO": logging.INFO,
            "WARNING": logging.WARNING,
            "ERROR": logging.ERROR,
            "CRITICAL": logging.CRITICAL,
        }
        verbosity = level_map.get(verbosity.upper(), logging.INFO)
    elif not isinstance(verbosity, int):
        verbosity = logging.INFO

    # Get log level from environment or use default
    if log_level_env := os.environ.get("LOGLEVEL", None):
        level_map = {
            "DEBUG": logging.DEBUG,
            "INFO": logging.INFO,
            "WARNING": logging.WARNING,
            "ERROR": logging.ERROR,
            "CRITICAL": logging.CRITICAL,
        }
        log_level = level_map.get(log_level_env.upper(), verbosity)
    else:
        log_level = verbosity

    # Get the lm_eval logger directly
    logger = logging.getLogger("lm_eval")

    # Configure custom formatter
Baber Abbasi's avatar
Baber Abbasi committed
117
118
    class CustomFormatter(logging.Formatter):
        def format(self, record):
Baber's avatar
Baber committed
119
            record.name = record.name.removeprefix("lm_eval.")
Baber Abbasi's avatar
Baber Abbasi committed
120
121
122
123
124
125
126
            return super().format(record)

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

Baber's avatar
Baber committed
127
128
129
130
131
    # Check if handler already exists to prevent duplicates
    has_stream_handler = any(
        isinstance(h, logging.StreamHandler) for h in logger.handlers
    )
    if not has_stream_handler:
Baber Abbasi's avatar
Baber Abbasi committed
132
133
        handler = logging.StreamHandler()
        handler.setFormatter(formatter)
Baber's avatar
Baber committed
134
135
136
        logger.addHandler(handler)
        # For CLI use, we disable propagation to avoid duplicate messages
        logger.propagate = False
Baber Abbasi's avatar
Baber Abbasi committed
137

Baber's avatar
Baber committed
138
139
    # Set the logger level
    logger.setLevel(log_level)
Baber Abbasi's avatar
Baber Abbasi committed
140

Baber's avatar
Baber committed
141
142
143
144
145
146
147
    # Optionally suppress verbose third-party library logs
    if suppress_third_party and log_level == logging.DEBUG:
        third_party_loggers = ["urllib3", "filelock", "fsspec"]
        for logger_name in third_party_loggers:
            logging.getLogger(logger_name).setLevel(logging.INFO)

    return logger
Lintang Sutawika's avatar
Lintang Sutawika committed
148
149


150
151
152
153
def hash_string(string: str) -> str:
    return hashlib.sha256(string.encode("utf-8")).hexdigest()


154
155
156
157
158
159
160
161
162
163
164
165
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
166
167
168
    assert len(sep_char) == 1, (
        "separation string must be a single character for escaped splitting"
    )
169
170
171
172
173

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

Baber's avatar
Baber committed
174
    return re.split(r"(?<!\\)" + sep_char, text, maxsplit=maxsplit)
175
176


haileyschoelkopf's avatar
haileyschoelkopf committed
177
178
179
180
181
def handle_arg_string(arg):
    if arg.lower() == "true":
        return True
    elif arg.lower() == "false":
        return False
182
183
184
185
186
187
    elif arg.isnumeric():
        return int(arg)
    try:
        return float(arg)
    except ValueError:
        return arg
haileyschoelkopf's avatar
haileyschoelkopf committed
188
189


190
def handle_non_serializable(o):
Baber's avatar
Baber committed
191
    if isinstance(o, np.integer):
192
193
194
195
196
197
198
        return int(o)
    elif isinstance(o, set):
        return list(o)
    else:
        return str(o)


199
200
201
202
203
204
205
206
207
208
209
210
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
211
def simple_parse_args_string(args_string: str | None) -> dict:
Jason Phang's avatar
gpt3  
Jason Phang committed
212
213
214
215
216
    """
    Parses something like
        args1=val1,arg2=val2
    Into a dictionary
    """
Baber Abbasi's avatar
Baber Abbasi committed
217
218
    if args_string is None:
        return {}
Jason Phang's avatar
Jason Phang committed
219
    args_string = args_string.strip()
Jason Phang's avatar
gpt3  
Jason Phang committed
220
221
    if not args_string:
        return {}
222
    arg_list = [arg for arg in args_string.split(",") if arg]
haileyschoelkopf's avatar
haileyschoelkopf committed
223
    args_dict = {
224
225
        kv[0]: handle_arg_string("=".join(kv[1:]))
        for kv in [arg.split("=") for arg in arg_list]
haileyschoelkopf's avatar
haileyschoelkopf committed
226
    }
Jason Phang's avatar
gpt3  
Jason Phang committed
227
    return args_dict
Leo Gao's avatar
Leo Gao committed
228

Fabrizio Milo's avatar
Fabrizio Milo committed
229

Leo Gao's avatar
Leo Gao committed
230
231
def join_iters(iters):
    for iter in iters:
Leo Gao's avatar
Leo Gao committed
232
        yield from iter
Leo Gao's avatar
Leo Gao committed
233
234


235
236
237
238
239
def group(arr, fn):
    res = collections.defaultdict(list)

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

241
242
    return list(res.values())

Fabrizio Milo's avatar
Fabrizio Milo committed
243

gakada's avatar
gakada committed
244
245
# Returns a list containing all values of the source_list that
# match at least one of the patterns
Baber's avatar
Baber committed
246
def pattern_match(patterns: list[str], source_list: list[str]) -> list[str]:
247
    if isinstance(patterns, str):
248
249
        patterns = [patterns]

gakada's avatar
gakada committed
250
251
252
253
254
255
256
    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
257
def softmax(x) -> np.ndarray:
Lintang Sutawika's avatar
Lintang Sutawika committed
258
259
260
261
262
    """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
263
def general_detokenize(string: str) -> str:
Leo Gao's avatar
Leo Gao committed
264
265
266
    string = string.replace(" n't", "n't")
    string = string.replace(" )", ")")
    string = string.replace("( ", "(")
Fabrizio Milo's avatar
Fabrizio Milo committed
267
268
    string = string.replace('" ', '"')
    string = string.replace(' "', '"')
Leo Gao's avatar
Fix  
Leo Gao committed
269
    string = re.sub(r" (['.,])", r"\1", string)
270
271
272
    return string


273
274
275
276
277
278
279
280
281
282
283
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.
    """
284
    return filename[filename.rfind("_") + 1 :].replace(".jsonl", "")
285
286
287
288
289
290


def sanitize_model_name(model_name: str) -> str:
    """
    Given the model name, returns a sanitized version of it.
    """
Baber's avatar
Baber committed
291
    return re.sub(r"[\"<>:/|\\?*\[\]]+", "__", model_name)
292
293
294
295
296
297
298
299
300


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
301
def get_latest_filename(filenames: list[str]) -> str:
302
303
304
305
306
307
    """
    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
308
def get_results_filenames(filenames: list[str]) -> list[str]:
309
310
311
312
313
314
    """
    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
315
def get_sample_results_filenames(filenames: list[str]) -> list[str]:
316
317
318
319
320
321
    """
    Extracts filenames that correspond to sample results.
    """
    return [f for f in filenames if "/samples_" in f and ".json" in f]


322
def get_rolling_token_windows(
Baber's avatar
Baber committed
323
324
    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
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
    """
    - 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))
351
    yield [prefix_token] + token_list[: first_seq_len - 1], token_list[:first_seq_len]
Jason Phang's avatar
Jason Phang committed
352
353
354
355
356
    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
357

Jason Phang's avatar
Jason Phang committed
358
        yield (
lintangsutawika's avatar
lintangsutawika committed
359
360
            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
361
362
363
        )
        predicted += window_pred_len

Fabrizio Milo's avatar
Fabrizio Milo committed
364

365
def make_disjoint_window(
Baber's avatar
Baber committed
366
367
    pair: tuple[list[int], list[int]],
) -> tuple[list[int], list[int]]:
Fabrizio Milo's avatar
Fabrizio Milo committed
368
    """Takes output from get_rolling_token_windows and makes the context not overlap with the continuation"""
Leo Gao's avatar
Leo Gao committed
369
    a, b = pair
370
    return a[: len(a) - (len(b) - 1)], b
Fabrizio Milo's avatar
Fabrizio Milo committed
371

Jason Phang's avatar
Jason Phang committed
372

373
374
375
376
377
378
379
380
381
382
383
384
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)


385
class Reorderer:
Baber's avatar
Baber committed
386
    def __init__(self, arr: list[Any], fn: Callable) -> None:
baberabb's avatar
baberabb committed
387
388
389
390
391
392
        """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
        """
393
394
395
        self.size = len(arr)
        arr = list(enumerate(arr))
        arr = group(arr, lambda x: fn(x[1]))
396
397
398
        # 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]
399
400
401
        arr.sort(key=lambda x: fn(x[1]))

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

403
    def get_reordered(self):
baberabb's avatar
baberabb committed
404
405
406
407
408
        """Gets the reordered array

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

411
    def get_original(self, newarr):
baberabb's avatar
baberabb committed
412
413
414
415
416
417
418
419
        """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
        """
420
421
422
423
        res = [None] * self.size
        cov = [False] * self.size

        for (inds, _), v in zip(self.arr, newarr):
Fabrizio Milo's avatar
Fabrizio Milo committed
424
            for ind in inds:
425
426
                res[ind] = v
                cov[ind] = True
Fabrizio Milo's avatar
Fabrizio Milo committed
427

428
        assert all(cov)
Fabrizio Milo's avatar
Fabrizio Milo committed
429

430
431
        return res

Fabrizio Milo's avatar
Fabrizio Milo committed
432

Lintang Sutawika's avatar
Lintang Sutawika committed
433
def make_table(result_dict, column: str = "results", sort_results: bool = False):
434
    """Generate table of results."""
435
    from pytablewriter import LatexTableWriter, MarkdownTableWriter
436

lintangsutawika's avatar
lintangsutawika committed
437
    if column == "results":
lintangsutawika's avatar
lintangsutawika committed
438
439
440
        column_name = "Tasks"
    elif column == "groups":
        column_name = "Groups"
lintangsutawika's avatar
lintangsutawika committed
441

lintangsutawika's avatar
lintangsutawika committed
442
    all_headers = [
lintangsutawika's avatar
lintangsutawika committed
443
        column_name,
lintangsutawika's avatar
lintangsutawika committed
444
445
        "Version",
        "Filter",
446
        "n-shot",
lintangsutawika's avatar
lintangsutawika committed
447
        "Metric",
448
        "",
lintangsutawika's avatar
lintangsutawika committed
449
450
451
452
        "Value",
        "",
        "Stderr",
    ]
453

lintangsutawika's avatar
lintangsutawika committed
454
455
456
457
458
    md_writer = MarkdownTableWriter()
    latex_writer = LatexTableWriter()
    md_writer.headers = all_headers
    latex_writer.headers = all_headers

459
460
    values = []

461
462
    keys = result_dict[column].keys()
    if sort_results:
Lintang Sutawika's avatar
Lintang Sutawika committed
463
464
465
        # 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
466
467
468
        keys = sorted(keys)
    for k in keys:
        dic = result_dict[column][k]
Lintang Sutawika's avatar
Lintang Sutawika committed
469
470
        version = result_dict["versions"].get(k, "    N/A")
        n = str(result_dict.get("n-shot", " ").get(k, " "))
Baber's avatar
Baber committed
471
472
        # TODO: fix this
        # higher_is_better = result_dict.get("higher_is_better", {}).get(k, {})
473
474
475
476

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

477
        metric_items = dic.items()
Lintang Sutawika's avatar
Lintang Sutawika committed
478
        metric_items = sorted(metric_items)
479
480

        for (mf), v in metric_items:
481
            m, _, f = mf.partition(",")
482
483
484
            if m.endswith("_stderr"):
                continue

Baber's avatar
Baber committed
485
486
487
            # hib = HIGHER_IS_BETTER_SYMBOLS.get(higher_is_better.get(m), "")
            # TODO: fix
            hib = "↑"
488

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

491
492
            if m + "_stderr" + "," + f in dic:
                se = dic[m + "_stderr" + "," + f]
Baber's avatar
Baber committed
493
                se = "   N/A" if se == "N/A" else f"{se:.4f}"
Lintang Sutawika's avatar
Lintang Sutawika committed
494
                values.append([k, version, f, n, m, hib, v, "±", se])
495
            else:
Lintang Sutawika's avatar
Lintang Sutawika committed
496
                values.append([k, version, f, n, m, hib, v, "", ""])
497
498
499
500
501
502
503
504
505
506
507
            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()


508
509
def positional_deprecated(fn):
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
510
    A decorator to nudge users into passing only keyword args (`kwargs`) to the
511
512
    wrapped function, `fn`.
    """
Fabrizio Milo's avatar
Fabrizio Milo committed
513

Baber's avatar
Baber committed
514
515
    wraps(fn)

516
    def _wrapper(*args, **kwargs):
Fabrizio Milo's avatar
Fabrizio Milo committed
517
518
519
        if len(args) != 1 if inspect.ismethod(fn) else 0:
            print(
                f"WARNING: using {fn.__name__} with positional arguments is "
520
                "deprecated and will be disallowed in a future version of "
Fabrizio Milo's avatar
Fabrizio Milo committed
521
522
                "lm-evaluation-harness!"
            )
523
        return fn(*args, **kwargs)
Fabrizio Milo's avatar
Fabrizio Milo committed
524

525
    return _wrapper
Stephen Hogg's avatar
Stephen Hogg committed
526

Fabrizio Milo's avatar
Fabrizio Milo committed
527

Baber's avatar
Baber committed
528
529
530
531
532
533
534
def create_iterator(
    raw_iterator: collections.Iterator,
    *,
    rank: int = 0,
    world_size: int = 1,
    limit: int | None = None,
) -> islice:
535
536
537
    """
    Method for creating a (potentially) sliced and limited
    iterator from a raw document iterator. Used for splitting data
538
539
540
    among ranks in multigpu setting or only pulling a sample of documents
    """
    return islice(raw_iterator, rank, limit, world_size)
541
542
543
544
545
546
547
548
549
550


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
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579


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
580
581
        from PIL import Image

582
583
584
585
586
587
588
589
590
591
592
593
594
595
596
597
598
599
600
601
        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
602
603
604
605
606
    return (
        {key: _process_value(val) for key, val in data_dict.items()}
        if importlib.util.find_spec("PIL")
        else data_dict
    )
Baber's avatar
Baber committed
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
626
627
628


def regex_replace(string, pattern, repl, count: int = 0):
    """Implements the `re.sub` function as a custom Jinja filter."""
    return re.sub(pattern, repl, string, count=count)


@functools.lru_cache(maxsize=256)
def _compile_tpl(src: str):
    return apply_template._env.from_string(src)


def apply_template(template: str, doc: dict) -> str:
    if not hasattr(apply_template, "_env"):
        apply_template._env = Environment(
            loader=BaseLoader(),
            undefined=StrictUndefined,
            keep_trailing_newline=True,
        )
        apply_template._env.filters["regex_replace"] = regex_replace

    return _compile_tpl(template).render(**doc)