samplers.py 3.05 KB
Newer Older
1
class Sampler:
2
3
4
5
6
7
8
9
10
11
    def __init__(self, docs, task, fewshot_indices=None, rnd=None):

        self.rnd = rnd
        assert self.rnd, "must pass rnd to FewShotSampler!"

        self.task = task
        self.config = task._config

        self.delimiter = self.config.delimiter

lintangsutawika's avatar
lintangsutawika committed
12
13
        self.docs = docs  # HF dataset split, provided by task._fewshot_docs()
        if fewshot_indices:  # subset few-shot docs from
14
15
16
17
            self.docs = self.docs.select(fewshot_indices)

    def get_context(self, doc, num_fewshot):

lintangsutawika's avatar
lintangsutawika committed
18
19
20
21
22
23
        # draw an extra fewshot sample if using same split as evaluating on
        n_samples = (
            num_fewshot + 1
            if self.config.fewshot_split == self.config.test_split
            else num_fewshot
        )
24

25
        # draw `n_samples` docs from fewshot_docs
26
27
28
        fewshotex = self.sample(n_samples)

        # get rid of the doc that's the one we're evaluating, if it's in the fewshot
29
        # TODO: should we just stop people from using fewshot from same split as evaluating?
30
        selected_docs = [x for x in fewshotex if x != doc][:num_fewshot]
lintangsutawika's avatar
lintangsutawika committed
31

32
        labeled_examples = (
lintangsutawika's avatar
lintangsutawika committed
33
34
            self.delimiter.join(
                [
35
36
                    # TODO: is separating doc_to_text and doc_to_target by one space always desired?
                    self.task.doc_to_text(doc) + " " + self.task.doc_to_target(doc)
lintangsutawika's avatar
lintangsutawika committed
37
38
                    for doc in selected_docs
                ]
39
            )
lintangsutawika's avatar
lintangsutawika committed
40
41
            + self.delimiter
        )
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56

        # only returns the fewshot context! Does not append the document, do this outside the object
        return labeled_examples

    def sample(self, n):
        """
        Draw `n` samples from our fewshot docs. This method should be overridden by subclasses.
        """

        return self.rnd.sample(self.docs, n)


class BalancedSampler(Sampler):
    def sample(self, n):
        """
lintangsutawika's avatar
lintangsutawika committed
57
        TODO: this should return approximately class-balanced samples from our fewshot examples.
58
        TODO: what order should they be in? maybe random?
59
60
61
62
63
        """

        pass


lintangsutawika's avatar
lintangsutawika committed
64
class ManualSampler(Sampler):
65
    def sample(self, n):
lintangsutawika's avatar
lintangsutawika committed
66
67
        """ """
        pass
68
69


lintangsutawika's avatar
lintangsutawika committed
70
# TODO: how should we do design here? might be better to have a single sampler and pass more kwargs at init.
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
# Depends what's easier for new user to add own functionality on top of

# types of sampler:
# - class-balanced, randomly shuffled
# - class-balanced, one particular set of fewshot examples for all evaled instances
# - hand-specify number of fewshot examples per class?
# - random, varies per example (check that this is curr. default in old repo)
# - random, unified per example
# - enforce a specific fixed fewshot string! (or should we not use this, in favor of including it in prompt template directly)


# - user-specified doc indices to restrict fewshot doc options to
# - user specifies split to use for drawing fewshot instances (TODO: manually prevent this from being same split you eval!)
# - user specifies a prepended "description"/string to add in front of the (prompted) input

# - user specifies a location to draw fewshot samples from? DO THIS IN TASK CLASS