task.py 74.9 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
)
Baber's avatar
rename  
Baber committed
40
from lm_eval.api.schemas import GenerateInput, LoglikelihoodInput
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
102
103
    metadata: Optional[dict] = (
        None  # by default, not used in the code. allows for users to pass arbitrary info to tasks
    )
104

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

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

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

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

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

haileyschoelkopf's avatar
haileyschoelkopf committed
149
150
151
152
153
154
155
156
157
158
        :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)
159
160
161
162
163
164
165
166
167
168
            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
169
        return cfg_dict
170

171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
    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)

187
188
189
190
191
192
193
194
195
196
197

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)
    """

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

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

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

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

209
210
    def __init__(
        self,
211
212
213
214
        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
215
    ) -> None:
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
        """
        :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)
238
239
240
        self._training_docs: Optional[list] = None
        self._fewshot_docs: Optional[list] = None
        self._instances: Optional[List[Instance]] = None
241

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

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

249
250
251
252
253
254
    def download(
        self,
        data_dir: Optional[str] = None,
        cache_dir: Optional[str] = None,
        download_mode=None,
    ) -> None:
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
        """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.
        """
279
280
281
282
283
284
285
        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,
        )
286

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

292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
    @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

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

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

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

328
    def fewshot_docs(self) -> Iterable:
329
330
331
332
333
334
335
336
337
        """
        :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
338
339
340
341
342
            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."
                )
343
344
            return self.test_docs()

345
    def _process_doc(self, doc: dict) -> dict:
346
347
348
349
350
351
352
353
354
        """
        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
355

356
    @property
357
    def instances(self) -> List[Instance]:
358
359
360
361
362
363
364
365
366
367
368
        """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)

369
370
    def doc_to_decontamination_query(self, doc):
        raise NotImplementedError(
371
372
373
374
375
376
377
378
379
380
381
            "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

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

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

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

392
393
    def build_all_requests(
        self,
394
        *,
395
        limit: Union[int, None] = None,
396
        samples: Optional[List[int]] = None,
397
398
399
400
401
402
403
404
405
        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 = "",
406
    ) -> None:
407
        """Build a set of Instances for a task, and store them in task.instances"""
408
409
410
411

        # used with caching
        og_limit = limit

412
        cache_key = f"requests-{self._config.task}-{self.config.num_fewshot}shot-rank{rank}-world_size{world_size}"
KonradSzafer's avatar
KonradSzafer committed
413
414
415
416
417
418
419
        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 ""
        )
420
        cache_key += f"-tokenizer{tokenizer_name}"
421

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

        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
436
        eval_logger.info(f"Building contexts for {self.config.task} on rank {rank}...")
437

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

        # 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(
449
450
451
            self.doc_iterator(
                rank=rank, limit=limit, samples=samples, world_size=world_size
            )
452
453
454
455
456
457
458
        )

        num_docs = len(doc_id_docs)

        for doc_id, doc in tqdm(
            doc_id_docs,
            total=num_docs,
lintangsutawika's avatar
lintangsutawika committed
459
        ):
460
            # sample fewshot context #TODO: need to offset doc_id by rank now!
461
            fewshot_ctx = self.fewshot_context(
462
                doc,
463
464
465
466
467
468
469
                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
470
                gen_prefix=self.doc_to_prefix(doc),
471
            )
472

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

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

485
486
487
488
489
490
491
492
493
494
495
496
497
            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
498

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

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

505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
    @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
521
            The number of times each instance in a dataset is inferred on. Defaults to 1,
522
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
            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

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

haileyschoelkopf's avatar
haileyschoelkopf committed
560
561
562
563
564
565
566
567
568
569
    @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))

570
    @utils.positional_deprecated
Baber Abbasi's avatar
Baber Abbasi committed
571
    def fewshot_context(self, doc, num_fewshot, rnd=None, description=None, **kwargs):
572
573
574
575
576
577
578
        """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
579
580
581
582
583
        :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.
