diff --git a/src/lighteval/tasks/prompt_manager.py b/src/lighteval/tasks/prompt_manager.py index f72b8050c..475064217 100644 --- a/src/lighteval/tasks/prompt_manager.py +++ b/src/lighteval/tasks/prompt_manager.py @@ -261,7 +261,7 @@ def _sample_from_pool(self, variance_seed: int, num_fewshot: int, sampler: rando def _init_fewshot_sampling_sequential(self, num_fewshot: int, variance_seed: int): # No balancing of the few-shot examples, we take the first items of the set # We rotate by num_fewshot * seed (seed >= 0) to be able to have different series of sequential few-shots - fewshotpool = self.task.fewshot_docs() + fewshotpool = list(self.task.fewshot_docs()) for _ in range(num_fewshot * variance_seed): fewshotpool.append(fewshotpool.pop(0)) self._fewshot_cache[variance_seed] = fewshotpool # Store few shot examples @@ -282,13 +282,17 @@ def _init_fewshot_sampling_balanced( ): fewshotpool = self.task.fewshot_docs() - random.seed(variance_seed) + rnd = random.Random(variance_seed) # Build up balanced selection based on fewshot_sorting_class # (or the gold target, if the class is undefined) label_to_instances = defaultdict(list) for instance in fewshotpool: - target = instance.fewshot_sorting_class or as_list(instance.get_golds())[0] + target = ( + instance.fewshot_sorting_class + if instance.fewshot_sorting_class is not None + else as_list(instance.get_golds())[0] + ) label_to_instances[target].append(instance) # Sort by counts of class labels @@ -301,7 +305,7 @@ def _init_fewshot_sampling_balanced( for count in sorted(counts_to_labels, reverse=True): labels = counts_to_labels[count] # Break ties by randomly shuffling labels that have the same number of Instances - random.shuffle(labels) + rnd.shuffle(labels) sorted_labels.extend(labels) examples = [] @@ -311,7 +315,7 @@ def _init_fewshot_sampling_balanced( labels_iterable = cycle(sorted_labels) while num_instances_to_sample > 0: next_label = next(labels_iterable, None) - if not next_label: + if next_label is None: break instances = label_to_instances[next_label] @@ -320,7 +324,7 @@ def _init_fewshot_sampling_balanced( continue # Randomly sample without replacement - examples.append(instances.pop(random.randrange(len(instances)))) + examples.append(instances.pop(rnd.randrange(len(instances)))) num_instances_to_sample -= 1 self._fewshot_cache[variance_seed] = examples # Store few shot examples diff --git a/tests/unit/prompt/test_prompt_manager.py b/tests/unit/prompt/test_prompt_manager.py index 6a0c6a326..0b6bb6079 100644 --- a/tests/unit/prompt/test_prompt_manager.py +++ b/tests/unit/prompt/test_prompt_manager.py @@ -51,9 +51,70 @@ def test_fewshot_sampler(fewshot_select: str): match task.fewshot_selection: case "sequential": - assert docs == task.fewshot_docs()[:20] + assert docs == task.fewshot_docs()[20:40] case "random": rnd = random.Random(seed) task_docs = task.fewshot_docs() rnd.shuffle(task_docs) assert docs == task_docs[:20] + + +def test_sequential_fewshot_sampling_keeps_each_seed_independent(): + config = LightevalTaskConfig( + name="test_sequential_fewshot_task", + prompt_function=lambda _, __: None, + hf_repo="", + hf_subset="default", + metrics=[], + few_shots_split="test", + few_shots_select="sequential", + ) + task = LightevalTask(config) + task._fewshot_docs = [Doc(str(i), ["A", "B"], 0) for i in range(10)] + sampler = FewShotSampler(task) + + sampled_by_seed = {seed: [doc.query for doc in sampler.sample_fewshot_examples(2, seed)] for seed in (0, 1, 2)} + + assert sampled_by_seed == {0: ["0", "1"], 1: ["2", "3"], 2: ["4", "5"]} + assert [doc.query for doc in task.fewshot_docs()] == [str(i) for i in range(10)] + + +def test_balanced_fewshot_sampling_accepts_falsy_labels(): + config = LightevalTaskConfig( + name="test_balanced_fewshot_task", + prompt_function=lambda _, __: None, + hf_repo="", + hf_subset="default", + metrics=[], + few_shots_split="test", + few_shots_select="balanced", + ) + task = LightevalTask(config) + task._fewshot_docs = [Doc(f"empty-{i}", ["", "x"], 0) for i in range(10)] + [ + Doc(f"value-{i}", ["", "x"], 1) for i in range(10) + ] + + sampled = FewShotSampler(task).sample_fewshot_examples(4, variance_seed=0) + + assert len(sampled) == 4 + + +def test_balanced_fewshot_sampling_does_not_mutate_global_random_state(): + config = LightevalTaskConfig( + name="test_balanced_fewshot_task", + prompt_function=lambda _, __: None, + hf_repo="", + hf_subset="default", + metrics=[], + few_shots_split="test", + few_shots_select="balanced", + ) + task = LightevalTask(config) + task._fewshot_docs = [Doc(str(i), ["A", "B"], i % 2) for i in range(10)] + + random.seed(12345) + expected_next_value = random.random() + random.seed(12345) + FewShotSampler(task).sample_fewshot_examples(4, variance_seed=7) + + assert random.random() == expected_next_value