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

Baber's avatar
Baber committed
3
import logging
Baber's avatar
Baber committed
4
from collections.abc import Iterable
Baber's avatar
Baber committed
5
from dataclasses import asdict, dataclass, field
Baber's avatar
Baber committed
6
from typing import TYPE_CHECKING, Callable
Baber's avatar
Baber committed
7

8
from lm_eval.api.filter import FilterEnsemble
Baber's avatar
Baber committed
9
10
11
12
13
14
15
from lm_eval.api.instance import OutputType
from lm_eval.config.metric import MetricConfig
from lm_eval.config.utils import maybe_serialize


if TYPE_CHECKING:
    from lm_eval.api.samplers import ContextSampler
16
    from lm_eval.api.task import Task
17
    from lm_eval.config.template import TemplateConfig
Baber's avatar
Baber committed
18
19
20
21
22
23
24
25
26

eval_logger = logging.getLogger(__name__)


@dataclass
class RepeatConfig:
    """Encapsulates information about a single repeat."""

    repeats: int = 1
Baber's avatar
Baber committed
27
28
    metric_fn: str | Callable = "pass@N"
    kwargs: dict | None = field(default_factory=dict)
Baber's avatar
Baber committed
29
30
31
32
33
34
35


@dataclass
class FilterConfig:
    """Encapsulates information about a single filter."""

    name: str
36
37
    ensemble: FilterEnsemble
    metric_list: list[MetricConfig]
Baber's avatar
Baber committed
38
39
40
41


@dataclass
class FewshotConfig:
42
43
44
    # hack: this returns task.config.num_fewshot
    # to keep in sync as it is runtime-modified
    num_fewshot: Callable[[], int]
Baber's avatar
Baber committed
45
46
47
48
49
    split: str | None = None
    sampler: str | Callable = "default"
    samples: Callable[[], list[dict]] | list[dict] | None = None
    process_docs: Callable[[list[dict]], Iterable[dict]] | None = None
    fewshot_indices: list[int] | None = None
Baber's avatar
Baber committed
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
    rnd: int = field(init=False, default=False)

    def __post_init__(self) -> None:
        if self.samples is not None and not (
            isinstance(self.samples, list) or callable(self.samples)
        ):
            raise TypeError(
                "samples must be either list[dict] or callable returning list[dict]"
            )

        if self.split is not None and self.samples is not None:
            eval_logger.warning(
                "Both split and samples are configured; split will take precedence"
            )

    @property
    def has_source(self) -> bool:
        """Check if any fewshot source is configured."""
        return self.split is not None or self.samples is not None

    def _get_raw_docs(
        self, dataset
Baber's avatar
Baber committed
72
    ) -> list[dict] | Callable[[], Iterable[dict]] | None:
Baber's avatar
Baber committed
73
74
75
76
77
        """Get raw documents from configured source."""
        if self.split is not None:
            return dataset[self.split]

        if self.samples is not None:
Baber's avatar
Baber committed
78
            if isinstance(self.samples, list) or callable(self.samples):
Baber's avatar
Baber committed
79
80
81
82
83
84
                return self.samples
            else:
                raise TypeError(
                    "samples must be either a list of dicts or a callable returning a list"
                )

Baber's avatar
Baber committed
85
    def get_docs(self, dataset) -> Iterable[dict] | None:
Baber's avatar
Baber committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
        """Get processed documents from configured source."""
        raw_docs = self._get_raw_docs(dataset)
        if raw_docs is None:
            return None

        if self.process_docs is not None:
            return self.process_docs(raw_docs)
        return raw_docs

    @property
    def get_sampler(self):
        from lm_eval.api import samplers

        if isinstance(self.sampler, str):
            return samplers.get_sampler(self.sampler)
        elif callable(self.sampler):
            return self.sampler

    def init_sampler(
Baber's avatar
Baber committed
105
106
        self, docs: list[dict], task: Task, rnd=None, fewshot_indices=None
    ) -> ContextSampler:
Baber's avatar
Baber committed
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
        """Initialize the sampler with the given documents and task."""
        if rnd is None:
            raise ValueError(
                "A `random.Random` generator argument must be provided to `rnd` of FewShotSampler!"
            )
        return self.get_sampler(
            docs,
            task,
            rnd=rnd,
            fewshot_indices=fewshot_indices
            if fewshot_indices
            else self.fewshot_indices,
        )


@dataclass
123
class TaskConfig:
Baber's avatar
Baber committed
124
    # task naming/registry
Baber's avatar
Baber committed
125
126
127
    task: str | None = None
    task_alias: str | None = None
    tag: str | list | None = None
