task.py 75 KB
Newer Older
1
import abc
2
import ast
Baber's avatar
TODO!  
Baber committed
3
import collections
lintangsutawika's avatar
lintangsutawika committed
4
import logging
5
import random
6
7
import re
from collections.abc import Callable
8
from copy import deepcopy
9
from dataclasses import asdict, dataclass
10
from inspect import getsource
11
12
13
14
15
16
17
18
19
20
21
22
from typing import (
    Any,
    Dict,
    Iterable,
    Iterator,
    List,
    Literal,
    Mapping,
    Optional,
    Tuple,
    Union,
)
23
24
25

import datasets
import numpy as np
26
from tqdm import tqdm
27
28

from lm_eval import utils
29
from lm_eval.api import samplers
30
31
from lm_eval.api.instance import Instance, OutputType
from lm_eval.api.metrics import bits_per_byte, mean, weighted_perplexity
lintangsutawika's avatar
lintangsutawika committed
32
from lm_eval.api.registry import (
33
34
    AGGREGATION_REGISTRY,
    DEFAULT_METRIC_REGISTRY,
haileyschoelkopf's avatar
haileyschoelkopf committed
35
    get_aggregation,
36
    get_metric,
37
    get_metric_aggregation,
haileyschoelkopf's avatar
haileyschoelkopf committed
38
    is_higher_better,
lintangsutawika's avatar
lintangsutawika committed
39
)
40
from lm_eval.api.schemas import GenerateInput, LoglikelihoodInput, MetricResult
41
from lm_eval.caching.cache import load_from_cache, save_to_cache
42
43
44
from lm_eval.filters import build_filter_ensemble
from lm_eval.prompts import get_prompt

45

46
47
48
49
ALL_OUTPUT_TYPES = [
    "loglikelihood",
    "multiple_choice",
    "loglikelihood_rolling",
50
    "generate_until",
51
52
]

Lintang Sutawika's avatar
Lintang Sutawika committed
53
eval_logger = logging.getLogger(__name__)
54

lintangsutawika's avatar
lintangsutawika committed
55

56
57
@dataclass
class TaskConfig(dict):
58
    # task naming/registry
59
60
    task: Optional[str] = None
    task_alias: Optional[str] = None
Lintang Sutawika's avatar
Lintang Sutawika committed
61
    tag: Optional[Union[str, list]] = None
62
63
64
    # HF dataset options.
    # which dataset to use,
    # and what splits for what purpose
Baber Abbasi's avatar
Baber Abbasi committed
65
    custom_dataset: Optional[Callable] = None
66
67
68
69
70
71
    dataset_path: Optional[str] = None
    dataset_name: Optional[str] = None
    dataset_kwargs: Optional[dict] = None
    training_split: Optional[str] = None
    validation_split: Optional[str] = None
    test_split: Optional[str] = None
72
    fewshot_split: Optional[str] = (
Baber Abbasi's avatar
Baber Abbasi committed
73
        None  # TODO: assert that this not None if num_fewshot > 0. (?) assert if this is same split as one evaluating (?)
74
    )
75
76
    # formatting / prompting options.
    # see docs/advanced_task_guide.md for more info
77
78
79
    process_docs: Optional[Callable] = None
    doc_to_text: Optional[Union[Callable, str]] = None
    doc_to_target: Optional[Union[Callable, str]] = None
80
    doc_to_image: Union[Callable, str] = None
81
    doc_to_audio: Union[Callable, str] = None
Hojin Lee's avatar
Hojin Lee committed
82
    unsafe_code: bool = False
83
84
85
    doc_to_choice: Optional[Union[Callable, str, dict, list]] = None
    process_results: Optional[Union[Callable, str]] = None
    use_prompt: Optional[str] = None
86
    description: str = ""
87
88
    target_delimiter: str = " "
    fewshot_delimiter: str = "\n\n"
89
    fewshot_config: Optional[dict] = None
90
    # runtime configuration options
91
    num_fewshot: Optional[int] = None
92
    # scoring options
93
94
95
    metric_list: Optional[list] = None
    output_type: OutputType = "generate_until"
    generation_kwargs: Optional[dict] = None
96
    repeats: int = 1
97
    filter_list: Optional[Union[str, list]] = None
98
    should_decontaminate: bool = False
99
    doc_to_decontamination_query: Optional[str] = None
Baber Abbasi's avatar
Baber Abbasi committed
100
    gen_prefix: Optional[str] = None
101
    repeat_agg: Optional[str] = None
102
103
104
    metadata: Optional[dict] = (
        None  # by default, not used in the code. allows for users to pass arbitrary info to tasks
    )
105

Ethan Smith's avatar
Ethan Smith committed
106
    def __post_init__(self) -> None:
Lintang Sutawika's avatar
Lintang Sutawika committed
107
        if self.generation_kwargs is not None:
108
            if self.output_type != "generate_until":
109
                eval_logger.warning(
110
                    f"[{self.task}] passed `generation_kwargs`, but not using `output_type: generate_until`!"
Lintang Sutawika's avatar
Lintang Sutawika committed
111
112
113
114
115
116
117
118
                )

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

            if "until" not in self.generation_kwargs:
Baber Abbasi's avatar
Baber Abbasi committed
119
120
121
                eval_logger.warning(
                    f"{self.task}: No `until` specified in `generation_kwargs`! Defaulting to the fewshot_delimiter={repr(self.fewshot_delimiter)}"
                )
122
                self.generation_kwargs["until"] = [self.fewshot_delimiter]
Lintang Sutawika's avatar
Lintang Sutawika committed
123
        else:
124
            if self.output_type == "generate_until":
Lintang Sutawika's avatar
Lintang Sutawika committed
125
126
                # ensure that we greedily generate in absence of explicit arguments otherwise
                self.generation_kwargs = {
127
128
129
130
131
                    "until": (
                        None
                        if self.fewshot_delimiter is None
                        else [self.fewshot_delimiter]
                    ),
Lintang Sutawika's avatar
Lintang Sutawika committed
132
                    "do_sample": False,
Baber Abbasi's avatar
Baber Abbasi committed
133
                    "temperature": 0,
Lintang Sutawika's avatar
Lintang Sutawika committed
134
                }
Baber Abbasi's avatar
Baber Abbasi committed
135
136
137
                eval_logger.warning(
                    f"{self.task}: No `generation_kwargs` specified in task config, defaulting to {self.generation_kwargs}"
                )
138

139
140
141
    def __getitem__(self, item):
        return getattr(self, item)

142
143
144
    def __setitem__(self, item, value):
        return setattr(self, item, value)

145
    def to_dict(self, keep_callable: bool = False) -> dict:
146
147
        """dumps the current config as a dictionary object, as a printable format.
        null fields will not be printed.
haileyschoelkopf's avatar
haileyschoelkopf committed
148
        Used for dumping results alongside full task configuration
149

haileyschoelkopf's avatar
haileyschoelkopf committed
150
151
152
153
154
155
156
157
158
159
        :return: dict
            A printable dictionary version of the TaskConfig object.

        # TODO: should any default value in the TaskConfig not be printed?
        """
        cfg_dict = asdict(self)
        # remove values that are `None`
        for k, v in list(cfg_dict.items()):
            if v is None:
                cfg_dict.pop(k)
160
161
162
163
164
165
166
167
168
169
            elif k == "metric_list":
                for metric_dict in v:
                    for metric_key, metric_value in metric_dict.items():
                        if callable(metric_value):
                            metric_dict[metric_key] = self.serialize_function(
                                metric_value, keep_callable=keep_callable
                            )
                cfg_dict[k] = v
            elif callable(v):
                cfg_dict[k] = self.serialize_function(v, keep_callable=keep_callable)
haileyschoelkopf's avatar
haileyschoelkopf committed
170
        return cfg_dict
171

172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
    def serialize_function(
        self, value: Union[Callable, str], keep_callable=False
    ) -> Union[Callable, str]:
        """Serializes a given function or string.

        If 'keep_callable' is True, the original callable is returned.
        Otherwise, attempts to return the source code of the callable using 'getsource'.
        """
        if keep_callable:
            return value
        else:
            try:
                return getsource(value)
            except (TypeError, OSError):
                return str(value)

188
189
190
191
192
193
194
195
196
197
198

class Task(abc.ABC):
    """A task represents an entire benchmark including its dataset, problems,
    answers, and evaluation methods. See BoolQ for a simple example implementation

    A `doc` can be any python object which represents one instance of evaluation.
    This is usually a dictionary e.g.
        {"question": ..., "answer": ...} or
        {"question": ..., question, answer)
    """

199
    VERSION: Optional[Union[int, str]] = None
200

201
202
    # The name of the `Task` benchmark as denoted in the HuggingFace datasets Hub
    # or a path to a custom `datasets` loading script.
203
    DATASET_PATH: Optional[str] = None
204
205

    # The name of a subset within `DATASET_PATH`.
206
    DATASET_NAME: Optional[str] = None
207

208
    OUTPUT_TYPE: Optional[OutputType] = None
lintangsutawika's avatar
lintangsutawika committed
209

210
211
    def __init__(
        self,
212
213
214
215
        data_dir: Optional[str] = None,
        cache_dir: Optional[str] = None,
        download_mode: Optional[datasets.DownloadMode] = None,
        config: Optional[Mapping] = None,  # Union[dict, TaskConfig]
Ethan Smith's avatar
Ethan Smith committed
216
    ) -> None:
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
        """
        :param data_dir: str
            Stores the path to a local folder containing the `Task`'s data files.
            Use this to specify the path to manually downloaded data (usually when
            the dataset is not publicly accessible).
        :param cache_dir: str
            The directory to read/write the `Task` dataset. This follows the
            HuggingFace `datasets` API with the default cache directory located at:
                `~/.cache/huggingface/datasets`
            NOTE: You can change the cache location globally for a given process
            to another directory:
                `export HF_DATASETS_CACHE="/path/to/another/directory"`
        :param download_mode: datasets.DownloadMode
            How to treat pre-existing `Task` downloads and data.
            - `datasets.DownloadMode.REUSE_DATASET_IF_EXISTS`
                Reuse download and reuse dataset.
            - `datasets.DownloadMode.REUSE_CACHE_IF_EXISTS`
                Reuse download with fresh dataset.
            - `datasets.DownloadMode.FORCE_REDOWNLOAD`
                Fresh download and fresh dataset.
        """
        self.download(data_dir, cache_dir, download_mode)
239
240
241
        self._training_docs: Optional[list] = None
        self._fewshot_docs: Optional[list] = None
        self._instances: Optional[List[Instance]] = None
242

243
        self._config: TaskConfig = TaskConfig({**config}) if config else TaskConfig()
244

