evaluator.py 23 KB
Newer Older
lintangsutawika's avatar
lintangsutawika committed
1
import random
Leo Gao's avatar
Leo Gao committed
2
import itertools
FarzanehNakhaee's avatar
FarzanehNakhaee committed
3
import json
lintangsutawika's avatar
lintangsutawika committed
4
import collections
FarzanehNakhaee's avatar
FarzanehNakhaee committed
5
import sys
lintangsutawika's avatar
lintangsutawika committed
6

7
8
import torch

9
import numpy as np
lintangsutawika's avatar
lintangsutawika committed
10
11

import lm_eval.api
12
import lm_eval.tasks
lintangsutawika's avatar
lintangsutawika committed
13
import lm_eval.models
lintangsutawika's avatar
lintangsutawika committed
14
import lm_eval.api.metrics
lintangsutawika's avatar
lintangsutawika committed
15
import lm_eval.api.registry
lintangsutawika's avatar
lintangsutawika committed
16

lintangsutawika's avatar
lintangsutawika committed
17
18
19
20
from lm_eval.utils import (
    positional_deprecated,
    run_task_tests,
    make_table,
21
    create_iterator,
lintangsutawika's avatar
lintangsutawika committed
22
    get_git_commit_hash,
lintangsutawika's avatar
lintangsutawika committed
23
    eval_logger,
lintangsutawika's avatar
lintangsutawika committed
24
)
25

Fabrizio Milo's avatar
Fabrizio Milo committed
26

27
@positional_deprecated
Fabrizio Milo's avatar
Fabrizio Milo committed
28
29
30
31
def simple_evaluate(
    model,
    model_args=None,
    tasks=[],
32
    num_fewshot=None,
Fabrizio Milo's avatar
Fabrizio Milo committed
33
    batch_size=None,
34
    max_batch_size=None,
Fabrizio Milo's avatar
Fabrizio Milo committed
35
    device=None,
haileyschoelkopf's avatar
haileyschoelkopf committed
36
    use_cache=None,
Fabrizio Milo's avatar
Fabrizio Milo committed
37
    limit=None,
Ethan Smith's avatar
Ethan Smith committed
38
39
    bootstrap_iters: int = 100000,
    check_integrity: bool = False,
Fabrizio Milo's avatar
Fabrizio Milo committed
40
    decontamination_ngrams_path=None,
Ethan Smith's avatar
Ethan Smith committed
41
42
    write_out: bool = False,
    log_samples: bool = True,
Fabrizio Milo's avatar
Fabrizio Milo committed
43
):
44
    """Instantiate and evaluate a model on a list of tasks.
45

46
47
48
    :param model: Union[str, LM]
        Name of model or LM object, see lm_eval.models.get_model
    :param model_args: Optional[str]
Fabrizio Milo's avatar
Fabrizio Milo committed
49
        String arguments for each model class, see LM.create_from_arg_string.
50
51
        Ignored if `model` argument is a LM object.
    :param tasks: list[Union[str, Task]]
Leo Gao's avatar
Leo Gao committed
52
        List of task names or Task objects. Task objects will be taken to have name task.EVAL_HARNESS_NAME if defined and type(task).__name__ otherwise.
53
54
    :param num_fewshot: int
        Number of examples in few-shot context
55
    :param batch_size: int or str, optional
56
        Batch size for model
57
58
    :param max_batch_size: int, optional
        Maximal batch size to try with automatic batch size detection
59
    :param device: str, optional
60
        PyTorch device (e.g. "cpu" or "cuda:0") for running models
haileyschoelkopf's avatar
haileyschoelkopf committed
61
62
    :param use_cache: str, optional
        A path to a sqlite db file for caching model responses. `None` if not caching.
63
64
    :param limit: int or float, optional
        Limit the number of examples per task (only use this for testing), If <1, limit is a percentage of the total number of examples.
65
66
    :param bootstrap_iters:
        Number of iterations for bootstrap statistics
Stephen Hogg's avatar
Stephen Hogg committed
67
68
    :param check_integrity: bool
        Whether to run the relevant part of the test suite for the tasks
69
    :param write_out: bool
70
71
72
        If True, write out an example document and model input for checking task integrity
    :param log_samples: bool
        If True, write out all model outputs and documents for per-sample measurement and post-hoc analysis
73
    :return
74
        Dictionary of results
75
    """
