instance.py 943 Bytes
Newer Older
1
from dataclasses import dataclass, field
2
from typing import Literal, Tuple
3

lintangsutawika's avatar
lintangsutawika committed
4

5
6
@dataclass
class Instance:
lintangsutawika's avatar
lintangsutawika committed
7
8
9
    request_type: str = Literal[
        "loglikelihood", "loglikelihood_rolling", "greedy_until"
    ]
10
11
    doc: dict = None
    arguments: tuple = None
12
    idx: int = None
lintangsutawika's avatar
lintangsutawika committed
13
    metadata: tuple = Tuple[str, int, int]  # TODO: better typehints here
14
15
16
    resps: list = field(default_factory=list)
    filtered_resps: dict = field(default_factory=dict)

17
    # initialized after init
18
19
20
21
22
    task_name: str = None
    doc_id: str = None
    repeats: str = None

    def __post_init__(self):
haileyschoelkopf's avatar
haileyschoelkopf committed
23
        # unpack metadata field
24
        self.task_name, self.doc_id, self.repeats = self.metadata
lintangsutawika's avatar
lintangsutawika committed
25

26
27
28
29
30
    @property
    def args(self):
        """
        Returns (string,) where `string` is the string to calculate loglikelihood over
        """
lintangsutawika's avatar
lintangsutawika committed
31
32
33
        return (
            self.arguments if isinstance(self.arguments, tuple) else (self.arguments,)
        )