lintangsutawika's avatar
lintangsutawika committed
245
        self._filters = [build_filter_ensemble("none", [["take_first", None]])]
246
247
248
        self.fewshot_rnd: Optional[random.Random] = (
            None  # purposely induce errors in case of improper usage
        )
249

250
251
252
253
254
255
    def download(
        self,
        data_dir: Optional[str] = None,
        cache_dir: Optional[str] = None,
        download_mode=None,
    ) -> None:
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
        """Downloads and returns the task dataset.
        Override this method to download the dataset from a custom API.

        :param data_dir: str
            Stores the path to a local folder containing the `Task`'s data files.
            Use this to specify the path to manually downloaded data (usually when
            the dataset is not publicly accessible).
        :param cache_dir: str
            The directory to read/write the `Task` dataset. This follows the
            HuggingFace `datasets` API with the default cache directory located at:
                `~/.cache/huggingface/datasets`
            NOTE: You can change the cache location globally for a given process
            by setting the shell environment variable, `HF_DATASETS_CACHE`,
            to another directory:
                `export HF_DATASETS_CACHE="/path/to/another/directory"`
        :param download_mode: datasets.DownloadMode
            How to treat pre-existing `Task` downloads and data.
            - `datasets.DownloadMode.REUSE_DATASET_IF_EXISTS`
                Reuse download and reuse dataset.
            - `datasets.DownloadMode.REUSE_CACHE_IF_EXISTS`
                Reuse download with fresh dataset.
            - `datasets.DownloadMode.FORCE_REDOWNLOAD`
                Fresh download and fresh dataset.
        """
280
281
282
283
284
285
286
        self.dataset = datasets.load_dataset(
            path=self.DATASET_PATH,
            name=self.DATASET_NAME,
            data_dir=data_dir,
            cache_dir=cache_dir,
            download_mode=download_mode,
        )
287

288
    @property
289
    def config(self) -> TaskConfig:
290
291
292
        """Returns the TaskConfig associated with this class."""
        return self._config

293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
    @abc.abstractmethod
    def has_training_docs(self):
        """Whether the task has a training set"""
        pass

    @abc.abstractmethod
    def has_validation_docs(self):
        """Whether the task has a validation set"""
        pass

    @abc.abstractmethod
    def has_test_docs(self):
        """Whether the task has a test set"""
        pass

308
    def training_docs(self) -> Iterable:
309
310
311
312
313
314
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
        return []

315
    def validation_docs(self) -> Iterable:
316
317
318
319
320
321
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
        return []

322
    def test_docs(self) -> Iterable:
323
324
325
326
327
328
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
        return []

329
    def fewshot_docs(self) -> Iterable:
330
331
332
333
334
335
336
337
338
        """
        :return: Iterable[obj]
            A iterable of any object, that doc_to_text can handle
        """
        if self.has_training_docs():
            return self.training_docs()
        elif self.has_validation_docs():
            return self.validation_docs()
        else:
Baber Abbasi's avatar
Baber Abbasi committed
339
340
341
342
343
            if self.config.get("num_fewshot", 0) > 0:
                eval_logger.warning(
                    f"[Task: {self.config.task}] has_training_docs and has_validation_docs are False"
                    ", using test_docs as fewshot_docs but this is not recommended."
                )
344
345
            return self.test_docs()

346
    def _process_doc(self, doc: dict) -> dict:
347
348
349
350
351
352
353
354
355
        """
        Override this to process (detokenize, strip, replace, etc.) individual
        documents. This can be used in a map over documents of a data split.
        E.g. `map(self._process_doc, self.dataset["validation"])`

        :return: dict
            The processed version of the specified `doc`.
        """
        return doc
lintangsutawika's avatar
lintangsutawika committed
356

357
    @property
358
    def instances(self) -> List[Instance]:
359
360
361
362
363
364
365
366
367
368
369
        """After calling `task.build_all_requests()`, tasks
        maintain a list of the dataset instances which will be evaluated.
        """
        return self._instances

    def fewshot_examples(self, k, rnd):
        if self._training_docs is None:
            self._training_docs = list(self.training_docs())

        return rnd.sample(self._training_docs, k)

370
371
    def doc_to_decontamination_query(self, doc):
        raise NotImplementedError(
372
373
374
375
376
377
378
379
380
381
382
            "Override doc_to_decontamination_query with document specific decontamination query."
        )

    @abc.abstractmethod
    def doc_to_text(self, doc):
        pass

    @abc.abstractmethod
    def doc_to_target(self, doc):
        pass

383
384
385
386
    # not an abstractmethod because not every language-only task has to implement this
    def doc_to_image(self, doc):
        raise NotImplementedError

387
388
389
    def doc_to_audio(self, doc):
        raise NotImplementedError

Baber Abbasi's avatar
Baber Abbasi committed
390
391
392
    def doc_to_prefix(self, doc):
        return ""

393
394
    def build_all_requests(
        self,
395
        *,
396
        limit: Union[int, None] = None,
397
        samples: Optional[List[int]] = None,
398
399
400
401
402
403
404
405
406
        rank: int = 0,
        world_size: int = 1,
        cache_requests: bool = False,
        rewrite_requests_cache: bool = False,
        system_instruction: Optional[str] = None,
        apply_chat_template: bool = False,
        fewshot_as_multiturn: bool = False,
        chat_template: Optional[Callable] = None,
        tokenizer_name: str = "",
407
    ) -> None:
408
        """Build a set of Instances for a task, and store them in task.instances"""
409
410
411
412

        # used with caching
        og_limit = limit

413
        cache_key = f"requests-{self._config.task}-{self.config.num_fewshot}shot-rank{rank}-world_size{world_size}"
KonradSzafer's avatar
KonradSzafer committed
414
415
416
417
418
419
420
        cache_key += "-chat_template" if apply_chat_template else ""
        cache_key += "-fewshot_as_multiturn" if fewshot_as_multiturn else ""
        cache_key += (
            f"-system_prompt_hash{utils.hash_string(system_instruction)}"
            if system_instruction is not None
            else ""
        )
421
        cache_key += f"-tokenizer{tokenizer_name}"
422

Baber Abbasi's avatar
Baber Abbasi committed
423
        cached_instances = load_from_cache(file_name=cache_key, cache=cache_requests)
424
425
426
427
428
429
430
431
432
433
434
435
436

        if cache_requests and cached_instances and not rewrite_requests_cache:
            cached_instances = cached_instances[:limit]

            flattened_instances = [
                instance
                for instance_group in cached_instances
                for instance in instance_group
            ]

            self._instances = flattened_instances
            return

Baber Abbasi's avatar
Baber Abbasi committed
437
        eval_logger.info(f"Building contexts for {self.config.task} on rank {rank}...")
438

439
        instances = []
440
441
442
443
444
445
446
447
448
449

        # process all documents when caching is specified for simplicity
        if (
            cache_requests
            and (not cached_instances or rewrite_requests_cache)
            and limit is not None
        ):
            limit = None

        doc_id_docs = list(
450
451
452
            self.doc_iterator(
                rank=rank, limit=limit, samples=samples, world_size=world_size
            )
453
454
455
456
457
458
459
        )

        num_docs = len(doc_id_docs)

        for doc_id, doc in tqdm(
            doc_id_docs,
            total=num_docs,
lintangsutawika's avatar
lintangsutawika committed
460
        ):
461
            # sample fewshot context #TODO: need to offset doc_id by rank now!
462
            fewshot_ctx = self.fewshot_context(
463
                doc,
464
465
466
467
468
469
470
                num_fewshot=0
                if self.config.num_fewshot is None
                else self.config.num_fewshot,
                system_instruction=system_instruction,
                apply_chat_template=apply_chat_template,
                fewshot_as_multiturn=fewshot_as_multiturn,
                chat_template=chat_template,
Baber Abbasi's avatar
Baber Abbasi committed
471
                gen_prefix=self.doc_to_prefix(doc),
472
            )
473

474
            # TODO: we should override self.config.repeats if doing greedy gen so users don't waste time+compute
lintangsutawika's avatar
lintangsutawika committed
475
476
477
            inst = self.construct_requests(
                doc=doc,
                ctx=fewshot_ctx,
478
                metadata=(self.config["task"], doc_id, self.config.repeats),
479
                apply_chat_template=apply_chat_template,
480
                chat_template=chat_template,
lintangsutawika's avatar
lintangsutawika committed
481
            )
482
483
484
485

            if not isinstance(inst, list):
                inst = [inst]

486
487
488
489
490
491
492
493
494
495
496
497
498
            instances.append(inst)

        # now flatten, this is to allow slicing to work with pickles

        sliced_instances = instances[:og_limit]

        flattened_instances = [
            instance
            for instance_group in sliced_instances
            for instance in instance_group
        ]

        self._instances = flattened_instances
499

500
501
        if len(self._instances) == 0:
            raise ValueError("task.build_requests() did not find any docs!")
502

503
504
505
        if cache_requests and (not cached_instances or rewrite_requests_cache):
            save_to_cache(file_name=cache_key, obj=instances)

506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
    @abc.abstractmethod
    def construct_requests(self, doc, ctx, **kwargs):
        """Uses RequestFactory to construct Requests and returns an iterable of
        Requests which will be sent to the LM.

        :param doc:
            The document as returned from training_docs, validation_docs, or test_docs.
        :param ctx: str
            The context string, generated by fewshot_context. This includes the natural
            language description, as well as the few shot examples, and the question
            part of the document for `doc`.
        :param doc_idx: int
            The index of a document within `self.test_docs()` or `self.validation_docs()`,
            whichever is the main split used.
        :param repeats: int
        TODO: update this docstring
lintangsutawika's avatar
lintangsutawika committed
522
            The number of times each instance in a dataset is inferred on. Defaults to 1,
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
            can be increased for techniques like majority voting.
        """
        pass

    @abc.abstractmethod
    def process_results(self, doc, results):
        """Take a single document and the LM results and evaluates, returning a
        dict where keys are the names of submetrics and values are the values of
        the metric for that one document

        :param doc:
            The document as returned from training_docs, validation_docs, or test_docs.
        :param results:
            The results of the requests created in construct_requests.
        """
        pass

    @abc.abstractmethod
    def aggregation(self):
        """
        :returns: {str: [metric_score] -> float}
            A dictionary where keys are the names of submetrics and values are
            functions that aggregate a list of metric scores
        """
        pass

    @abc.abstractmethod
    def higher_is_better(self):
        """
        :returns: {str: bool}
            A dictionary where keys are the names of submetrics and values are
            whether a higher value of the submetric is better
        """
        pass

558
559
560
    def get_config(self, key: str) -> Any:
        return getattr(self._config, key, None)