76
    random.seed(0)
77
    np.random.seed(1234)
78
79
80
    torch.manual_seed(
        1234
    )  # TODO: this may affect training runs that are run with evaluation mid-run.
81

82
83
84
    assert (
        tasks != []
    ), "No tasks specified, or no tasks found. Please verify the task names."
85
86

    if isinstance(model, str):
Fabrizio Milo's avatar
Fabrizio Milo committed
87
88
        if model_args is None:
            model_args = ""
lintangsutawika's avatar
lintangsutawika committed
89
        lm = lm_eval.api.registry.get_model(model).create_from_arg_string(
lintangsutawika's avatar
lintangsutawika committed
90
91
92
93
94
95
            model_args,
            {
                "batch_size": batch_size,
                "max_batch_size": max_batch_size,
                "device": device,
            },
Fabrizio Milo's avatar
Fabrizio Milo committed
96
        )
97
    else:
98
        assert isinstance(model, lm_eval.api.model.LM)
99
        lm = model
100

haileyschoelkopf's avatar
haileyschoelkopf committed
101
102
103
104
105
106
107
108
109
110
    if use_cache is not None:
        print(f"Using cache at {use_cache + '_rank' + str(lm.rank) + '.db'}")
        lm = lm_eval.api.model.CachingLM(
            lm,
            use_cache
            # each rank receives a different cache db.
            # necessary to avoid multiple writes to cache at once
            + "_rank" + str(lm.rank) + ".db",
        )

111
112
    task_dict = lm_eval.tasks.get_task_dict(tasks)
    for task_name in task_dict.keys():
lintangsutawika's avatar
lintangsutawika committed
113
114
115
        task_obj = task_dict[task_name]
        if type(task_obj) == tuple:
            group, task_obj = task_obj
116
117
            if task_obj is None:
                continue
lintangsutawika's avatar
lintangsutawika committed
118
119

        config = task_obj._config
120
121
122
123
124
125
126
        if num_fewshot is not None:
            if config["num_fewshot"] > 0:
                default_num_fewshot = config["num_fewshot"]
                eval_logger.warning(
                    f"Overwriting default num_fewshot of {task_name} from {default_num_fewshot} to {num_fewshot}"
                )

Lintang Sutawika's avatar
Lintang Sutawika committed
127
            task_obj._config["num_fewshot"] = num_fewshot
Jonathan Tow's avatar
Merge  
Jonathan Tow committed
128

Stephen Hogg's avatar
Stephen Hogg committed
129
    if check_integrity:
130
        run_task_tests(task_list=tasks)
Stephen Hogg's avatar
Stephen Hogg committed
131

132
133
134
135
    results = evaluate(
        lm=lm,
        task_dict=task_dict,
        limit=limit,
Niklas Muennighoff's avatar
Niklas Muennighoff committed
136
        bootstrap_iters=bootstrap_iters,
Fabrizio Milo's avatar
Fabrizio Milo committed
137
        decontamination_ngrams_path=decontamination_ngrams_path,
138
        write_out=write_out,
139
        log_samples=log_samples,
140
    )
141

142
143
144
    if lm.rank == 0:
        # add info about the model and few shot config
        results["config"] = {
lintangsutawika's avatar
lintangsutawika committed
145
146
147
            "model": model
            if isinstance(model, str)
            else model.model.config._name_or_path,
148
149
            "model_args": model_args,
            "batch_size": batch_size,
lintangsutawika's avatar
lintangsutawika committed
150
151
152
            "batch_sizes": list(lm.batch_sizes.values())
            if hasattr(lm, "batch_sizes")
            else [],
153
            "device": device,
haileyschoelkopf's avatar
haileyschoelkopf committed
154
            "use_cache": use_cache,
155
156
157
            "limit": limit,
            "bootstrap_iters": bootstrap_iters,
        }
158
        results["git_hash"] = get_git_commit_hash()
159
160
161
        return results
    else:
        return None
162

Leo Gao's avatar
Leo Gao committed
163

164
decontaminate_suffix = "_decontaminate"
Leo Gao's avatar
Leo Gao committed
165

Fabrizio Milo's avatar
Fabrizio Milo committed
166

