Fix few-shot sampler isolation

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

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