haileyschoelkopf's avatar
haileyschoelkopf committed
561
562
563
564
565
566
567
568
569
570
    @classmethod
    def count_bytes(cls, doc):
        """Used for byte-level perplexity metrics in rolling loglikelihood"""
        return len(doc.encode("utf-8"))

    @classmethod
    def count_words(cls, doc):
        """Downstream loglikelihood_rolling perplexity tasks with custom word boundaries should override this!"""
        return len(re.split(r"\s+", doc))

571
    @utils.positional_deprecated
Baber Abbasi's avatar
Baber Abbasi committed
572
    def fewshot_context(self, doc, num_fewshot, rnd=None, description=None, **kwargs):
573
574
575
576
577
578
579
        """Returns a fewshot context string that is made up of a prepended description
        (if provided), the `num_fewshot` number of examples, and an appended prompt example.

        :param doc: str
            The document as returned from training_docs, validation_docs, or test_docs.
        :param num_fewshot: int
            The number of fewshot examples to provide in the returned context string.
lintangsutawika's avatar
lintangsutawika committed
580
581
582
583
584
        :param rnd: random.Random
            The pseudo-random number generator used to randomly sample examples.
            WARNING: This is currently a required arg although it's optionalized with a default `None`.
        :param description: str
            The task's description that will be prepended to the fewshot examples.
585
586
587
        :returns: str
            The fewshot context.
        """
588
        if rnd is None:
589
590
591
592
593
594
            if self.fewshot_rnd is not None:
                rnd = self.fewshot_rnd
            else:
                raise ValueError(
                    "A `random.Random` generator argument must be provided to `rnd`"
                )
lintangsutawika's avatar
lintangsutawika committed
595

596
        description = description if description else ""
597
598

        if num_fewshot == 0:
lintangsutawika's avatar
lintangsutawika committed
599
            labeled_examples = ""
600
        else:
lintangsutawika's avatar
lintangsutawika committed
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
            # for sets with no training docs, draw from other set *but ensure no overlap with current doc*
            if self.has_training_docs():
                fewshotex = self.fewshot_examples(k=num_fewshot, rnd=rnd)
            else:
                if self._fewshot_docs is None:
                    self._fewshot_docs = list(
                        self.validation_docs()
                        if self.has_validation_docs()
                        else self.test_docs()
                    )

                fewshotex = rnd.sample(self._fewshot_docs, num_fewshot + 1)

                # get rid of the doc that's the one we're evaluating, if it's in the fewshot
                fewshotex = [x for x in fewshotex if x != doc][:num_fewshot]

            labeled_examples = (
                "\n\n".join(
                    [
                        self.doc_to_text(doc) + self.doc_to_target(doc)
                        for doc in fewshotex
                    ]
                )
                + "\n\n"
lintangsutawika's avatar
lintangsutawika committed
625
            )
626
627

        example = self.doc_to_text(doc)
lintangsutawika's avatar
lintangsutawika committed
628
        return description + labeled_examples + example
629

630
    def apply_filters(self) -> Optional[List[Instance]]:
Baber Abbasi's avatar
Baber Abbasi committed
631
        """Iterates over FilterEnsembles and applies them to instances"""
lintangsutawika's avatar
lintangsutawika committed
632
633
        if hasattr(self, "_filters"):
            for f in self._filters:
634
                f.apply(self._instances)
lintangsutawika's avatar
lintangsutawika committed
635
636
637
        else:
            eval_logger.warning("No filter defined, passing through instances")
            return self._instances
638

baberabb's avatar
baberabb committed
639
    def dump_config(self) -> dict:
Baber Abbasi's avatar
Baber Abbasi committed
640
        """Returns the config as a dictionary."""
641
        # TODO: this should only return the overrides applied to a non-YAML task's configuration.
642
        # (num_fewshot)
643
        return self.config.to_dict()
644

Baber Abbasi's avatar
Baber Abbasi committed
645
646
647
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
    def set_config(self, key: str, value: Any, update: bool = False) -> None:
        """Set or update the configuration for a given key."""
        if key is None:
            raise ValueError("Key must be provided.")

        if update:
            current_value = getattr(self._config, key, {})
            if not isinstance(current_value, dict):
                raise TypeError(
                    f"Expected a dict for key '{key}', got {type(current_value).__name__} instead."
                )
            current_value.update(value)
        else:
            setattr(self._config, key, value)

    def override_metric(self, metric_name: str) -> None:
        """
        Override the default metrics used for evaluation with custom metrics.

        Parameters:
        - metric_name (str): The name of the custom metric to override. Should be registered in api.metrics.
        """
        (
            self._metric_fn_list,
            self._aggregation_list,
            self._metric_fn_kwargs,
            self._higher_is_better,
        ) = ({}, {}, {}, {})
        self._metric_fn_list[metric_name] = get_metric(metric_name)
        self._aggregation_list[metric_name] = get_metric_aggregation(metric_name)
        self._higher_is_better[metric_name] = is_higher_better(metric_name)
        self._metric_fn_kwargs[metric_name] = {}
        if not isinstance(self, ConfigurableTask):
            self.process_results = lambda x, y: {metric_name: get_metric(metric_name)}
            self.aggregation = lambda: {
                metric_name: get_metric_aggregation(metric_name)
            }
        setattr(self._config, "metric_list", [{"metric": metric_name}])
        setattr(self._config, "process_results", None)

685
686
687
688
689
    def set_fewshot_seed(self, seed: Optional[int] = None) -> None:
        self.fewshot_rnd = random.Random(seed)
        if hasattr(self, "sampler"):
            self.sampler.rnd = self.fewshot_rnd

690
691
692
693
694
695
696
    @property
    def eval_docs(self) -> Union[datasets.Dataset, List[dict]]:
        if self.has_test_docs():
            return self.test_docs()
        elif self.has_validation_docs():
            return self.validation_docs()
        else:
697
698
699
            raise ValueError(
                f"Task dataset (path={self.DATASET_PATH}, name={self.DATASET_NAME}) must have valid or test docs!"
            )
700
701

    def doc_iterator(
702
703
704
705
706
707
        self,
        *,
        rank: int = 0,
        limit: Union[int, None] = None,
        world_size: int = 1,
        samples: Optional[List[int]] = None,
708
    ) -> Iterator[Tuple[int, Any]]:
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
730
        if samples:
            n = len(self.eval_docs)
            assert all([e < n for e in samples]), (
                f"Elements of --samples should be in the interval [0,k-1] where k is the number of total examples. In this case, k={n}."
            )
            eval_logger.info(
                f"{self.config.task}: Evaluating on {len(samples)} examples"
            )
            doc_iterator = utils.create_iterator(
                enumerate(x for i, x in enumerate(self.eval_docs) if i in samples),
                rank=int(rank),
                limit=None,  # limit does not matter here since we are selecting samples directly
                world_size=int(world_size),
            )
        else:
            limit = int(limit) if limit else None
            doc_iterator = utils.create_iterator(
                enumerate(self.eval_docs),
                rank=int(rank),
                limit=limit,
                world_size=int(world_size),
            )
731
732
        return doc_iterator

733
734

class ConfigurableTask(Task):
735
    VERSION = "Yaml"
736
    OUTPUT_TYPE = None
737
    CONFIG = None
738
739

    def __init__(
740
741
742
743
744
        self,
        data_dir=None,
        cache_dir=None,
        download_mode=None,
        config: Optional[dict] = None,
Ethan Smith's avatar
Ethan Smith committed
745
    ) -> None:  # TODO no super() call here
746
        # Get pre-configured attributes
747
        self._config = self.CONFIG
748

749
        # Use new configurations if there was no preconfiguration
750
        if self.config is None:
751
            self._config = TaskConfig(**config)
752
753
        # Overwrite configs
        else:
lintangsutawika's avatar
lintangsutawika committed
754
            if config is not None:
755
                self._config.__dict__.update(config)
756

757
        if self.config is None:
lintangsutawika's avatar
lintangsutawika committed
758
759
760
            raise ValueError(
                "Must pass a config to ConfigurableTask, either in cls.CONFIG or `config` kwarg"
            )
761

762
763
764
765
        if isinstance(self.config.metadata, dict):
            if "version" in self.config.metadata:
                self.VERSION = self.config.metadata["version"]

766
        if self.config.output_type is not None:
767
768
769
770
            if self.config.output_type not in ALL_OUTPUT_TYPES:
                raise ValueError(
                    f"Got invalid output_type '{self.config.output_type}', must be in '{','.join(ALL_OUTPUT_TYPES)}'"
                )
771
            self.OUTPUT_TYPE = self.config.output_type
772

773
774
775
776
        if self.config.doc_to_image is not None:
            # mark the task as requiring multimodality.
            self.MULTIMODAL = True

777
778
779
780
        if self.config.doc_to_audio:
            # mark the task as requiring multimodality.
            self.MULTIMODAL = True

Hojin Lee's avatar
Hojin Lee committed
781
782
783
        if self.config.unsafe_code is not False:
            self.UNSAFE_CODE = True

784
785
        if self.config.dataset_path is not None:
            self.DATASET_PATH = self.config.dataset_path
786

787
788
        if self.config.dataset_name is not None:
            self.DATASET_NAME = self.config.dataset_name
789

790
791
792
793
        self._metric_fn_list = {}
        self._metric_fn_kwargs = {}
        self._aggregation_list = {}
        self._higher_is_better = {}
794

795
        if self.config.metric_list is None:
796
            # TODO: handle this in TaskConfig.__post_init__ ?
797
798
            _metric_list = DEFAULT_METRIC_REGISTRY[self.config.output_type]

799
            for metric_name in _metric_list:
haileyschoelkopf's avatar
haileyschoelkopf committed
800
                self._metric_fn_list[metric_name] = get_metric(metric_name)
lintangsutawika's avatar
lintangsutawika committed
801
                self._metric_fn_kwargs[metric_name] = {}
802
803
804
                self._aggregation_list[metric_name] = get_metric_aggregation(
                    metric_name
                )
haileyschoelkopf's avatar
haileyschoelkopf committed
805
                self._higher_is_better[metric_name] = is_higher_better(metric_name)
806
        else:
807
            for metric_config in self.config.metric_list:
808
809
810
811
                if "metric" not in metric_config:
                    raise ValueError(
                        "'metric' key not provided for an entry in 'metric_list', must be specified!"
                    )
812
813
814
815
                metric_name = metric_config["metric"]
                kwargs = {
                    key: metric_config[key]
                    for key in metric_config
Chris's avatar
Chris committed
816
817
                    if key
                    not in ["metric", "aggregation", "higher_is_better", "hf_evaluate"]
818
                }
Chris's avatar
Chris committed
819
820
821
822
                hf_evaluate_metric = (
                    "hf_evaluate" in metric_config
                    and metric_config["hf_evaluate"] is True
                )