167
@positional_deprecated
Fabrizio Milo's avatar
Fabrizio Milo committed
168
169
170
171
def evaluate(
    lm,
    task_dict,
    limit=None,
Ethan Smith's avatar
Ethan Smith committed
172
    bootstrap_iters: int = 100000,
Fabrizio Milo's avatar
Fabrizio Milo committed
173
    decontamination_ngrams_path=None,
Ethan Smith's avatar
Ethan Smith committed
174
175
    write_out: bool = False,
    log_samples: bool = True,
Fabrizio Milo's avatar
Fabrizio Milo committed
176
):
177
178
179
180
181
    """Instantiate and evaluate a model on a list of tasks.

    :param lm: obj
        Language Model
    :param task_dict: dict[str, Task]
haileyschoelkopf's avatar
haileyschoelkopf committed
182
        Dictionary of tasks. Tasks will be taken to have name type(task).config.task .
183
184
185
186
    :param limit: int, optional
        Limit the number of examples per task (only use this for testing)
    :param bootstrap_iters:
        Number of iterations for bootstrap statistics
187
    :param write_out: bool
188
189
190
        If True, write out an example document and model input for checking task integrity
    :param log_samples: bool
        If True, write out all model outputs and documents for per-sample measurement and post-hoc analysis
191
192
193
    :return
        Dictionary of results
    """
194

lintangsutawika's avatar
lintangsutawika committed
195
    # decontaminate = decontamination_ngrams_path is not None
196

197
    # stores the final result for each task, for each metric/filter pair.
Leo Gao's avatar
Leo Gao committed
198
    results = collections.defaultdict(dict)
199
    # Tracks each task's version.
Leo Gao's avatar
Leo Gao committed
200
    versions = collections.defaultdict(dict)
201
    # Tracks the YAML configs of all chosen tasks.
202
    configs = collections.defaultdict(dict)
203
    # logs info about each document evaluated.
lintangsutawika's avatar
lintangsutawika committed
204
    samples = collections.defaultdict(list)
205
    # tracks all Instances/requests a model must generate output on.
Leo Gao's avatar
Leo Gao committed
206
    requests = collections.defaultdict(list)
207
    # Aggregated task scores presented with groups
208
    results_agg = collections.defaultdict(dict)
209
    # Aggregated groups scores only
lintangsutawika's avatar
lintangsutawika committed
210
    groups_agg = collections.defaultdict(dict)
211
212
    # stores the amount to pad out reqs per req. type so that
    # number of fwd passes per distributed rank is equal
213
    padding_requests = collections.defaultdict(int)
lintangsutawika's avatar
lintangsutawika committed
214
    # store the hierarchy to do proper ordering
lintangsutawika's avatar
lintangsutawika committed
215
    task_hierarchy = collections.defaultdict(list)
lintangsutawika's avatar
lintangsutawika committed
216
    # store the ordering of tasks and groups
lintangsutawika's avatar
lintangsutawika committed
217
    task_order = collections.defaultdict(int)
lintangsutawika's avatar
lintangsutawika committed
218
    task_group_alias = collections.defaultdict(dict)
219

220
    # get lists of each type of request
221
    for task_name, task in task_dict.items():
222
        if type(task) == tuple:
lintangsutawika's avatar
lintangsutawika committed
223
224
            group_name, task = task
            task_hierarchy[group_name].append(task_name)
225
            versions[group_name] = "N/A"
lintangsutawika's avatar
lintangsutawika committed
226

227
        else:
228
            group_name = None
lintangsutawika's avatar
lintangsutawika committed
229
230
231
232
            task_hierarchy[task_name] = []

        if task is None:
            continue
233

Leo Gao's avatar
Leo Gao committed
234
        versions[task_name] = task.VERSION
haileyschoelkopf's avatar
haileyschoelkopf committed
235
236
        configs[task_name] = dict(task.dump_config())

lintangsutawika's avatar
lintangsutawika committed
237
238
239
        if "task_alias" in configs[task_name]:
            task_group_alias[task_name] = configs[task_name]["task_alias"]

lintangsutawika's avatar
format  
lintangsutawika committed
240
241
242
243
        if (
            ("group_alias" in configs[task_name])
            and (group_name not in task_group_alias)
            and (group_name is not None)
lintangsutawika's avatar
lintangsutawika committed
244
245
246
        ):
            task_group_alias[group_name] = configs[task_name]["group_alias"]

Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
247
        if limit is not None:
248
249
250
251
252
253
            if task.has_test_docs():
                task_docs = task.test_docs()
            elif task.has_validation_docs():
                task_docs = task.validation_docs()
            else:
                raise RuntimeError("Task has neither test_docs nor validation_docs")
254
            limit = int(len(task_docs) * limit) if limit < 1.0 else int(limit)
255

256
257
        task.build_all_requests(limit=limit, rank=lm.rank, world_size=lm.world_size)

258
        eval_logger.debug(
haileyschoelkopf's avatar
haileyschoelkopf committed
259
260
261
262
263
264
            f"Task: {task_name}; number of requests on this rank: {len(task.instances)}"
        )

        if write_out:
            for inst in task.instances:
                # print the prompt for the first few documents
Hailey Schoelkopf's avatar
Hailey Schoelkopf committed
265
266
                if inst.doc_id < 1:
                    eval_logger.info(
haileyschoelkopf's avatar
haileyschoelkopf committed
267
268
                        f"Task: {task_name}; document {inst.doc_id}; context prompt (starting on next line):\
\n{inst.args[0]}\n(end of prompt on previous line)\ntarget string or answer choice index (starting on next line):\n{task.doc_to_target(inst.doc)}\n(end of target on previous line)"
haileyschoelkopf's avatar
haileyschoelkopf committed
269
                    )
haileyschoelkopf's avatar
haileyschoelkopf committed
270
                    eval_logger.info(f"Request: {str(inst)}")
haileyschoelkopf's avatar
haileyschoelkopf committed
271

272
        # aggregate Instances by LM method requested to get output.
lintangsutawika's avatar
lintangsutawika committed
273
274
275
        for instance in task.instances:
            reqtype = instance.request_type
            requests[reqtype].append(instance)
276
277

        if lm.world_size > 1:
278
279
280
281
            instances_rnk = torch.tensor(len(task._instances), device=lm.device)
            gathered_item = (
                lm.accelerator.gather(instances_rnk).cpu().detach().numpy().tolist()
            )
282

283
            # compute number of pseudobatches to pad with (FSDP/DDP require even batches among ranks)
284
            numpad = max(gathered_item) - gathered_item[lm.rank]
285
            padding_requests[task.OUTPUT_TYPE] += numpad
286

287
    ### Run LM on inputs, get all outputs ###
Leo Gao's avatar
Leo Gao committed
288
289
    # execute each type of request
    for reqtype, reqs in requests.items():
lintangsutawika's avatar
lintangsutawika committed
290
        eval_logger.info("Running {} requests".format(reqtype))
291
292
293
294
        # create `K` copies of each request `req` based off `K = req.repeats`
        cloned_reqs = []
        for req in reqs:
            cloned_reqs.extend([req] * req.repeats)
lintangsutawika's avatar
lintangsutawika committed
295

296
297
        if (lm.world_size > 1) and (padding_requests[reqtype] > 0):
            for _ in range(padding_requests[reqtype]):
298
299
                cloned_reqs.extend([req] * req.repeats)

300
301
302
303
304
305
306
        # run requests through model
        resps = getattr(lm, reqtype)(cloned_reqs)

        # put responses from model into a list of length K for each request.
        for x, req in zip(resps, cloned_reqs):
            req.resps.append(x)

307
308
        if lm.world_size > 1:
            lm.accelerator.wait_for_everyone()
309

310
311
312
    ### Postprocess outputs ###
    # TODO: del model here, maybe (idea: allow user to specify device of e.g. reward model separately)
    for task_name, task in task_dict.items():
313
314
        if type(task) == tuple:
            group, task = task
315
316
            if task is None:
                continue
317
318
319
        task.apply_filters()

    ### Collect values of metrics on all datapoints ###
Leo Gao's avatar
Leo Gao committed
320
321
322
    vals = collections.defaultdict(list)

    # unpack results and sort back in order and return control to Task
323
    for task_name, task in task_dict.items():
324
325
        if type(task) == tuple:
            group, task = task
326
327
            if task is None:
                continue
haileyschoelkopf's avatar
haileyschoelkopf committed
328
329
        # TODO: make it possible to use a different metric per filter
        # iterate over different filters used
330
        for key in task.instances[0].filtered_resps.keys():