584
585
586
        :returns: str
            The fewshot context.
        """
587
        if rnd is None:
588
589
590
591
592
593
            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
594

595
        description = description if description else ""
596
597

        if num_fewshot == 0:
lintangsutawika's avatar
lintangsutawika committed
598
            labeled_examples = ""
599
        else:
lintangsutawika's avatar
lintangsutawika committed
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
            # 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
624
            )
625
626

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

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

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

Baber Abbasi's avatar
Baber Abbasi committed
644
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
    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)

684
685
686
687
688
    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

689
690
691
692
693
694
695
    @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:
696
697
698
            raise ValueError(
                f"Task dataset (path={self.DATASET_PATH}, name={self.DATASET_NAME}) must have valid or test docs!"
            )
699
700

    def doc_iterator(
701
702
703
704
705
706
        self,
        *,
        rank: int = 0,
        limit: Union[int, None] = None,
        world_size: int = 1,
        samples: Optional[List[int]] = None,
707
    ) -> Iterator[Tuple[int, Any]]:
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
729
        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),
            )
730
731
        return doc_iterator

732
733

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

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

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

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

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

765
        if self.config.output_type is not None:
766
767
768
769
            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)}'"
                )
770
            self.OUTPUT_TYPE = self.config.output_type
771

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

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

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

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

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

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

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

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

823
                if self.config.process_results is not None:
824
825
                    self._metric_fn_list[metric_name] = None
                    self._metric_fn_kwargs[metric_name] = {}
826
827
828
829
830
831
                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
832
833
834
                    self._metric_fn_list[metric_name] = get_metric(
                        metric_name, hf_evaluate_metric
                    )
835
                    self._metric_fn_kwargs[metric_name] = kwargs
lintangsutawika's avatar
lintangsutawika committed
836

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

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

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

871
        if self.config.filter_list is not None:
lintangsutawika's avatar
lintangsutawika committed
872
            self._filters = []
873
            for filter_config in self.config.filter_list:
874
875
876
877
878
879
880
881
882
                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
883
                self._filters.append(filter_pipeline)
lintangsutawika's avatar
lintangsutawika committed
884
        else:
Baber Abbasi's avatar
Baber Abbasi committed
885
886
887
888
            # 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."
            )
889
            # self._filters = [build_filter_ensemble("none", [["take_first", None]])]
890

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

lintangsutawika's avatar
lintangsutawika committed
899
        if self.fewshot_docs() is not None:
900
901
902
903
            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
904
905
906
                self.config.fewshot_config.get("sampler", "default")
                if self.config.fewshot_config
                else "default"
907
908
909
910
911
912
913
914
915
916
917
918
919
920
921
922
            )
            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)}"
                )
923

924
        self.task_docs = self.eval_docs
925

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

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

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

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

960
961
962
        if test_choice is not None:
            check_choices = test_choice
        else:
lintangsutawika's avatar
lintangsutawika committed
963
            check_choices = [test_target]
964
965
966
967
        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 = (
968
969
                    True
                    if self.config.target_delimiter.rstrip()
970
                    != self.config.target_delimiter
971
                    else False
972
                )
973

974
                if delimiter_has_whitespace and choice_has_whitespace:
975
976
                    eval_logger.debug(
                        f'Both target_delimiter "{self.config.target_delimiter}" and target choice: "{choice}" have whitespace'
977
978
                    )
                elif (not delimiter_has_whitespace) and (not choice_has_whitespace):
979
                    eval_logger.debug(
980
                        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'
981
982
                    )

Baber Abbasi's avatar
Baber Abbasi committed
983
984
985
986
987
988
989
990
991
992
993
994
995
996
997
998
999
    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 {},
            )
1000

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

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

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

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

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

baberabb's avatar
baberabb committed
1035
    def test_docs(self) -> datasets.Dataset:
1036
        if self.has_test_docs():
1037
1038
1039
            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]
1040

1041
    def fewshot_docs(self):
1042
        if self.config.fewshot_split is not None:
1043
1044
            if self.config.process_docs is not None:
                return self.config.process_docs(self.dataset[self.config.fewshot_split])
1045
            return self.dataset[self.config.fewshot_split]
1046
1047
1048
1049
1050
1051
1052
1053
1054
1055
1056
1057
        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."
                )
1058
        else:
1059
            if (self.config.num_fewshot is not None) and (self.config.num_fewshot > 0):
1060
                eval_logger.warning(
Lintang Sutawika's avatar
Lintang Sutawika committed
1061
                    f"[Task: {self.config.task}] "
1062
1063
1064
1065
                    "num_fewshot > 0 but fewshot_split is None. "
                    "using preconfigured rule."
                )
            return super().fewshot_docs()
1066

KonradSzafer's avatar
KonradSzafer committed
1067
1068
1069
1070
1071
    @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
1072
        gen_prefix: Optional[str] = None,
KonradSzafer's avatar
KonradSzafer committed
1073
1074
1075
1076
1077
1078
1079
1080
    ) -> 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":
1081
                labeled_examples.append({"role": "user", "content": question})
KonradSzafer's avatar
KonradSzafer committed
1082
1083
            # if last message is user, append to it to avoid two user messages in a row
            else:
1084
                labeled_examples[-1]["content"] += question
KonradSzafer's avatar
KonradSzafer committed
1085
1086
        else:
            # if fewshot_as_multiturn is True, append as next user entry (last is always assistant)
1087
            labeled_examples.append({"role": "user", "content": question})
Baber Abbasi's avatar
Baber Abbasi committed
1088
1089
        if gen_prefix:
            labeled_examples.append({"role": "assistant", "content": gen_prefix})
KonradSzafer's avatar
KonradSzafer committed
1090

lintangsutawika's avatar
lintangsutawika committed
1091
    @utils.positional_deprecated
KonradSzafer's avatar
KonradSzafer committed
1092
1093
    def fewshot_context(
        self,
Baber Abbasi's avatar
Baber Abbasi committed
1094
        doc: dict,
KonradSzafer's avatar
KonradSzafer committed
1095
1096
1097
1098
        num_fewshot: int,
        system_instruction: Optional[str] = None,
        apply_chat_template: bool = False,
        fewshot_as_multiturn: bool = False,
1099
        chat_template: Optional[Callable] = None,
Baber Abbasi's avatar
Baber Abbasi committed
1100
        gen_prefix: Optional[str] = None,
Baber Abbasi's avatar
Baber Abbasi committed
1101
    ) -> Union[str, List[str]]:
lintangsutawika's avatar
lintangsutawika committed
1102
1103
1104
1105
1106
1107
1108
        """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