823

824
                if self.config.process_results is not None:
825
826
                    self._metric_fn_list[metric_name] = None
                    self._metric_fn_kwargs[metric_name] = {}
827
828
829
830
831
832
                elif callable(metric_name):
                    metric_fn = metric_name.__call__
                    metric_name = metric_name.__name__
                    self._metric_fn_list[metric_name] = metric_fn
                    self._metric_fn_kwargs[metric_name] = kwargs
                else:
Chris's avatar
Chris committed
833
834
835
                    self._metric_fn_list[metric_name] = get_metric(
                        metric_name, hf_evaluate_metric
                    )
836
                    self._metric_fn_kwargs[metric_name] = kwargs
lintangsutawika's avatar
lintangsutawika committed
837

838
                if "aggregation" in metric_config:
839
                    agg_name = metric_config["aggregation"]
840
                    if isinstance(agg_name, str):
haileyschoelkopf's avatar
haileyschoelkopf committed
841
                        self._aggregation_list[metric_name] = get_aggregation(agg_name)
842
                    elif callable(agg_name):  # noqa: E721
843
844
845
                        self._aggregation_list[metric_name] = metric_config[
                            "aggregation"
                        ]
846
                else:
847
                    INV_AGG_REGISTRY = {v: k for k, v in AGGREGATION_REGISTRY.items()}
lintangsutawika's avatar
lintangsutawika committed
848
                    metric_agg = get_metric_aggregation(metric_name)
849
                    eval_logger.warning(
850
                        f"[Task: {self.config.task}] metric {metric_name} is defined, but aggregation is not. "
851
852
                        f"using default "
                        f"aggregation={INV_AGG_REGISTRY[metric_agg]}"
853
                    )
854
                    self._aggregation_list[metric_name] = metric_agg
lintangsutawika's avatar
lintangsutawika committed
855

856
857
858
859
860
861
                if "higher_is_better" in metric_config:
                    self._higher_is_better[metric_name] = metric_config[
                        "higher_is_better"
                    ]
                else:
                    eval_logger.warning(
862
                        f"[Task: {self.config.task}] metric {metric_name} is defined, but higher_is_better is not. "
863
                        f"using default "
haileyschoelkopf's avatar
haileyschoelkopf committed
864
                        f"higher_is_better={is_higher_better(metric_name)}"
865
                    )
haileyschoelkopf's avatar
haileyschoelkopf committed
866
                    self._higher_is_better[metric_name] = is_higher_better(metric_name)
867

868
        self.download(self.config.dataset_kwargs)
869
870
871
        self._training_docs = None
        self._fewshot_docs = None

872
        if self.config.filter_list is not None:
lintangsutawika's avatar
lintangsutawika committed
873
            self._filters = []
874
            for filter_config in self.config.filter_list:
875
876
877
878
879
880
881
882
883
                filter_name = filter_config["name"]
                filter_functions = filter_config["filter"]
                components = []
                for function in filter_functions:
                    kwargs = {
                        key: function[key] for key in function if key != "function"
                    }
                    components.append([function["function"], kwargs])
                filter_pipeline = build_filter_ensemble(filter_name, components)
lintangsutawika's avatar
lintangsutawika committed
884
                self._filters.append(filter_pipeline)
lintangsutawika's avatar
lintangsutawika committed
885
        else:
Baber Abbasi's avatar
Baber Abbasi committed
886
887
888
889
            # TODO: handle repeats in a more general way rather than just discarding
            eval_logger.debug(
                "No custom filters defined. Using default 'take_first' filter for handling repeats."
            )
890
            # self._filters = [build_filter_ensemble("none", [["take_first", None]])]
891

892
893
        if self.config.use_prompt is not None:
            eval_logger.info(f"loading prompt {self.config.use_prompt}")
894
            self.prompt = get_prompt(
895
                self.config.use_prompt, self.DATASET_PATH, self.DATASET_NAME
lintangsutawika's avatar
lintangsutawika committed
896
            )
897
898
899
        else:
            self.prompt = None

lintangsutawika's avatar
lintangsutawika committed
900
        if self.fewshot_docs() is not None:
901
902
903
904
            self.fewshot_rnd = (
                random.Random()
            )  # setting with no seed, to be overridden at a later time
            config_sampler: Union[str, Callable] = (
haileyschoelkopf's avatar
haileyschoelkopf committed
905
906
907
                self.config.fewshot_config.get("sampler", "default")
                if self.config.fewshot_config
                else "default"
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
923
            )
            if isinstance(config_sampler, str):
                self.sampler = samplers.get_sampler(config_sampler)(
                    list(self.fewshot_docs()), self, rnd=self.fewshot_rnd
                )
            elif callable(config_sampler) and issubclass(
                config_sampler, samplers.ContextSampler
            ):
                self.sampler = config_sampler(
                    docs=list(self.fewshot_docs()), task=self, rnd=self.fewshot_rnd
                )
            else:
                raise TypeError(
                    f"fewshot_config.sampler should be a string or callable of ContextSampler type, "
                    f"not {type(config_sampler)}"
                )
924

925
        self.task_docs = self.eval_docs
926

927
        # Test One Doc
928
        self.features = list(self.task_docs.features.keys())
929
930
        self.multiple_input = 0
        self.multiple_target = 0
931
        test_doc = self.task_docs[0]
932
        test_text = self.doc_to_text(test_doc)
933
        test_target = self.doc_to_target(test_doc)
lintangsutawika's avatar
lintangsutawika committed
934

935
        if self.config.doc_to_choice is not None:
lintangsutawika's avatar
lintangsutawika committed
936
            test_choice = self.doc_to_choice(test_doc)
937
            if not isinstance(test_choice, list):
lintangsutawika's avatar
lintangsutawika committed
938
                eval_logger.error("doc_to_choice must return list")
939
940
            else:
                num_choice = len(test_choice)
941

942
            if isinstance(test_text, int):
Baber Abbasi's avatar
Baber Abbasi committed
943
944
945
                eval_logger.debug(
                    "doc_to_text returned an int. Assuming multiple inputs."
                )
946
                self.multiple_input = num_choice
947
948
        else:
            test_choice = None
949

950
        if isinstance(test_target, list):
Baber Abbasi's avatar
Baber Abbasi committed
951
952
953
            eval_logger.debug(
                "doc_to_target returned a list. Assuming multiple targets."
            )
954
            self.multiple_target = len(test_target)
955
        else:
956
            if (isinstance(test_target, int)) and (test_choice is not None):
lintangsutawika's avatar
lintangsutawika committed
957
                test_target = test_choice[test_target]
958
            else:
lintangsutawika's avatar
lintangsutawika committed
959
                test_target = str(test_target)
960

961
962
963
        if test_choice is not None:
            check_choices = test_choice
        else:
lintangsutawika's avatar
lintangsutawika committed
964
            check_choices = [test_target]
965
966
967
968
        if self.config.doc_to_choice is not None:
            for choice in check_choices:
                choice_has_whitespace = True if choice[0].isspace() else False
                delimiter_has_whitespace = (
969
970
                    True
                    if self.config.target_delimiter.rstrip()
971
                    != self.config.target_delimiter
972
                    else False
973
                )
974

975
                if delimiter_has_whitespace and choice_has_whitespace:
976
977
                    eval_logger.debug(
                        f'Both target_delimiter "{self.config.target_delimiter}" and target choice: "{choice}" have whitespace'
978
979
                    )
                elif (not delimiter_has_whitespace) and (not choice_has_whitespace):
980
                    eval_logger.debug(
981
                        f'Both target_delimiter "{self.config.target_delimiter}" and target choice: "{choice}" do not have whitespace, ignore if the language you are evaluating on does not require/use whitespace'
982
983
                    )

Baber Abbasi's avatar
Baber Abbasi committed
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
1000
    def download(
        self, dataset_kwargs: Optional[Dict[str, Any]] = None, **kwargs
    ) -> None:
        if isinstance(self.config.custom_dataset, Callable):
            eval_logger.warning(
                f"{self.config.task}: Custom kwargs can be passed to `--metadata` in console (as json string) or to the TaskManager."
                + "\nFor example --metadata='{\"max_seq_lengths\":[4096, 8192]}'. For details see task Readme."
            )
            self.dataset = self.config.custom_dataset(
                **(self.config.metadata or {}), **(self.config.dataset_kwargs or {})
            )
        else:
            self.dataset = datasets.load_dataset(
                path=self.DATASET_PATH,
                name=self.DATASET_NAME,
                **dataset_kwargs if dataset_kwargs is not None else {},
            )
1001

baberabb's avatar
baberabb committed
1002
    def has_training_docs(self) -> bool:
1003
        if self.config.training_split is not None:
1004
1005
1006
1007
            return True
        else:
            return False

baberabb's avatar
baberabb committed
1008
    def has_validation_docs(self) -> bool:
1009
        if self.config.validation_split is not None:
1010
1011
1012
1013
            return True
        else:
            return False

baberabb's avatar
baberabb committed
1014
    def has_test_docs(self) -> bool:
1015
        if self.config.test_split is not None:
1016
1017
1018
1019
            return True
        else:
            return False

baberabb's avatar
baberabb committed
1020
    def training_docs(self) -> datasets.Dataset:
1021
        if self.has_training_docs():
1022
1023
1024
            if self.config.process_docs is not None:
                return self.config.process_docs(
                    self.dataset[self.config.training_split]
1025
                )
1026
            return self.dataset[self.config.training_split]
1027

baberabb's avatar
baberabb committed
1028
    def validation_docs(self) -> datasets.Dataset:
1029
        if self.has_validation_docs():
1030
1031
1032
            if self.config.process_docs is not None:
                return self.config.process_docs(
                    self.dataset[self.config.validation_split]
1033
                )
1034
            return self.dataset[self.config.validation_split]
1035

baberabb's avatar
baberabb committed
1036
    def test_docs(self) -> datasets.Dataset:
1037
        if self.has_test_docs():
1038
1039
1040
            if self.config.process_docs is not None:
                return self.config.process_docs(self.dataset[self.config.test_split])
            return self.dataset[self.config.test_split]
1041

1042
    def fewshot_docs(self):
1043
        if self.config.fewshot_split is not None:
1044
1045
            if self.config.process_docs is not None:
                return self.config.process_docs(self.dataset[self.config.fewshot_split])
1046
            return self.dataset[self.config.fewshot_split]
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
1058
        elif (
            self.config.fewshot_config is not None
            and self.config.fewshot_config.get("samples", None) is not None
        ):
            if isinstance(self.config.fewshot_config["samples"], list):
                return self.config.fewshot_config["samples"]
            elif callable(self.config.fewshot_config["samples"]):
                return self.config.fewshot_config["samples"]()
            else:
                raise Exception(
                    "`fewshot_config['samples']` was incorrectly defined in the configuration. It should be either a list of samples as a dict, or function returning this list."
                )