331
332
333
334
            doc_iterator = (
                itertools.islice(
                    enumerate(task.test_docs()), lm.rank, limit, lm.world_size
                )
lintangsutawika's avatar
lintangsutawika committed
335
                if task.has_test_docs()
336
337
338
339
                else itertools.islice(
                    enumerate(task.validation_docs()), lm.rank, limit, lm.world_size
                )
            )
340
            for doc_id, doc in doc_iterator:
341
342
                # subset instances to only this document id ; sort by idx
                requests = list(filter(lambda x: x.doc_id == doc_id, task.instances))
343
                requests.sort(key=lambda x: x.idx)
lintangsutawika's avatar
lintangsutawika committed
344
345
346
                metrics = task.process_results(
                    doc, [req.filtered_resps[key] for req in requests]
                )
347
348
349
350
351
352
353
354
355
356
357
358
                if log_samples:
                    target = task.doc_to_target(doc)
                    example = {
                        "doc_id": doc_id,
                        "doc": doc,
                        "target": target,
                        "arguments": [req.args for req in requests],
                        "resps": [req.resps for req in requests],
                        "filtered_resps": [req.filtered_resps[key] for req in requests],
                    }
                    example.update(metrics)
                    samples[task_name].append(example)
359
360
361
                for metric, value in metrics.items():
                    vals[(task_name, key, metric)].append(value)

362
    if lm.world_size > 1:
363
        # if multigpu, then gather data across all ranks
364
365
366
367
368
369
370
371
        # first gather logged samples across all ranks
        for task_name, task_samples in list(samples.items()):
            full_samples = [None] * lm.world_size
            torch.distributed.all_gather_object(full_samples, task_samples)

            samples[task_name] = list(itertools.chain.from_iterable(full_samples))

        # then collect metrics across all ranks
372
373
        vals_torch = collections.defaultdict(list)
        for (task_name, key, metric), items in vals.items():
374
            numitem = 0
375
            if type(items[0]) == tuple:
376
377
                numitem = len(items[0])

378
379
380
381
            if isinstance(items[0], (str, list)):
                # handle the string case
                gathered_items = [None] * lm.accelerator.num_processes
                torch.distributed.all_gather_object(gathered_items, items)
382

383
                gathered_item = list(itertools.chain.from_iterable(gathered_items))
384
            else:
385
386
387
388
389
390
391
392
393
394
                # distributed gather requires all ranks to have same dimensions
                # so we pad out with float32 min value
                pad_value = torch.finfo(torch.float32).min
                metrics_tensor = torch.tensor(items, device=lm.device)

                original_dtype = metrics_tensor.dtype  # store original dtype
                torch_device_tensor = lm.accelerator.pad_across_processes(
                    metrics_tensor.to(torch.float32), pad_index=pad_value
                )
                gathered_item = lm.accelerator.gather(torch_device_tensor)
395

396
397
398
399
400
401
402
403
404
405
406
                if numitem > 0:
                    gathered_filtered = gathered_item[gathered_item[:, 0] != pad_value]
                else:
                    gathered_filtered = gathered_item[gathered_item != pad_value]

                gathered_item = (
                    gathered_filtered.to(original_dtype).cpu().detach().numpy().tolist()
                )
                # reconvert if we were passed a tuple of values
                if numitem > 0:
                    gathered_item = [tuple(g) for g in gathered_item]
407

408
409
            if lm.rank == 0:
                vals_torch[(task_name, key, metric)] = gathered_item
410

411
        vals = vals_torch
412

413
    if lm.rank == 0:
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441

        ### Get task ordering for correct sample-wide aggregation
        group_to_task = {}
        for group in task_hierarchy.keys():
            if group not in task_order:
                task_order[group] = 0

            if len(task_hierarchy[group]) > 0:
                group_to_task[group] = task_hierarchy[group].copy()

            for task in task_hierarchy[group]:

                if task in task_order:
                    task_order[task] += 1
                else:
                    task_order[task] = 1 + task_order[group]

                if task in task_hierarchy:
                    group_to_task[group].remove(task)
                    group_to_task[group].extend(task_hierarchy[task])

        task_to_group = {}
        for group in group_to_task:
            for task in group_to_task[group]:
                if task in task_to_group:
                    task_to_group[task].append(group)
                else:
                    task_to_group[task] = [group]