1109
1110
1111
1112
1113
1114
        :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.
1115
1116
        :param chat_template:
            callable (from lm.apply_chat_template) that takes in a list[Dict] chat transcript and renders it into a string.
1117
1118
        :param gen_prefix:
            String to append after the <|assistant|> token.
lintangsutawika's avatar
lintangsutawika committed
1119
1120
1121
        :returns: str
            The fewshot context.
        """
KonradSzafer's avatar
KonradSzafer committed
1122
1123
1124
1125
1126
1127
        if apply_chat_template:
            labeled_examples = []
        else:
            labeled_examples = ""

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

KonradSzafer's avatar
KonradSzafer committed
1131
1132
1133
1134
1135
1136
1137
1138
1139
        # 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
1140
        else:
KonradSzafer's avatar
KonradSzafer committed
1141
1142
1143
1144
1145
1146
1147
1148
1149
1150
1151
1152
1153
            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
1154
1155
1156
                        doc,
                        num_fewshot,
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1157
                        gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1158
1159
1160
                    )
                )
            else:
Baber Abbasi's avatar
Baber Abbasi committed
1161
                labeled_examples += self.sampler.get_context(
Baber Abbasi's avatar
Baber Abbasi committed
1162
                    doc, num_fewshot, gen_prefix=gen_prefix
Baber Abbasi's avatar
Baber Abbasi committed
1163
                )
lintangsutawika's avatar
lintangsutawika committed
1164
1165

        example = self.doc_to_text(doc)
KonradSzafer's avatar
KonradSzafer committed
1166
1167
        if apply_chat_template:
            if self.multiple_input:
Baber Abbasi's avatar
Baber Abbasi committed
1168
                # TODO: append prefill?
1169
1170
                if not labeled_examples:
                    return ""
1171
                return chat_template(labeled_examples)
KonradSzafer's avatar
KonradSzafer committed
1172
1173
            if isinstance(example, str):
                self.append_target_question(
Baber Abbasi's avatar
Baber Abbasi committed
1174
1175
1176
                    labeled_examples,
                    example,
                    fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1177
                    gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1178
1179
1180
1181
1182
1183
1184
                )
            # 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
1185
1186
1187
1188
                    self.append_target_question(
                        chat,
                        ex,
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1189
                        gen_prefix=gen_prefix,
Baber Abbasi's avatar
Baber Abbasi committed
1190
1191
1192
1193
1194
                    )
                    # TODO: append prefill?
                    labeled_examples_list.append(
                        chat_template(
                            chat,
Baber Abbasi's avatar
Baber Abbasi committed
1195
                            add_generation_prompt=False if gen_prefix else True,
Baber Abbasi's avatar
Baber Abbasi committed
1196
1197
                        )
                    )
KonradSzafer's avatar
KonradSzafer committed
1198
1199
1200
1201
1202
1203
                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
1204
1205
1206
                        labeled_examples,
                        choices[example],
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1207
                        gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1208
1209
1210
                    )
                else:
                    self.append_target_question(
Baber Abbasi's avatar
Baber Abbasi committed
1211
1212
1213
                        labeled_examples,
                        str(example),
                        fewshot_as_multiturn,
Baber Abbasi's avatar
Baber Abbasi committed
1214
                        gen_prefix=gen_prefix,
KonradSzafer's avatar
KonradSzafer committed
1215
1216
                    )
                # return lm.apply_chat_template(labeled_examples)
Baber Abbasi's avatar
Baber Abbasi committed
1217
1218
            return chat_template(
                labeled_examples,
Baber Abbasi's avatar
Baber Abbasi committed
1219
                add_generation_prompt=False if gen_prefix else True,
Baber Abbasi's avatar
Baber Abbasi committed
1220
            )
1221
        else:
Baber Abbasi's avatar
Baber Abbasi committed
1222
            prefix = (
Baber Abbasi's avatar
Baber Abbasi committed
1223
1224
                self.config.target_delimiter + gen_prefix
                if gen_prefix is not None
Baber Abbasi's avatar
Baber Abbasi committed
1225
1226
                else ""
            )
KonradSzafer's avatar
KonradSzafer committed
1227
1228
            if self.multiple_input:
                return labeled_examples
1229
            if isinstance(example, str):
Baber Abbasi's avatar
Baber Abbasi committed
1230
                return labeled_examples + example + prefix
1231
            elif isinstance(example, list):
Baber Abbasi's avatar
Baber Abbasi committed
1232
                return [labeled_examples + ex + prefix for ex in example]
1233
1234
1235
            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
1236
                    return labeled_examples + choices[example] + prefix
1237
                else:
Baber Abbasi's avatar
Baber Abbasi committed
1238
                    return labeled_examples + str(example) + prefix
lintangsutawika's avatar
lintangsutawika committed
1239

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

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

Baber Abbasi's avatar
Baber Abbasi committed
1252
    def doc_to_decontamination_query(self, doc: dict):
1253
        if self.config.should_decontaminate:
1254
1255
            if self.config.doc_to_decontamination_query is None:
                return self.doc_to_text(doc)
1256
            else:
1257
1258
1259
1260
1261
1262
1263
1264
1265
1266
1267
                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
                        )
                    )
1268

1269
    def _process_doc(self, doc: dict) -> dict:
1270
1271
1272
1273
1274
1275
1276
1277
1278
1279
        """
        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