1059
        else:
1060
            if (self.config.num_fewshot is not None) and (self.config.num_fewshot > 0):
1061
                eval_logger.warning(
Lintang Sutawika's avatar
Lintang Sutawika committed
1062
                    f"[Task: {self.config.task}] "
1063
1064
1065
1066
                    "num_fewshot > 0 but fewshot_split is None. "
                    "using preconfigured rule."
                )
            return super().fewshot_docs()
1067

KonradSzafer's avatar
KonradSzafer committed
1068
1069
1070
1071
1072
    @staticmethod
    def append_target_question(
        labeled_examples: List[Dict[str, str]],
        question: str,
        fewshot_as_multiturn: bool = False,
Baber Abbasi's avatar
Baber Abbasi committed
1073
        gen_prefix: Optional[str] = None,
KonradSzafer's avatar
KonradSzafer committed
1074
1075
1076
1077
1078
1079
1080
1081
    ) -> None:
        """Adds a target question to the labeled examples list.
        If fewshot_as_multiturn is True, or labeled_examples is empty, or the last entry is a system turn, appends the question as a new user entry.
        Otherwise, it is appended to the last user entry, ensuring that the conversation alternates between the user and the assistant.
        """
        if not fewshot_as_multiturn:
            # if no messages or last message is system, append as new user entry
            if len(labeled_examples) == 0 or labeled_examples[-1]["role"] == "system":
1082
                labeled_examples.append({"role": "user", "content": question})
KonradSzafer's avatar
KonradSzafer committed
1083
1084
            # if last message is user, append to it to avoid two user messages in a row
            else:
1085
                labeled_examples[-1]["content"] += question
KonradSzafer's avatar
KonradSzafer committed
1086
1087
        else:
            # if fewshot_as_multiturn is True, append as next user entry (last is always assistant)
1088
            labeled_examples.append({"role": "user", "content": question})
Baber Abbasi's avatar
Baber Abbasi committed
1089
1090
        if gen_prefix:
            labeled_examples.append({"role": "assistant", "content": gen_prefix})
KonradSzafer's avatar
KonradSzafer committed
1091

lintangsutawika's avatar
lintangsutawika committed
1092
    @utils.positional_deprecated
KonradSzafer's avatar
KonradSzafer committed
1093
1094
    def fewshot_context(
        self,
Baber Abbasi's avatar
Baber Abbasi committed
1095
        doc: dict,
KonradSzafer's avatar
KonradSzafer committed
1096
1097
1098
1099
        num_fewshot: int,
        system_instruction: Optional[str] = None,
        apply_chat_template: bool = False,
        fewshot_as_multiturn: bool = False,
1100
        chat_template: Optional[Callable] = None,
Baber Abbasi's avatar
Baber Abbasi committed
1101
        gen_prefix: Optional[str] = None,
Baber Abbasi's avatar
Baber Abbasi committed
1102
    ) -> Union[str, List[str]]:
lintangsutawika's avatar
lintangsutawika committed
1103
1104
1105
1106
1107
1108
1109
        """Returns a fewshot context string that is made up of a prepended description
        (if provided), the `num_fewshot` number of examples, and an appended prompt example.

        :param doc: str
            The document as returned from training_docs, validation_docs, or test_docs.
        :param num_fewshot: int
            The number of fewshot examples to provide in the returned context string.
KonradSzafer's avatar
KonradSzafer committed
1110
1111
1112
1113
1114
1115
        :param  system_instruction: str
            System instruction to be applied to the prompt.
        :param apply_chat_template: bool
            Whether to apply the chat template to the fewshot context.
        :param fewshot_as_multiturn: bool
            Whether to provide the fewshot examples as a multiturn conversation or a single user turn.
1116
1117
        :param chat_template:
            callable (from lm.apply_chat_template) that takes in a list[Dict] chat transcript and renders it into a string.
1118
1119
        :param gen_prefix:
            String to append after the <|assistant|> token.
lintangsutawika's avatar
lintangsutawika committed
1120
1121
1122
        :returns: str
            The fewshot context.
        """
KonradSzafer's avatar
KonradSzafer committed
1123
1124
1125
1126
1127
1128
        if apply_chat_template:
            labeled_examples = []
        else:
            labeled_examples = ""

        # get task description
1129
1130
        if description := self.config.description:
            description = utils.apply_template(self.config.description, doc)
lintangsutawika's avatar
lintangsutawika committed
1131

KonradSzafer's avatar
KonradSzafer committed
1132
1133
1134
1135
1136
1137
1138
1139
1140
        # create system prompt based on the provided system instruction and description
        if system_instruction is not None and description:
            system_prompt = (
                f"{system_instruction}{self.sampler.fewshot_delimiter}{description}"
            )
        elif system_instruction is not None:
            system_prompt = system_instruction
        elif description:
            system_prompt = description
lintangsutawika's avatar
lintangsutawika committed
1141
        else:
KonradSzafer's avatar
KonradSzafer committed
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
1154
            system_prompt = ""

        # add system prompt if specified
        if system_prompt:
            if apply_chat_template:
                labeled_examples.append({"role": "system", "content": system_prompt})
            else:
                labeled_examples = system_prompt
        # if few-shot - append examples after the system prompt
        if num_fewshot > 0:
            if apply_chat_template:
                labeled_examples.extend(
                    self.sampler.get_chat_context(
Baber Abbasi's avatar
Baber Abbasi committed
1155
1156
1157
                        doc,
                        num_fewshot,
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1158
                        gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1159
1160
1161
                    )
                )
            else:
Baber Abbasi's avatar
Baber Abbasi committed
1162
                labeled_examples += self.sampler.get_context(
Baber Abbasi's avatar
Baber Abbasi committed
1163
                    doc, num_fewshot, gen_prefix=gen_prefix
Baber Abbasi's avatar
Baber Abbasi committed
1164
                )
lintangsutawika's avatar
lintangsutawika committed
1165
1166

        example = self.doc_to_text(doc)
KonradSzafer's avatar
KonradSzafer committed
1167
1168
        if apply_chat_template:
            if self.multiple_input:
Baber Abbasi's avatar
Baber Abbasi committed
1169
                # TODO: append prefill?
1170
1171
                if not labeled_examples:
                    return ""
1172
                return chat_template(labeled_examples)
KonradSzafer's avatar
KonradSzafer committed
1173
1174
            if isinstance(example, str):
                self.append_target_question(
Baber Abbasi's avatar
Baber Abbasi committed
1175
1176
1177
                    labeled_examples,
                    example,
                    fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1178
                    gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1179
1180
1181
1182
1183
1184
1185
                )
            # for loglikelihood create a list of questions with appended choices
            elif isinstance(example, list):
                labeled_examples_list = []
                # copy chat history for each example and append the answer
                for ex in example:
                    chat = deepcopy(labeled_examples)
Baber Abbasi's avatar
Baber Abbasi committed
1186
1187
1188
1189
                    self.append_target_question(
                        chat,
                        ex,
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1190
                        gen_prefix=gen_prefix,
Baber Abbasi's avatar
Baber Abbasi committed
1191
1192
1193
1194
1195
                    )
                    # TODO: append prefill?
                    labeled_examples_list.append(
                        chat_template(
                            chat,
Baber Abbasi's avatar
Baber Abbasi committed
1196
                            add_generation_prompt=False if gen_prefix else True,
Baber Abbasi's avatar
Baber Abbasi committed
1197
1198
                        )
                    )
KonradSzafer's avatar
KonradSzafer committed
1199
1200
1201
1202
1203
1204
                return labeled_examples_list
            # if example is an integer, append the choice or convert to string
            elif isinstance(example, int):
                if self.config.doc_to_choice is not None:
                    choices = self.doc_to_choice(doc)
                    self.append_target_question(
Baber Abbasi's avatar
Baber Abbasi committed
1205
1206
1207
                        labeled_examples,
                        choices[example],
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1208
                        gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1209
1210
1211
                    )
                else:
                    self.append_target_question(
Baber Abbasi's avatar
Baber Abbasi committed
1212
1213
1214
                        labeled_examples,
                        str(example),
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1215
                        gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1216
1217
                    )
                # return lm.apply_chat_template(labeled_examples)
Baber Abbasi's avatar
Baber Abbasi committed
1218
1219
            return chat_template(
                labeled_examples,
Baber Abbasi's avatar
Baber Abbasi committed
1220
                add_generation_prompt=False if gen_prefix else True,
Baber Abbasi's avatar
Baber Abbasi committed
1221
            )
1222
        else:
Baber Abbasi's avatar
Baber Abbasi committed
1223
            prefix = (
Baber Abbasi's avatar
Baber Abbasi committed
1224
1225
                self.config.target_delimiter + gen_prefix
                if gen_prefix is not None
Baber Abbasi's avatar
Baber Abbasi committed
1226
1227
                else ""
            )
KonradSzafer's avatar
KonradSzafer committed
1228
1229
            if self.multiple_input:
                return labeled_examples
1230
            if isinstance(example, str):
Baber Abbasi's avatar
Baber Abbasi committed
1231
                return labeled_examples + example + prefix
1232
            elif isinstance(example, list):
Baber Abbasi's avatar
Baber Abbasi committed
1233
                return [labeled_examples + ex + prefix for ex in example]
1234
1235
1236
            elif isinstance(example, int):
                if self.config.doc_to_choice is not None:
                    choices = self.doc_to_choice(doc)
Baber Abbasi's avatar
Baber Abbasi committed
1237
                    return labeled_examples + choices[example] + prefix
1238
                else:
Baber Abbasi's avatar
Baber Abbasi committed
1239
                    return labeled_examples + str(example) + prefix
lintangsutawika's avatar
lintangsutawika committed
1240

Baber Abbasi's avatar
Baber Abbasi committed
1241
    def apply_filters(self) -> Optional[List[Instance]]:
Baber Abbasi's avatar
Baber Abbasi committed
1242
        """Iterates over FilterEnsembles and applies them to instances"""
1243
1244
        if hasattr(self, "_filters"):
            for f in self._filters:
1245
                f.apply(self._instances)
1246
1247
1248
1249
        else:
            eval_logger.warning("No filter defined, passing through instances")
            return self._instances

1250
    def should_decontaminate(self):
1251
        return self.config.should_decontaminate
1252

Baber Abbasi's avatar
Baber Abbasi committed
1253
    def doc_to_decontamination_query(self, doc: dict):
1254
        if self.config.should_decontaminate:
1255
1256
            if self.config.doc_to_decontamination_query is None:
                return self.doc_to_text(doc)
1257
            else:
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
1268
                doc_to_decontamination_query = self.config.doc_to_decontamination_query
                if doc_to_decontamination_query in self.features:
                    return doc[doc_to_decontamination_query]
                elif callable(doc_to_decontamination_query):
                    return doc_to_decontamination_query(doc)
                else:
                    return ast.literal_eval(
                        utils.apply_template(
                            self.config.doc_to_decontamination_query, doc
                        )
                    )
1269

1270
    def _process_doc(self, doc: dict) -> dict:
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280
        """
        Override this to process (detokenize, strip, replace, etc.) individual
        documents. This can be used in a map over documents of a data split.
        E.g. `map(self._process_doc, self.dataset["validation"])`

        :return: dict
            The processed version of the specified `doc`.
        """
        return doc

Yu Shi Jie's avatar
Yu Shi Jie committed
1281
    def doc_to_text(self, doc, doc_to_text=None):
1282
1283
        if self.prompt is not None:
            doc_to_text = self.prompt
Yu Shi Jie's avatar
Yu Shi Jie committed
1284
1285
        elif doc_to_text is not None:
            doc_to_text = doc_to_text
1286
        else:
1287
            doc_to_text = self.config.doc_to_text
lintangsutawika's avatar
lintangsutawika committed
1288

1289
        if isinstance(doc_to_text, int):
1290
            return doc_to_text
1291
        elif isinstance(doc_to_text, str):
1292
            if doc_to_text in self.features:
1293
                # if self.config.doc_to_choice is not None:
1294
1295
                #     return self.doc_to_choice(doc)[doc[doc_to_text]]
                # else:
1296
1297
                return doc[doc_to_text]
            else:
lintangsutawika's avatar
lintangsutawika committed
1298
                text_string = utils.apply_template(doc_to_text, doc)
lintangsutawika's avatar
lintangsutawika committed
1299
                if text_string.isdigit() and self._config.doc_to_choice is not None:
lintangsutawika's avatar
lintangsutawika committed
1300
1301
1302
                    return ast.literal_eval(text_string)
                else:
                    return text_string
1303
        elif callable(doc_to_text):
1304
            return doc_to_text(doc)
lintangsutawika's avatar
lintangsutawika committed
1305
        # Used when applying a Promptsource template
1306
        elif hasattr(doc_to_text, "apply"):
1307
1308
1309
1310
1311
            applied_prompt = doc_to_text.apply(doc)
            if len(applied_prompt) == 2:
                return applied_prompt[0]
            else:
                eval_logger.warning("Applied prompt returns empty string")
1312
                return self.config.fewshot_delimiter
1313
        else:
1314
            print(type(doc_to_text))
1315
            raise TypeError
1316

Yu Shi Jie's avatar
Yu Shi Jie committed
1317
    def doc_to_target(self, doc: Mapping, doc_to_target=None) -> Union[int, str, list]:
1318
1319
        if self.prompt is not None:
            doc_to_target = self.prompt
Yu Shi Jie's avatar
Yu Shi Jie committed
1320
1321
        elif doc_to_target is not None:
            doc_to_target = doc_to_target
1322
        else:
1323
            doc_to_target = self.config.doc_to_target
1324

1325
        if isinstance(doc_to_target, int):
1326
            return doc_to_target
1327
        elif isinstance(doc_to_target, str):
1328
            if doc_to_target in self.features:
1329
                # if self.config.doc_to_choice is not None:
1330
1331
1332
                #     return self.doc_to_choice(doc)[doc[doc_to_target]]
                # else:
                return doc[doc_to_target]
1333
            else:
lintangsutawika's avatar
lintangsutawika committed
1334
                target_string = utils.apply_template(doc_to_target, doc)
lintangsutawika's avatar
lintangsutawika committed
1335
                if target_string.isdigit() and self._config.doc_to_choice is not None:
lintangsutawika's avatar
lintangsutawika committed
1336
                    return ast.literal_eval(target_string)
lintangsutawika's avatar
lintangsutawika committed
1337
1338
1339
1340
1341
                elif (
                    len(target_string) >= 2
                    and (target_string[0] == "[")
                    and (target_string[-1] == "]")
                ):
baberabb's avatar
baberabb committed
1342
1343
1344
1345
                    try:
                        return ast.literal_eval(target_string)
                    except (SyntaxError, ValueError):
                        return target_string
lintangsutawika's avatar
lintangsutawika committed
1346
1347
                else:
                    return target_string
1348
        elif isinstance(doc_to_target, list):
1349
            return doc_to_target
1350
        elif callable(doc_to_target):
1351
            return doc_to_target(doc)
lintangsutawika's avatar
lintangsutawika committed
1352
        # Used when applying a Promptsource template
1353
        elif hasattr(doc_to_target, "apply"):
1354
            applied_prompt = doc_to_target.apply(doc)
1355
1356
1357
1358
            if len(applied_prompt) == 2:
                return applied_prompt[1]
            else:
                eval_logger.warning("Applied prompt returns empty string")
1359
                return self.config.fewshot_delimiter
1360
1361
        else:
            raise TypeError
1362

Yu Shi Jie's avatar
Yu Shi Jie committed
1363
    def doc_to_choice(self, doc: Any, doc_to_choice=None) -> List[str]:
1364
1365
        if self.prompt is not None:
            doc_to_choice = self.prompt
Yu Shi Jie's avatar
Yu Shi Jie committed
1366
1367
        elif doc_to_choice is not None:
            doc_to_choice = doc_to_choice
1368
        elif self.config.doc_to_choice is None:
1369
1370
            eval_logger.error("doc_to_choice was called but not set in config")
        else:
1371
            doc_to_choice = self.config.doc_to_choice
1372

1373
        if isinstance(doc_to_choice, str):
1374
1375
1376
1377
            if doc_to_choice in self.features:
                return doc[doc_to_choice]
            else:
                return ast.literal_eval(utils.apply_template(doc_to_choice, doc))
1378
        elif isinstance(doc_to_choice, list):
1379
            return doc_to_choice
1380
        elif isinstance(doc_to_choice, dict):
1381
1382
1383
1384
1385
1386
1387
            return list(doc_to_choice.values())
        elif callable(doc_to_choice):
            return doc_to_choice(doc)
        elif hasattr(doc_to_choice, "get_answer_choices_list"):
            return doc_to_choice.get_answer_choices_list(doc)
        else:
            raise TypeError
1388

1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
    def doc_to_image(self, doc: Any, doc_to_image=None) -> Union[int, str, list]:
        if doc_to_image is not None:
            doc_to_image = doc_to_image
        elif self.config.doc_to_image is not None:
            doc_to_image = self.config.doc_to_image
        else:
            return None

        if isinstance(doc_to_image, list):
            image_feature = [
                self.doc_to_image(doc, feature) for feature in doc_to_image
            ]
            return [feature for feature in image_feature if feature is not None]
        elif isinstance(doc_to_image, str):
            if doc_to_image in self.features:
                return doc[doc_to_image]
            else:
                return ast.literal_eval(utils.apply_template(doc_to_image, doc))
        elif callable(doc_to_image):
            return doc_to_image(doc)
        else:
            return None

1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
1434
    def doc_to_audio(self, doc: Any, doc_to_audio=None) -> Union[int, str, list]:
        if doc_to_audio is not None:
            doc_to_audio = doc_to_audio
        elif self.config.doc_to_audio is not None:
            doc_to_audio = self.config.doc_to_audio
        else:
            return None

        if isinstance(doc_to_audio, list):
            audio_feature = [
                self.doc_to_audio(doc, feature) for feature in doc_to_audio
            ]
            return [feature for feature in audio_feature if feature is not None]
        elif isinstance(doc_to_audio, str):
            if doc_to_audio in self.features:
                return doc[doc_to_audio]
            else:
                return ast.literal_eval(utils.apply_template(doc_to_audio, doc))
        elif callable(doc_to_audio):
            return doc_to_audio(doc)
        else:
            return None

Baber Abbasi's avatar
Baber Abbasi committed
1435
1436
1437
1438
1439
1440
1441
1442
    def doc_to_prefix(self, doc):
        if (gen_prefix := self.config.gen_prefix) is not None:
            if gen_prefix in self.features:
                return doc[gen_prefix]
            else:
                return utils.apply_template(gen_prefix, doc)
        return None

baberabb's avatar
baberabb committed
1443
1444
1445
    def construct_requests(
        self, doc: dict, ctx: str, **kwargs
    ) -> Union[List[Instance], Instance]:
1446
        apply_chat_template = kwargs.pop("apply_chat_template", False)
1447
        chat_template: Callable | None = kwargs.pop("chat_template", None)
1448

1449
1450
        aux_arguments = None

1451
        if self.OUTPUT_TYPE == "loglikelihood":
lintangsutawika's avatar
lintangsutawika committed
1452
            arguments = (ctx, self.doc_to_target(doc))
1453
        elif self.OUTPUT_TYPE == "loglikelihood_rolling":
lintangsutawika's avatar
lintangsutawika committed
1454
            arguments = (self.doc_to_target(doc),)
1455
        elif self.OUTPUT_TYPE == "multiple_choice":
1456
            choices = self.doc_to_choice(doc)
1457
            target_delimiter = self.config.target_delimiter
1458
1459
            if apply_chat_template:
                target_delimiter = ""
1460
1461
            if self.multiple_input:
                # If there are multiple inputs, choices are placed in the ctx
1462
                # apply chat_template to choices if apply_chat_template
1463
                cont = self.doc_to_target(doc)
1464

1465
                arguments = [
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475
                    (
                        ctx
                        + (
                            chat_template([{"role": "user", "content": choice}])
                            if apply_chat_template
                            else choice
                        ),
                        f"{target_delimiter}{cont}",
                    )
                    for choice in choices
1476
                ]
1477
            else:
1478
                # Otherwise they are placed in the continuation
1479
                arguments = [(ctx, f"{target_delimiter}{cont}") for cont in choices]
1480

1481
1482
1483
1484
1485
1486
1487
1488
            # TODO: we should raise a warning telling users this will at most ~2x runtime.
            if "acc_mutual_info" in self._metric_fn_list.keys():
                # if we are calculating multiple choice accuracy
                # using mutual information instead of raw loglikelihood as metric, need unconditional lls.

                # here mutual info refers to calculating
                # log(P(choice|ctx) / P(choice)) = log(P(choice|ctx)) - log(P(choice))
                # in other words normalizing by subtracting the unconditional logprob of each choice.
1489
1490
1491
1492
                # TODO: should these be strided? will have to modify the processing in process_results if so
                aux_arguments = [
                    ("", f"{target_delimiter}{choice}") for choice in choices
                ]
