Add support for selecting few-shot examples by id

src/lighteval/tasks/lighteval_task.py

@@ -79,6 +79,14 @@ class LightevalTaskConfig:
             from. Defaults to None.
         few_shots_select (str | None, optional): Method for selecting few-shot
             examples. Defaults to None.
+        few_shots_id_column (str | None, optional): Name of a column in the dataset
+            that provides a stable identifier for each row. When set, this value is
+            used as the `id` of the corresponding few-shot `Doc` instead of the row's
+            position in the split. Defaults to None (the row's position is used).
+        few_shots_id_list (ListLike[str] | None, optional): Explicit list of ids
+            (matching `few_shots_id_column`, or the row position if it is not set)
+            to use as few-shot examples, in the given order, instead of applying
+            `few_shots_select`. Defaults to None.
 
     Generation Parameters:
         generation_size (int | None, optional): Maximum token length for generated
@@ -133,6 +141,8 @@ class LightevalTaskConfig:
     evaluation_splits: ListLike[str] = field(default_factory=lambda: ["validation"])
     few_shots_split: str | None = None
     few_shots_select: str | None = None
+    few_shots_id_column: str | None = None
+    few_shots_id_list: ListLike[str] | None = None
 
     # Generation args
     generation_size: int | None = None
@@ -157,6 +167,7 @@ def __post_init__(self):
         self.hf_avail_splits = tuple(self.hf_avail_splits)
         self.evaluation_splits = tuple(self.evaluation_splits)
         self.stop_sequence = self.stop_sequence if self.stop_sequence is not None else ()
+        self.few_shots_id_list = tuple(str(x) for x in self.few_shots_id_list) if self.few_shots_id_list else None
         self.full_name = f"{self.name}|{self.num_fewshots}"  # todo clefourrier: this is likely incorrect
 
     def __str__(self, lite: bool = False):  # noqa: C901
@@ -235,6 +246,8 @@ def __init__(
             config.hf_avail_splits or []
         )
         self.fewshot_selection = config.few_shots_select
+        self.few_shots_id_column = config.few_shots_id_column
+        self.few_shots_id_list = config.few_shots_id_list
         self.must_remove_duplicate_docs = config.must_remove_duplicate_docs
 
         self.formatter = config.prompt_function
@@ -311,7 +324,10 @@ def _get_docs_from_split(self, splits: list[str], few_shots=False) -> list[Doc]:
                 if doc is None or doc == []:
                     continue
 
-                doc.id = str(ix)
+                if few_shots and self.few_shots_id_column:
+                    doc.id = str(item[self.few_shots_id_column])
+                else:
+                    doc.id = str(ix)
 
                 # Transfer task-level generation parameters to the document
                 doc.generation_grammar = self.generation_grammar

src/lighteval/tasks/prompt_manager.py

@@ -211,6 +211,7 @@ def __init__(self, task: "LightevalTask"):
 
         self.few_shots_select = FewShotSelection[few_shots_select]
         self.few_shots_split = task.fewshot_split
+        self.few_shots_id_list = getattr(task, "few_shots_id_list", None)
 
         self._fewshot_cache = {}
 
@@ -237,7 +238,9 @@ def _init_fewshot_pool(
     ):
         # If there is no cache, we initialize it
         if variance_seed not in self._fewshot_cache:
-            if self.few_shots_select.value.sorting == "sequential":
+            if self.few_shots_id_list:
+                self._init_fewshot_sampling_by_id(variance_seed=variance_seed)
+            elif self.few_shots_select.value.sorting == "sequential":
                 self._init_fewshot_sampling_sequential(num_fewshot=num_fewshot, variance_seed=variance_seed)
             elif self.few_shots_select.value.sorting == "random":
                 self._init_fewshot_sampling_random(variance_seed=variance_seed)
@@ -247,6 +250,10 @@ def _init_fewshot_pool(
                 raise Exception("No correct few shot strategy selected - but this point should not be reachable.")
 
     def _sample_from_pool(self, variance_seed: int, num_fewshot: int, sampler: random.Random) -> list:
+        if self.few_shots_id_list:
+            # The pool was already built in the exact requested order; no further
+            # sampling strategy should be applied on top of an explicit id list.
+            return self._fewshot_cache[variance_seed]
         if self.few_shots_select.value.with_sampling and sampler is not None:
             if self.few_shots_select.value.fewshotpool_unique:
                 # This functionality is here for compatibility with the harness few shot system.
@@ -258,6 +265,19 @@ def _sample_from_pool(self, variance_seed: int, num_fewshot: int, sampler: rando
         else:
             return self._fewshot_cache[variance_seed]
 
+    def _init_fewshot_sampling_by_id(self, variance_seed: int):
+        fewshotpool = self.task.fewshot_docs()
+        pool_by_id = {doc.id: doc for doc in fewshotpool}
+
+        missing_ids = [doc_id for doc_id in self.few_shots_id_list if doc_id not in pool_by_id]
+        if missing_ids:
+            raise ValueError(
+                f"few_shots_id_list for task {self.task.name} references ids that are not present "
+                f"in the few-shot pool: {missing_ids}"
+            )
+
+        self._fewshot_cache[variance_seed] = [pool_by_id[doc_id] for doc_id in self.few_shots_id_list]
+
     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

tests/unit/prompt/test_prompt_manager.py

@@ -57,3 +57,41 @@ def test_fewshot_sampler(fewshot_select: str):
             task_docs = task.fewshot_docs()
             rnd.shuffle(task_docs)
             assert docs == task_docs[:20]
+
+
+def test_fewshot_sampler_with_id_list():
+    config = LightevalTaskConfig(
+        name="test_fewshot_id_list_task",
+        prompt_function=lambda _, __: None,
+        hf_repo="",
+        hf_subset="default",
+        metrics=[],
+        few_shots_split="test",
+        few_shots_id_list=["7", "2", "5"],
+    )
+    task = LightevalTask(config)
+    task._fewshot_docs = [Doc(query=str(i), choices=["A", "B"], gold_index=0, id=str(i)) for i in range(10)]
+    sampler = FewShotSampler(task)
+
+    docs = sampler.sample_fewshot_examples(3, variance_seed=0)
+
+    # The requested ids are honored, in the exact order they were given.
+    assert [doc.id for doc in docs] == ["7", "2", "5"]
+
+
+def test_fewshot_sampler_with_id_list_raises_on_missing_id():
+    config = LightevalTaskConfig(
+        name="test_fewshot_id_list_missing_task",
+        prompt_function=lambda _, __: None,
+        hf_repo="",
+        hf_subset="default",
+        metrics=[],
+        few_shots_split="test",
+        few_shots_id_list=["1", "does-not-exist"],
+    )
+    task = LightevalTask(config)
+    task._fewshot_docs = [Doc(query=str(i), choices=["A", "B"], gold_index=0, id=str(i)) for i in range(3)]
+    sampler = FewShotSampler(task)
+
+    with pytest.raises(ValueError, match="does-not-exist"):
+        sampler.sample_fewshot_examples(2, variance_seed=0)

tests/unit/tasks/test_lighteval_task.py

@@ -20,6 +20,7 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
 
+import json
 
 from lighteval.tasks.lighteval_task import LightevalTask, LightevalTaskConfig
 from lighteval.tasks.requests import Doc
@@ -84,3 +85,30 @@ def test_hf_data_files(tmp_path):
 
     eval_docs = task.eval_docs()
     assert [doc.query for doc in eval_docs] == src_docs
+
+
+def test_few_shots_id_column_sets_doc_id_for_fewshot_docs(tmp_path):
+    # create a small jsonl dataset where each row has a stable, non-positional identifier
+    data_file = tmp_path / "data.jsonl"
+    rows = [{"text": f"document {i}", "row_id": f"custom-{i}"} for i in range(3)]
+    data_file.write_text("\n".join(json.dumps(row) for row in rows))
+
+    cfg = LightevalTaskConfig(
+        name="test_few_shots_id_column",
+        prompt_function=dummy_prompt_function,
+        hf_repo="json",
+        hf_subset="default",
+        metrics=[],
+        evaluation_splits=["train"],
+        few_shots_split="train",
+        few_shots_id_column="row_id",
+        hf_data_files=str(data_file),
+    )
+    task = LightevalTask(cfg)
+
+    fewshot_docs = task.fewshot_docs()
+    assert [doc.id for doc in fewshot_docs] == ["custom-0", "custom-1", "custom-2"]
+
+    # Evaluation docs are unaffected: they keep using the row's position as id.
+    eval_docs = task.eval_docs()
+    assert [doc.id for doc in eval_docs] == ["0", "1", "2"]