Baber's avatar
Baber committed
128
129
130
    # HF dataset options.
    # which dataset to use,
    # and what splits for what purpose
Baber's avatar
Baber committed
131
132
133
134
135
136
137
138
    custom_dataset: Callable | None = None
    dataset_path: str | None = None
    dataset_name: str | None = None
    dataset_kwargs: dict | None = field(default_factory=dict)
    training_split: str | None = None
    validation_split: str | None = None
    test_split: str | None = None
    fewshot_split: str | None = (
Baber's avatar
Baber committed
139
140
141
142
        None  # TODO: assert that this not None if num_fewshot > 0. (?) assert if this is same split as one evaluating (?)
    )
    # formatting / prompting options.
    # see docs/advanced_task_guide.md for more info
Baber's avatar
Baber committed
143
144
145
146
147
    process_docs: Callable | None = None
    doc_to_text: Callable | str | None = None
    doc_to_target: Callable | str | None = None
    doc_to_image: Callable | str | None = None
    doc_to_audio: Callable | str | None = None
Baber's avatar
Baber committed
148
    unsafe_code: bool = False
Baber's avatar
Baber committed
149
150
151
    doc_to_choice: Callable | str | dict | list | None = None
    process_results: Callable | str | None = None
    use_prompt: str | None = None
Baber's avatar
Baber committed
152
153
154
    description: str = ""
    target_delimiter: str = " "
    fewshot_delimiter: str = "\n\n"
Baber's avatar
Baber committed
155
    fewshot_config: dict | None = None
Baber's avatar
Baber committed
156
    # runtime configuration options
Baber's avatar
Baber committed
157
158
    num_fewshot: int | None = 0
    generation_kwargs: dict | None = None
Baber's avatar
Baber committed
159
    # scoring options
Baber's avatar
Baber committed
160
    metric_list: list | None = None
Baber's avatar
Baber committed
161
162
    output_type: OutputType = "generate_until"
    repeats: int = 1
Baber's avatar
Baber committed
163
    filter_list: list[dict] | None = None
Baber's avatar
Baber committed
164
    should_decontaminate: bool = False
Baber's avatar
Baber committed
165
166
167
    doc_to_decontamination_query: str | None = None
    gen_prefix: str | None = None
    metadata: dict | None = field(
Baber's avatar
nit  
Baber committed
168
169
170
        default_factory=dict
    )  # by default, not used in the code. allows for users to pass arbitrary info to tasks

171
    _metric_list: list[MetricConfig] = field(default_factory=list)
Baber's avatar
Baber committed
172
    _filter_list: list[FilterConfig] = field(default_factory=list)
173
    # ds_cfg: DatasetConfig = field(init=False)
Baber's avatar
nit  
Baber committed
174
    fewshot_cfg: FewshotConfig = field(init=False)
Baber's avatar
Baber committed
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211

    def __post_init__(self) -> None:
        ### ---setup generation kwargs--- ###
        if self.generation_kwargs is not None:
            if self.output_type != "generate_until":
                eval_logger.warning(
                    f"[{self.task}] passed `generation_kwargs`, but not using `output_type: generate_until`!"
                )

            if "temperature" in self.generation_kwargs:
                self.generation_kwargs["temperature"] = float(
                    self.generation_kwargs["temperature"]
                )

            if "until" not in self.generation_kwargs:
                eval_logger.warning(
                    f"{self.task}: No `until` specified in `generation_kwargs`! Defaulting to the fewshot_delimiter={repr(self.fewshot_delimiter)}"
                )
                self.generation_kwargs["until"] = [self.fewshot_delimiter]
        else:
            if self.output_type == "generate_until":
                # ensure that we greedily generate in absence of explicit arguments otherwise
                self.generation_kwargs = {
                    "until": (
                        None
                        if self.fewshot_delimiter is None
                        else [self.fewshot_delimiter]
                    ),
                    "do_sample": False,
                    "temperature": 0,
                }
                eval_logger.warning(
                    f"{self.task}: No `generation_kwargs` specified in task config, defaulting to {self.generation_kwargs}"
                )
        # ---setup fewshot config--- #
        _fewshot_cfg = self.fewshot_config if self.fewshot_config is not None else {}
        self.fewshot_cfg = FewshotConfig(
Baber's avatar
Baber committed
212
            num_fewshot=lambda: self.num_fewshot or _fewshot_cfg.get("num_fewshot", 0),
Baber's avatar
Baber committed
213
214
215
216
217
218
219
            split=self.fewshot_split,
            sampler=_fewshot_cfg.get("sampler", "default"),
            samples=_fewshot_cfg.get("samples", None),
            process_docs=_fewshot_cfg.get("process_docs", None),
            fewshot_indices=_fewshot_cfg.get("fewshot_indices", None),
        )