1493
1494
1495
1496
1497
1498

                arguments.extend(aux_arguments)

        elif self.OUTPUT_TYPE == "generate_until":
            arguments = (ctx, deepcopy(self.config.generation_kwargs))

Baber's avatar
Baber committed
1499
1500
1501
1502
1503
1504
1505
        else:
            raise ValueError(
                f"Unsupported OUTPUT_TYPE: '{self.OUTPUT_TYPE}'. "
                f"Expected one of: 'loglikelihood', 'loglikelihood_rolling', "
                f"'multiple_choice', 'generate_until'"
            )

1506
1507
1508
1509
1510
1511
1512
1513
1514
        multimodal_arg = {}
        if (
            self.config.doc_to_image
        ):  # TODO: ensure that non-multimodal tasks aren't getting visual args
            multimodal_arg = {
                **multimodal_arg,
                **{"visual": self.doc_to_image(doc)},
            }

1515
1516
1517
1518
1519
1520
1521
1522
        if (
            self.config.doc_to_audio
        ):  # TODO: ensure that non-multimodal tasks aren't getting audio args
            multimodal_arg = {
                **multimodal_arg,
                **{"audio": self.doc_to_audio(doc)},
            }

1523
1524
1525
1526
1527
1528
1529
        if bool(multimodal_arg):
            if isinstance(arguments, list):
                arguments = [arg + (multimodal_arg,) for arg in arguments]
            else:
                arguments = arguments + (multimodal_arg,)

        if self.OUTPUT_TYPE == "multiple_choice":
1530
            request_list = [
1531
1532
                Instance(
                    request_type="loglikelihood",
lintangsutawika's avatar
lintangsutawika committed
1533
                    doc=doc,
Baber's avatar
Baber committed
1534
                    arguments=LoglikelihoodInput(context=arg[0], continuation=arg[1]),
1535
                    idx=i,
1536
1537
                    **kwargs,
                )
1538
                for i, arg in enumerate(arguments)
1539
            ]
1540
1541

            return request_list
lintangsutawika's avatar
lintangsutawika committed
1542

lintangsutawika's avatar
lintangsutawika committed
1543
        return Instance(
1544
1545
            request_type=self.OUTPUT_TYPE,
            doc=doc,
Baber's avatar
Baber committed
1546
1547
1548
            arguments=LoglikelihoodInput(*arguments)
            if self.OUTPUT_TYPE in ["loglikelihood", "loglikelihood_rolling"]
            else GenerateInput(*arguments),
1549
1550
            idx=0,
            **kwargs,
lintangsutawika's avatar
lintangsutawika committed
1551
        )
1552
1553

    def process_results(self, doc, results):
1554
1555
        if callable(self.config.process_results):
            return self.config.process_results(doc, results)
lintangsutawika's avatar
lintangsutawika committed
1556

1557
        result_dict = {}
1558
        use_metric = list(self._metric_fn_list.keys())
1559
1560
1561
        if self.OUTPUT_TYPE == "loglikelihood":
            results = results[0]
            ll, is_greedy = results
1562
1563
1564
1565
            return {
                **({"perplexity": ll} if "perplexity" in use_metric else {}),
                **({"acc": int(is_greedy)} if "acc" in use_metric else {}),
            }
1566
        elif self.OUTPUT_TYPE == "loglikelihood_rolling":
haileyschoelkopf's avatar
haileyschoelkopf committed
1567
            (loglikelihood,) = results
1568
1569
            _words = self.count_words(self.doc_to_target(doc))
            _bytes = self.count_bytes(self.doc_to_target(doc))
haileyschoelkopf's avatar
haileyschoelkopf committed
1570
            return {
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
                **(
                    {"word_perplexity": (loglikelihood, _words)}
                    if "word_perplexity" in use_metric
                    else {}
                ),
                **(
                    {"byte_perplexity": (loglikelihood, _bytes)}
                    if "byte_perplexity" in use_metric
                    else {}
                ),
                **(
                    {"bits_per_byte": (loglikelihood, _bytes)}
                    if "bits_per_byte" in use_metric
                    else {}
                ),
haileyschoelkopf's avatar
haileyschoelkopf committed
1586
            }
1587
        elif self.OUTPUT_TYPE == "multiple_choice":
1588
            lls, is_greedy = zip(*results)
lintangsutawika's avatar
lintangsutawika committed
1589

1590
            # retrieve choices in List[str] form, to compute choice lengths, etc.
1591
            choices = self.doc_to_choice(doc)
1592
1593
            completion_len = np.array([float(len(i)) for i in choices])

1594
1595
            if (
                2 * len(choices) == len(lls)
1596
                and "acc_mutual_info" in self._metric_fn_list.keys()
1597
1598
1599
            ):
                # then we are doing mutual info.
                # this stores the "dryrun" / unconditional answer loglikelihoods
1600
1601
                # as we extend the args list with unconditional ("", continuation) pairs
                lls_unconditional = lls[len(choices) :]
1602
1603
                if len(lls_unconditional) != len(choices):
                    raise ValueError
1604
                # and this stores our "regular" conditional loglikelihoods
1605
                lls = lls[: len(choices)]
1606

1607
1608
            pred = np.argmax(lls)
            pred_norm = np.argmax(lls / completion_len)
lintangsutawika's avatar
lintangsutawika committed
1609

1610
1611
            if self.multiple_input:
                gold = self.doc_to_text(doc)
1612
            else:
1613
                gold = self.doc_to_target(doc)
1614
1615

            gold_index_error = False
1616
            if isinstance(gold, list):
Lintang Sutawika's avatar
Lintang Sutawika committed
1617
1618
                gold = [i if i < len(choices) else -100 for i in gold]
                if -100 in gold:
1619
1620
                    gold_index_error = True
            else:
1621
                if isinstance(gold, int):
Lintang Sutawika's avatar
Lintang Sutawika committed
1622
                    gold = gold if gold < len(choices) else -100
1623
                elif isinstance(gold, str):
Lintang Sutawika's avatar
Lintang Sutawika committed
1624
                    gold = choices.index(gold) if gold in choices else -100
lintangsutawika's avatar
lintangsutawika committed
1625

Lintang Sutawika's avatar
Lintang Sutawika committed
1626
                if gold == -100:
1627
1628
1629
1630
                    gold_index_error = True

            if gold_index_error:
                eval_logger.warning(
lintangsutawika's avatar
lintangsutawika committed
1631
                    f"Label index was not in within range of available choices,"
1632
1633
                    f"Sample:\n\n{doc}\n\n"
                )
lintangsutawika's avatar
lintangsutawika committed
1634

1635
            if self.multiple_target:
lintangsutawika's avatar
lintangsutawika committed
1636
1637
                acc = 1.0 if pred in gold else 0.0
                acc_norm = 1.0 if pred_norm in gold else 0.0
Lintang Sutawika's avatar
Lintang Sutawika committed
1638
                exact_match = int(any([is_greedy[i] if i != -100 else 0 for i in gold]))
lintangsutawika's avatar
lintangsutawika committed
1639
1640
1641
            else:
                acc = 1.0 if pred == gold else 0.0
                acc_norm = 1.0 if pred_norm == gold else 0.0
1642
                # TODO: this gets score of 0 on arc_challenge for pythia-70m. need to test that this works properly
Lintang Sutawika's avatar
Lintang Sutawika committed
1643
                exact_match = int(is_greedy[gold]) if gold != -100 else 0
1644

Lintang Sutawika's avatar
Lintang Sutawika committed
1645
1646
1647
1648
            prob_norm = utils.softmax(lls)

            # TODO use keyword arguments to the metric?
            # gold, pred, norm stuff, the original lls,
1649
            result_dict = {
1650
                **({"acc": acc} if "acc" in use_metric else {}),
1651
1652
                **({"f1": (gold, pred)} if "f1" in use_metric else {}),
                **({"mcc": (gold, pred)} if "mcc" in use_metric else {}),
1653
                **({"acc_norm": acc_norm} if "acc_norm" in use_metric else {}),
1654
                **({"exact_match": exact_match} if "exact_match" in use_metric else {}),
Lintang Sutawika's avatar
Lintang Sutawika committed
1655
1656
1657
1658
1659
                **(
                    {"brier_score": (gold, prob_norm)}
                    if "brier_score" in use_metric
                    else {}
                ),
1660
1661
            }

1662
            if "acc_mutual_info" in use_metric:
lintangsutawika's avatar
lintangsutawika committed
1663
1664
1665
                lls_mutual_info = [
                    ll_c - ll_u for ll_c, ll_u in zip(lls, lls_unconditional)
                ]
1666
1667
1668
                acc_mutual_info = 1.0 if np.argmax(lls_mutual_info) == gold else 0.0
                result_dict["acc_mutual_info"] = acc_mutual_info

1669
        elif self.OUTPUT_TYPE == "generate_until":
1670
            gold = self.doc_to_target(doc)
Chris's avatar
Chris committed
1671
            result = results[0]
1672
            if self.config.doc_to_choice is not None:
lintangsutawika's avatar
lintangsutawika committed
1673
                # If you set doc_to_choice,
lintangsutawika's avatar
lintangsutawika committed
1674
                # it assumes that doc_to_target returns a number.
1675
1676
                choices = self.doc_to_choice(doc)
                gold = choices[gold]
1677
1678
            # we expect multiple_targets to be a list.
            elif self.multiple_target:
baberabb's avatar
baberabb committed
1679
                gold = list(gold)
Hojin Lee's avatar
Hojin Lee committed
1680
1681
1682
            # TODO: handle this better
            elif type(gold) is not type(result) and not (
                "bypass" in self._metric_fn_list.keys() or isinstance(result, list)
1683
            ):
Chris's avatar
Chris committed
1684
1685
                # cast gold to the same type as result
                gold = type(result)(gold)
1686

lintangsutawika's avatar
lintangsutawika committed
1687
            for metric in self._metric_fn_list.keys():
haileyschoelkopf's avatar
haileyschoelkopf committed
1688
1689
1690
1691
1692
                if self.multiple_target:
                    # in the case where we have multiple targets,
                    # return true if any are true
                    # TODO: this may break for multipLe_target, non zero-or-1 metrics
                    scores = []
haileyschoelkopf's avatar
haileyschoelkopf committed
1693
1694
1695
1696
                    if not isinstance(gold, list):
                        # sometimes, a multiple_target dataset has exceptions where one doc has only one string answer
                        # print(gold)
                        gold = [gold]
1697
1698
1699
1700
1701
1702
1703
1704
                    if metric == "exact_match":
                        result = [result for _ in range(len(gold))]
                        scores = self._metric_fn_list[metric](
                            references=gold,
                            predictions=result,
                            **self._metric_fn_kwargs[metric],
                        )[metric]
                        result_score = 1.0 if scores > 0.0 else 0.0