lintangsutawika's avatar
lintangsutawika committed
442

443
444
445
446
        ### Aggregate results over all datapoints ###
        # aggregate results ; run bootstrap CIs
        for (task_name, key, metric), items in vals.items():
            task = task_dict[task_name]
lintangsutawika's avatar
lintangsutawika committed
447
448
            metric_key = metric + "," + key

449
            if type(task) == tuple:
lintangsutawika's avatar
lintangsutawika committed
450
451
452
453
                group_name, task = task
            else:
                group_name = None

454
            agg_fn = task.aggregation()[metric]
455
456
            results[task_name][metric_key] = agg_fn(items)
            results[task_name]["samples"] = len(items)
lintangsutawika's avatar
lintangsutawika committed
457

458
459
            # hotfix: bleu, chrf, ter seem to be really expensive to bootstrap
            # so we run them less iterations. still looking for a cleaner way to do this
haileyschoelkopf's avatar
haileyschoelkopf committed
460
            if bootstrap_iters > 0:
haileyschoelkopf's avatar
haileyschoelkopf committed
461
462
                stderr = lm_eval.api.metrics.stderr_for_metric(
                    metric=task.aggregation()[metric],
haileyschoelkopf's avatar
haileyschoelkopf committed
463
                    bootstrap_iters=min(bootstrap_iters, 100)
haileyschoelkopf's avatar
haileyschoelkopf committed
464
465
466
                    if metric in ["bleu", "chrf", "ter"]
                    else bootstrap_iters,
                )
467

haileyschoelkopf's avatar
haileyschoelkopf committed
468
469
                if stderr is not None:
                    results[task_name][metric + "_stderr" + "," + key] = stderr(items)
lintangsutawika's avatar
lintangsutawika committed
470
471
                else:
                    results[task_name][metric + "_stderr" + "," + key] = 0
Fabrizio Milo's avatar
Fabrizio Milo committed
472

lintangsutawika's avatar
lintangsutawika committed
473
        if bool(results):
474
475

            for group, task_list in reversed(task_hierarchy.items()):
476

477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
                if task_list == []:
                    total_size = results[group]["samples"]
                else:
                    total_size = 0

                    for task in task_list:
                        metrics = results[task]

                        current_size = metrics.pop("samples")
                        # TODO: There should be a way for users
                        #       to toggle between weighted and
                        #       unweighted averaging
                        # For unweighted averaging, use:
                        #     current_size = 1

                        all_stderr = []
                        for metric in [
                            key for key in metrics.keys() if "_stderr" not in key
                        ]:

                            stderr = "_stderr,".join(metric.split(","))
                            stderr_score = results[task][stderr]
lintangsutawika's avatar
lintangsutawika committed
499
                            var_score = stderr_score**2
500
501
502
503
504
505
506
507
508
509
510
511
                            metric_score = results[task][metric]

                            all_stderr.append(stderr)

                            if metric in results[group]:
                                results[group][metric] = (
                                    results[group][metric] * total_size
                                    + metric_score * current_size
                                ) / (total_size + current_size)
                                # $$s_z^2 = \frac{(n-1) s_x^2 + (m-1) s_y^2}{n+m-1} + \frac{nm(\bar x - \bar y)^2}{(n+m)(n+m-1)}.$$
                                results[group][stderr] = (
                                    (total_size - 1) * results[group][stderr]
lintangsutawika's avatar
lintangsutawika committed
512
                                    + (current_size - 1) * var_score
513
514
515
516
517
518
519
520
521
522
                                ) / (
                                    total_size + current_size - 1
                                ) + total_size * current_size / (
                                    (total_size + current_size)
                                    * (total_size + current_size - 1)
                                ) * (
                                    results[group][metric] - metric_score
                                ) ** 2
                            else:
                                results[group][metric] = metric_score
lintangsutawika's avatar
lintangsutawika committed
523
                                results[group][stderr] = var_score
524
525
526
527
528

                        total_size += current_size

                    for stderr in all_stderr:
                        results[group][stderr] = np.sqrt(results[group][stderr])
lintangsutawika's avatar
lintangsutawika committed
529

530
                results[group]["samples"] = total_size
lintangsutawika's avatar
lintangsutawika committed
531