Baber's avatar
Baber committed
220
    def _get_metric(self, metric_list: list[dict] | None = None) -> list[MetricConfig]:
Baber's avatar
nit  
Baber committed
221
222
223
224
225
226
227
228
229
        from lm_eval.api.registry import (
            AGGREGATION_REGISTRY,
            DEFAULT_METRIC_REGISTRY,
            get_aggregation,
            get_metric,
            get_metric_aggregation,
            is_higher_better,
        )

Baber's avatar
Baber committed
230
        # if metric_list defined inside a filter, use that; otherwise use the task's metric_list
231
        metric_list = metric_list or self.metric_list
Baber's avatar
Baber committed
232
        metrics = []
233
        if not metric_list:
Baber's avatar
Baber committed
234
235
236
237
238
239
240
241
242
243
            # ---------- 1. If no metrics defined, use defaults for output type ----------
            _metric_list = DEFAULT_METRIC_REGISTRY[self.output_type]
            eval_logger.info(
                f"No metrics defined in config, using default metrics for {self.output_type}={_metric_list}"
            )
            metrics.extend(
                MetricConfig(
                    name=metric_name,
                    fn=get_metric(metric_name),
                    aggregation_fn=get_metric_aggregation(metric_name),
244
                    higher_is_better=is_higher_better(metric_name) or True,
Baber's avatar
Baber committed
245
246
247
248
                )
                for metric_name in _metric_list
            )
        else:
Baber's avatar
nit  
Baber committed
249
            # ---------- 2. Process user-defined metrics from config ----------
250
            for metric_config in metric_list:
Baber's avatar
Baber committed
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
                metric_name = metric_config["metric"]
                _metric_fn_kwargs = {
                    key: metric_config[key]
                    for key in metric_config
                    if key
                    not in ["metric", "aggregation", "higher_is_better", "hf_evaluate"]
                }
                _hf_evaluate_metric: bool = metric_config.get("hf_evaluate", False)
                _metric_fn = None
                _aggregation = None

                if self.process_results is not None:
                    # User will compute metrics inside `process_results()`
                    _metric_name = None
                    _metric_fn_kwargs = {}
                elif callable(metric_name):
                    # User passed a function object
                    _metric_name = metric_name.__name__
                    _metric_fn = metric_name.__call__
                else:
                    # Normal: look up by name
Baber's avatar
Baber committed
272
273
                    _metric_name = metric_name
                    _metric_fn = get_metric(metric_name, _hf_evaluate_metric)
Baber's avatar
Baber committed
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310

                # ---------- 3. Decide how to aggregate examples ----------
                if "aggregation" in metric_config:
                    if isinstance(_agg_name := metric_config["aggregation"], str):
                        _aggregation = get_aggregation(_agg_name)
                    elif callable(_agg_name):  # noqa: E721
                        _aggregation = metric_config["aggregation"]
                else:
                    INV_AGG_REGISTRY = {v: k for k, v in AGGREGATION_REGISTRY.items()}
                    _aggregation = get_metric_aggregation(metric_name)
                    eval_logger.warning(
                        f"[Task: {self.task}] metric {metric_name} is defined, but aggregation is not. "
                        f"using default "
                        f"aggregation={INV_AGG_REGISTRY[_aggregation]}"
                    )

                # ---------- 4. Determine “higher-is-better” semantics ----------
                if "higher_is_better" in metric_config:
                    _higher_is_better = metric_config["higher_is_better"]
                else:
                    eval_logger.warning(
                        f"[Task: {self.task}] metric {metric_name} is defined, but higher_is_better is not. "
                        f"using default "
                        f"higher_is_better={is_higher_better(metric_name)}"
                    )
                    _higher_is_better = is_higher_better(metric_name)

                metrics.append(
                    MetricConfig(
                        name=_metric_name,
                        fn=_metric_fn,
                        kwargs=_metric_fn_kwargs,
                        aggregation_fn=_aggregation,
                        higher_is_better=_higher_is_better,
                        hf_evaluate=_hf_evaluate_metric,
                    )
                )
311
312
        for m in metrics:
            if m not in self._metric_list:
Baber's avatar
fixup  
Baber committed
313
                self._metric_list.append(m)
Baber's avatar
Baber committed
314
315
316
        return metrics

    @property
Baber's avatar
Baber committed
317
    def get_filters(self) -> list[FilterConfig]:
Baber's avatar
nit  
Baber committed
318
319
        from lm_eval.filters import build_filter_ensemble

Baber's avatar
Baber committed
320
321
322
323
        if not self.filter_list:
            eval_logger.debug(
                "No custom filters defined; falling back to 'take_first' for handling repeats."
            )
324
325
326
327
            return [
                FilterConfig(
                    name="none",
                    ensemble=build_filter_ensemble("none", [("take_first", None)]),
Baber's avatar
fixup  
Baber committed
328
                    metric_list=self._get_metric(metric_list=None),
329
330
                )
            ]
Baber's avatar
Baber committed
331
332
        else:

333
            def _strip_fn(d: dict) -> tuple[str, dict]:
334
335
336
                return d["function"], {
                    k: v for k, v in d.items() if k not in ["function", "metric_list"]
                }
Baber's avatar
Baber committed
337
338
339
340
341
342

            configs = (
                self.filter_list.values()
                if isinstance(self.filter_list, dict)
                else self.filter_list
            )
343
344
345
346
347
348
349
            x = [
                FilterConfig(
                    name=cfg["name"],
                    ensemble=build_filter_ensemble(
                        filter_name=cfg["name"],
                        components=[_strip_fn(f) for f in cfg["filter"]],
                    ),
Baber's avatar
fixup  
Baber committed
350
                    metric_list=self._get_metric(metric_list=cfg.get("metric_list")),
Baber's avatar
Baber committed
351
352
353
                )
                for cfg in configs
            ]
354
            return x
Baber's avatar
Baber committed
355

356
    @classmethod
Baber's avatar
Baber committed
357
    def from_yaml(cls, data: dict) -> TaskConfig:
358
359
360
        """Create a TaskConfig instance from a YAML-like dictionary."""
        return cls(**data)

361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
    @classmethod
    def from_template(cls, template: TemplateConfig, **kwargs) -> TaskConfig:
        """Create a TaskConfig instance from a template.

        Args:
            template: TemplateConfig instance (MCQTemplateConfig or ClozeTemplateConfig)
            **kwargs: Additional arguments to override template defaults

        Returns:
            TaskConfig instance configured from the template
        """
        from lm_eval.config.template import (
            ClozeTemplateConfig,
            MCQTemplateConfig,
        )

        # Extract base configuration from template
        config_dict = {
            "task": template.task,
            "doc_to_text": template.doc_to_text,
            "doc_to_choice": template.doc_to_choice,
            "doc_to_target": template.doc_to_target,
            "description": template.description,
            "target_delimiter": template.target_delimiter,
            "fewshot_delimiter": template.fewshot_delimiter,
            "metric_list": template.metric_list,
        }

        # Add common template attributes if they exist
        if hasattr(template, "answer_suffix"):
            config_dict["target_delimiter"] = (
                template.answer_suffix + template.target_delimiter
            )

        # Handle template-specific configurations
        if isinstance(template, MCQTemplateConfig):
            # For MCQ templates, set up multiple choice specific config
            config_dict["output_type"] = "multiple_choice"

            # MCQ templates typically use accuracy metrics
            if template.metric_list is None:
                config_dict["metric_list"] = [{"metric": "acc"}]

        elif isinstance(template, ClozeTemplateConfig):
            # For Cloze templates, set up generation config
            config_dict["output_type"] = "generate_until"

            # Cloze templates typically use accuracy and normalized accuracy
            if template.metric_list is None:
                config_dict["metric_list"] = [{"metric": "acc"}, {"metric": "acc_norm"}]
        else:
            # Generic template - try to infer output type
            if hasattr(template, "template"):
                if template.template == "mcq":
                    config_dict["output_type"] = "multiple_choice"
                elif template.template == "cloze":
                    config_dict["output_type"] = "generate_until"

        # Override with any user-provided kwargs
        config_dict.update(kwargs)

        # Create and return TaskConfig instance
        return cls(**config_dict)

Baber's avatar
Baber committed
425
426
427
428
429
430
431
    def __getitem__(self, item):
        return getattr(self, item)

    def __setitem__(self, item, value):
        return setattr(self, item, value)

    def to_dict(self, keep_callable: bool = False) -> dict:
Baber's avatar
Baber committed
432
433
434
435
436
437
438
439
        def _ser(x):
            if isinstance(x, dict):
                return {k: _ser(v) for k, v in x.items()}
            if isinstance(x, (list, tuple, set)):
                return type(x)(_ser(i) for i in x)
            return maybe_serialize(x, keep_callable)

        return {k: _ser(v) for k, v in asdict(self).items() if v is not None}