haileyschoelkopf's avatar
haileyschoelkopf committed
1705
                    else:
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
1726
                        for gold_option in gold:
                            try:
                                result_score = self._metric_fn_list[metric](
                                    references=[gold_option],
                                    predictions=[result],
                                    **self._metric_fn_kwargs[metric],
                                )
                            except (
                                TypeError
                            ):  # TODO: this is hacky and I don't want to do it
                                result_score = self._metric_fn_list[metric](
                                    [gold_option, result]
                                )
                            if isinstance(result_score, dict):
                                # TODO: this handles the case where HF evaluate returns a dict.
                                result_score = result_score[metric]
                            scores.append(result_score)
                        if any(scores):
                            result_score = 1.0
                        else:
                            result_score = 0.0
haileyschoelkopf's avatar
haileyschoelkopf committed
1727
                else:
1728
                    try:
1729
                        result_score = self._metric_fn_list[metric](
1730
1731
                            references=[gold],
                            predictions=[result],
1732
                            **self._metric_fn_kwargs[metric],
1733
                        )
1734
                    except TypeError:  # needed for now in order to use a different interface between our own metrics and HF Evaluate metrics
1735
                        result_score = self._metric_fn_list[metric]([gold, result])
1736
1737
1738
1739
1740
1741
1742
                if isinstance(result_score, dict):
                    # TODO: this handles the case where HF evaluate returns a dict.
                    # This allows for multiple metrics to be returned from the same function
                    for k, v in result_score.items():
                        result_dict[k] = v
                else:
                    result_dict[metric] = result_score
1743
        else:
lintangsutawika's avatar
lintangsutawika committed
1744
1745
            raise ValueError(
                f"Passed invalid output_type '{self.OUTPUT_TYPE}' ! Please use one of ",
1746
                "'loglikelihood', 'loglikelihood_rolling', 'generate_until' or 'multiple_choice'",
1747
            )
1748
1749
1750

        return result_dict

Baber Abbasi's avatar
Baber Abbasi committed
1751
    def aggregation(self) -> dict:
1752
1753
        return self._aggregation_list

Baber Abbasi's avatar
Baber Abbasi committed
1754
    def higher_is_better(self) -> dict:
haileyschoelkopf's avatar
haileyschoelkopf committed
1755
        return self._higher_is_better
1756

Baber Abbasi's avatar
Baber Abbasi committed
1757
1758
1759
    def get_config(self, key: str) -> Any:
        return getattr(self._config, key, None)

Lintang Sutawika's avatar
Lintang Sutawika committed
1760
1761
1762
1763
    @property
    def task_name(self) -> Any:
        return getattr(self.config, "task", None)

1764
1765
1766
1767
1768
    def __repr__(self):
        return (
            f"ConfigurableTask(task_name={getattr(self.config, 'task', None)},"
            f"output_type={self.OUTPUT_TYPE},"
            f"num_fewshot={getattr(self.config, 'num_fewshot', None)},"
Baber Abbasi's avatar
Baber Abbasi committed
1769
            f"num_samples={len(self.eval_docs)})"
1770
1771
        )

1772
    def calculate_metrics(
Baber's avatar
TODO!  
Baber committed
1773
1774
1775
1776
1777
1778
1779
1780
        self,
        instances_by_doc_id,
        filter_keys=None,
        samples=None,
        rank=1,
        limit=None,
        world_size=1,
    ) -> dict[str, list[dict]]:
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
        """Calculate metrics for all datapoints in the task.

        Args:
            instances_by_doc_id (dict): Dictionary mapping doc_ids to lists of instances.
            filter_key (str): The filter key to use for filtered responses.
            samples (dict, optional): Dictionary of sample indices to evaluate.
            rank (int): The process rank.
            limit (int, optional): Limit on number of examples to evaluate.
            world_size (int): Total number of processes.

        Returns:
            list: A list of metrics calculated for each document.
        """
Baber's avatar
TODO!  
Baber committed
1794
1795
1796
1797
1798
        if filter_keys is None:
            filter_keys = [x.name for x in self._filters]
        if isinstance(filter_keys, str):
            filter_keys = [filter_keys]
        all_metrics = collections.defaultdict(list)
1799
        # indices = samples.get(self.config.task, None) if samples is not None else None
Baber's avatar
TODO!  
Baber committed
1800
1801
1802
1803
1804
1805
1806
        for filter_key in filter_keys:
            doc_iterator = self.doc_iterator(
                rank=rank,
                limit=limit,
                world_size=world_size,
                # samples=indices,
            )
1807

Baber's avatar
TODO!  
Baber committed
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
1820
            for doc_id, doc in doc_iterator:
                # doc_id_true = indices[doc_id] if indices else doc_id
                requests = instances_by_doc_id[doc_id]

                metrics = [
                    self.process_results(doc, response)
                    for req in requests
                    for response in (
                        req.filtered_resps[filter_key]
                        if isinstance(req.filtered_resps[filter_key], list)
                        else [req.filtered_resps[filter_key]]
                    )
                ]
1821

1822
1823
1824
                all_metrics[filter_key].append(
                    MetricResult(scores=metrics, doc_id=doc_id, filter_key=filter_key)
                )
1825
1826
1827

        return all_metrics

1828
1829

class MultipleChoiceTask(Task):
1830
    OUTPUT_TYPE = "loglikelihood"
1831

baberabb's avatar
baberabb committed
1832
    def doc_to_target(self, doc: dict) -> str:
1833
1834
        return " " + doc["choices"][doc["gold"]]

baberabb's avatar
baberabb committed
1835
    def construct_requests(self, doc: dict, ctx: str, **kwargs) -> List[Instance]:
1836
        # TODO: add mutual info here?
lintangsutawika's avatar
lintangsutawika committed
1837
1838
        return [
            Instance(
haileyschoelkopf's avatar
haileyschoelkopf committed
1839
                request_type="loglikelihood",
lintangsutawika's avatar
lintangsutawika committed
1840
                doc=doc,
1841
                arguments=(ctx, " {}".format(choice)),
1842
                idx=i,
1843
1844
                **kwargs,
            )
lintangsutawika's avatar
lintangsutawika committed
1845
1846
            for i, choice in enumerate(doc["choices"])
        ]
1847

1848
    def process_results(self, doc: dict, results: Iterable[Tuple[float, bool]]) -> dict:
lintangsutawika's avatar
lintangsutawika committed
1849
1850
1851
        results = [
            res[0] for res in results
        ]  # only retain loglikelihoods, discard is_greedy TODO: do we need is_greedy anywhere?
1852
1853
1854
1855
1856
1857
1858
1859
1860
1861
1862
        gold = doc["gold"]

        acc = 1.0 if np.argmax(results) == gold else 0.0
        completion_len = np.array([float(len(i)) for i in doc["choices"]])
        acc_norm = 1.0 if np.argmax(results / completion_len) == gold else 0.0

        return {
            "acc": acc,
            "acc_norm": acc_norm,
        }

baberabb's avatar
baberabb committed
1863
    def higher_is_better(self) -> dict:
1864
1865
1866
1867
1868
        return {
            "acc": True,
            "acc_norm": True,
        }

baberabb's avatar
baberabb committed
1869
    def aggregation(self) -> dict:
1870
1871
1872
1873
1874
1875
        return {
            "acc": mean,
            "acc_norm": mean,
        }


lintangsutawika's avatar
lintangsutawika committed
1876
class PerplexityTask(Task):
Baber's avatar
Baber committed
1877
    OUTPUT_TYPE: OutputType = "loglikelihood_rolling"
1878

baberabb's avatar
baberabb committed
1879
    def has_training_docs(self) -> bool:
1880
1881
        return False

baberabb's avatar
baberabb committed
1882
    def fewshot_examples(self, k: int, rnd) -> List:
1883
1884
1885
1886
        if k != 0:
            raise ValueError(
                "The number of fewshot examples must be 0 for perplexity tasks."
            )
1887
1888
        return []

baberabb's avatar
baberabb committed
1889
    def fewshot_context(self, doc: dict, num_fewshot: int) -> Literal[""]:
1890
1891
1892
1893
        if num_fewshot != 0:
            raise ValueError(
                "The number of fewshot examples must be 0 for perplexity tasks."
            )
1894
1895
1896

        return ""

baberabb's avatar
baberabb committed
1897
    def higher_is_better(self) -> dict:
1898
1899
1900
1901
1902
1903
1904
1905
1906
        return {
            "word_perplexity": False,
            "byte_perplexity": False,
            "bits_per_byte": False,
        }

    def doc_to_decontamination_query(self, doc):
        return doc

Ethan Smith's avatar
Ethan Smith committed
1907
    def doc_to_text(self, doc) -> str:
1908
1909
1910
1911
1912
        return ""

    def doc_to_target(self, doc):
        return doc

1913
1914
1915
    def construct_requests(self, doc: dict, ctx: Optional[str], **kwargs):
        if bool(ctx):
            raise ValueError
1916

lintangsutawika's avatar
lintangsutawika committed
1917
1918
1919
1920
1921
1922
1923
        return Instance(
            request_type=self.OUTPUT_TYPE,
            doc=doc,
            arguments=(self.doc_to_target(doc),),
            idx=0,
            **kwargs,
        )
1924

1925
    def process_results(self, doc: dict, results: Tuple[float]) -> dict:
1926
        (loglikelihood,) = results
haileyschoelkopf's avatar
haileyschoelkopf committed
1927
1928
        words = self.count_words(self.doc_to_target(doc))
        bytes_ = self.count_bytes(self.doc_to_target(doc))
1929
1930
1931
1932
1933
1934
        return {
            "word_perplexity": (loglikelihood, words),
            "byte_perplexity": (loglikelihood, bytes_),
            "bits_per_byte": (loglikelihood, bytes_),
        }

baberabb's avatar
baberabb committed
1935
    def aggregation(self) -> dict:
1936
1937
1938
1939
1940
1941
1942
        return {
            "word_perplexity": weighted_perplexity,
            "byte_perplexity": weighted_perplexity,
            "bits_per_byte": bits_per_byte,
        }

    @classmethod
baberabb's avatar
baberabb committed
1943
    def count_bytes(cls, doc) -> int:
1944
1945
1946
        return len(doc.encode("utf-8"))

    @classmethod
baberabb's avatar
baberabb committed
1947
    def count_words(cls, doc) -> int:
1948
        """Downstream tasks with custom word boundaries should override this!"""
Lintang Sutawika's avatar
Lintang Sutawika committed
1949
        return len(re.split(r"\s+", doc))