lintangsutawika's avatar
lintangsutawika committed
532
        def print_tasks(task_hierarchy, task_order, task_version, task_group_alias):
533
534
535
536
537

            results_agg = collections.defaultdict(dict)
            groups_agg = collections.defaultdict(dict)
            for group_name, task_list in task_hierarchy.items():

lintangsutawika's avatar
lintangsutawika committed
538
                order = task_order[group_name]
lintangsutawika's avatar
lintangsutawika committed
539
                results_agg[group_name] = results[group_name].copy()
lintangsutawika's avatar
lintangsutawika committed
540
                results_agg[group_name]["tab"] = order
541
542

                if (order < max(task_order.values())) and (len(task_list) > 0):
lintangsutawika's avatar
lintangsutawika committed
543
                    groups_agg[group_name] = results[group_name].copy()
lintangsutawika's avatar
lintangsutawika committed
544
                    groups_agg[group_name]["tab"] = order
545
546
547
548
549
550
551
552
553

                if task_list != []:
                    for task in sorted(task_list):
                        if task in task_hierarchy:
                            _task_hierarchy = {task: task_hierarchy[task]}
                        else:
                            _task_hierarchy = {task: []}

                        _results_agg, _groups_agg, task_version = print_tasks(
lintangsutawika's avatar
lintangsutawika committed
554
                            _task_hierarchy, task_order, task_version, task_group_alias
555
556
557
558
559
560
561
562
                        )

                        results_agg = {**results_agg, **_results_agg}
                        groups_agg = {**groups_agg, **_groups_agg}

            return results_agg, groups_agg, task_version

        results_agg, groups_agg, versions = print_tasks(
lintangsutawika's avatar
lintangsutawika committed
563
            task_hierarchy, task_order, versions, task_group_alias
564
        )
lintangsutawika's avatar
lintangsutawika committed
565

lintangsutawika's avatar
lintangsutawika committed
566
567
        for task in results_agg:
            task_results = results_agg[task]
lintangsutawika's avatar
lintangsutawika committed
568
569
570
571

            if "samples" in task_results:
                task_results.pop("samples")

lintangsutawika's avatar
lintangsutawika committed
572
            tab_string = ""
lintangsutawika's avatar
lintangsutawika committed
573
574
            if "tab" in task_results:
                tab = task_results.pop("tab")
lintangsutawika's avatar
lintangsutawika committed
575
                tab_string = " " * tab + "- " if tab > 0 else ""
lintangsutawika's avatar
lintangsutawika committed
576
577
578

            if task in task_group_alias:
                task_alias = task_group_alias[task]
579
                results_agg[task]["alias"] = tab_string + task_alias
lintangsutawika's avatar
lintangsutawika committed
580
            else:
581
                results_agg[task]["alias"] = tab_string + task
lintangsutawika's avatar
lintangsutawika committed
582
583
584

        for group in groups_agg:
            group_results = groups_agg[group]
lintangsutawika's avatar
lintangsutawika committed
585
586
587
588

            if "samples" in group_results:
                group_results.pop("samples")

lintangsutawika's avatar
lintangsutawika committed
589
            tab_string = ""
lintangsutawika's avatar
lintangsutawika committed
590
591
            if "tab" in group_results:
                tab = group_results.pop("tab")
lintangsutawika's avatar
lintangsutawika committed
592
                tab_string = " " * tab + "- " if tab > 0 else ""
lintangsutawika's avatar
lintangsutawika committed
593
594
595

            if group in task_group_alias:
                group_alias = task_group_alias[group]
596
                groups_agg[group]["alias"] = tab_string + group_alias
lintangsutawika's avatar
lintangsutawika committed
597
            else:
598
                groups_agg[group]["alias"] = tab_string + group
lintangsutawika's avatar
lintangsutawika committed
599

600
        results_dict = {
601
            "results": dict(results_agg.items()),
lintangsutawika's avatar
lintangsutawika committed
602
            **({"groups": dict(groups_agg.items())} if bool(groups_agg) else {}),
603
604
            "configs": dict(sorted(configs.items())),
            "versions": dict(sorted(versions.items())),
605
        }
606
607
608
609
        if log_samples:
            results_dict["samples"] = dict(samples)

        return results_dict
Fabrizio Milo's avatar
Fabrizio Milo committed
610

611
612
    else:
        return None