1280
    def doc_to_text(self, doc, doc_to_text=None):
1281
1282
        if self.prompt is not None:
            doc_to_text = self.prompt
Yu Shi Jie's avatar
Yu Shi Jie committed
1283
1284
        elif doc_to_text is not None:
            doc_to_text = doc_to_text
1285
        else:
1286
            doc_to_text = self.config.doc_to_text
lintangsutawika's avatar
lintangsutawika committed
1287

1288
        if isinstance(doc_to_text, int):
1289
            return doc_to_text
1290
        elif isinstance(doc_to_text, str):
1291
            if doc_to_text in self.features:
1292
                # if self.config.doc_to_choice is not None:
1293
1294
                #     return self.doc_to_choice(doc)[doc[doc_to_text]]
                # else:
1295
1296
                return doc[doc_to_text]
            else:
lintangsutawika's avatar
lintangsutawika committed
1297
                text_string = utils.apply_template(doc_to_text, doc)
lintangsutawika's avatar
lintangsutawika committed
1298
                if text_string.isdigit() and self._config.doc_to_choice is not None:
lintangsutawika's avatar
lintangsutawika committed
1299
1300
1301
                    return ast.literal_eval(text_string)
                else:
                    return text_string
1302
        elif callable(doc_to_text):
1303
            return doc_to_text(doc)
lintangsutawika's avatar
lintangsutawika committed
1304
        # Used when applying a Promptsource template
1305
        elif hasattr(doc_to_text, "apply"):
1306
1307
1308
1309
1310
            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")
1311
                return self.config.fewshot_delimiter
1312
        else:
1313
            print(type(doc_to_text))
1314
            raise TypeError
1315

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

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

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

1372
        if isinstance(doc_to_choice, str):
1373
1374
1375
1376
            if doc_to_choice in self.features:
                return doc[doc_to_choice]
            else:
                return ast.literal_eval(utils.apply_template(doc_to_choice, doc))
1377
        elif isinstance(doc_to_choice, list):
1378
            return doc_to_choice
1379
        elif isinstance(doc_to_choice, dict):
1380
1381
1382
1383
1384
1385
1386
            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
1387

1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
    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

1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427
1428
1429
1430
1431
1432
1433
    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
1434
1435
1436
1437
1438
1439
1440
1441
    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
1442
1443
1444
    def construct_requests(
        self, doc: dict, ctx: str, **kwargs
    ) -> Union[List[Instance], Instance]:
1445
        apply_chat_template = kwargs.pop("apply_chat_template", False)
1446
        chat_template: Callable | None = kwargs.pop("chat_template", None)
1447

1448
1449
        aux_arguments = None

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

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

1480
1481
1482
1483
1484
1485
1486
1487
            # 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.
1488
1489
1490
1491
                # 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
                ]
1492
1493
1494
1495
1496
1497

                arguments.extend(aux_arguments)

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

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

1505
1506
1507
1508
1509
1510
1511
1512
1513
        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)},
            }

1514
1515
1516
1517
1518
1519
1520
1521
        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)},
            }

1522
1523
1524
1525
1526
1527
1528
        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":
1529
            request_list = [
1530
1531
                Instance(
                    request_type="loglikelihood",
lintangsutawika's avatar
lintangsutawika committed
1532
                    doc=doc,
Baber's avatar
Baber committed
1533
                    arguments=LoglikelihoodInput(context=arg[0], continuation=arg[1]),
1534
                    idx=i,
1535
1536
                    **kwargs,
                )
1537
                for i, arg in enumerate(arguments)
1538
            ]
1539
1540

            return request_list
lintangsutawika's avatar
lintangsutawika committed
1541

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

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

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

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

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

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

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

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

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

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

1634
            if self.multiple_target:
lintangsutawika's avatar
lintangsutawika committed
1635
1636
                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
1637
                exact_match = int(any([is_greedy[i] if i != -100 else 0 for i in gold]))
lintangsutawika's avatar
lintangsutawika committed
1638
1639
1640
            else:
                acc = 1.0 if pred == gold else 0.0
                acc_norm = 1.0 if pred_norm == gold else 0.0
1641
                # 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
1642
                exact_match = int(is_greedy[gold]) if gold != -100 else 0
1643

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

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

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

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

lintangsutawika's avatar
lintangsutawika committed
1686
            for metric in self._metric_fn_list.keys():
haileyschoelkopf's avatar
haileyschoelkopf committed
1687
1688
1689
1690
1691
                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
1692
1693
1694
1695
                    if not isinstance(gold, list):
                        # sometimes, a multiple_target dataset has exceptions where one doc has only one string answer
                        # print(gold)
                        gold = [gold]
1696
1697
1698
1699
1700
1701
1702
1703
                    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
1704
                    else:
1705
1706
1707
1708
1709
1710
1711
1712
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
1725
                        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
1726
                else:
1727
                    try:
1728
                        result_score = self._metric_fn_list[metric](
1729
1730
                            references=[gold],
                            predictions=[result],
1731
                            **self._metric_fn_kwargs[metric],
1732
                        )
1733
                    except TypeError:  # needed for now in order to use a different interface between our own metrics and HF Evaluate metrics
1734
                        result_score = self._metric_fn_list[metric]([gold, result])
1735
1736
1737
1738
1739
1740
1741
                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
1742
        else:
lintangsutawika's avatar
lintangsutawika committed
1743
1744
            raise ValueError(
                f"Passed invalid output_type '{self.OUTPUT_TYPE}' ! Please use one of ",
1745
                "'loglikelihood', 'loglikelihood_rolling', 'generate_until' or 'multiple_choice'",
1746
            )
1747
1748
1749

        return result_dict

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

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

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

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

1763
1764
1765
1766
1767
    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
1768
            f"num_samples={len(self.eval_docs)})"
1769
1770
        )

1771
    def calculate_metrics(
Baber's avatar
TODO!  
Baber committed
1772
1773
1774
1775
1776
1777
1778
1779
        self,
        instances_by_doc_id,
        filter_keys=None,
        samples=None,
        rank=1,
        limit=None,
        world_size=1,
    ) -> dict[str, list[dict]]:
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
        """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
1793
1794
1795
1796
1797
        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)
1798
        # indices = samples.get(self.config.task, None) if samples is not None else None
Baber's avatar
TODO!  
Baber committed
1799
1800
1801
1802
1803
1804
1805
        for filter_key in filter_keys:
            doc_iterator = self.doc_iterator(
                rank=rank,
                limit=limit,
                world_size=world_size,
                # samples=indices,
            )
1806

Baber's avatar
TODO!  
Baber committed
1807
1808
1809
1810
1811
1812
1813
1814
1815
1816
1817
1818
1819
            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]]
                    )
                ]
1820

Baber's avatar
TODO!  
Baber committed
1821
                all_metrics[filter_key].append(metrics)
1822
1823
1824

        return all_metrics

1825
1826

class MultipleChoiceTask(Task):
1827
    OUTPUT_TYPE = "loglikelihood"
1828

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

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

1845
    def process_results(self, doc: dict, results: Iterable[Tuple[float, bool]]) -> dict:
lintangsutawika's avatar
lintangsutawika committed
1846
1847
1848
        results = [
            res[0] for res in results
        ]  # only retain loglikelihoods, discard is_greedy TODO: do we need is_greedy anywhere?
1849
1850
1851
1852
1853
1854
1855
1856
1857
1858
1859
        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
1860
    def higher_is_better(self) -> dict:
1861
1862
1863
1864
1865
        return {
            "acc": True,
            "acc_norm": True,
        }

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


lintangsutawika's avatar
lintangsutawika committed
1873
class PerplexityTask(Task):
Baber's avatar
Baber committed
1874
    OUTPUT_TYPE: OutputType = "loglikelihood_rolling"
1875

baberabb's avatar
baberabb committed
1876
    def has_training_docs(self) -> bool:
1877
1878
        return False

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

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

        return ""

baberabb's avatar
baberabb committed
1894
    def higher_is_better(self) -> dict:
1895
1896
1897
1898
1899
1900
1901
1902
1903
        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
1904
    def doc_to_text(self, doc) -> str:
1905
1906
1907
1908
1909
        return ""

    def doc_to_target(self, doc):
        return doc

1910
1911
1912
    def construct_requests(self, doc: dict, ctx: Optional[str], **kwargs):
        if bool(ctx):
            raise ValueError
1913

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

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

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

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

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