feat: add ID-based few-shot example selection (#634)

src/lighteval/tasks/lighteval_task.py

@@ -134,6 +134,11 @@ class LightevalTaskConfig:
     few_shots_split: str | None = None
     few_shots_select: str | None = None
 
+    # ID-based few-shot selection: specify exact dataset rows to use as few-shot examples.
+    # See: https://github.com/huggingface/lighteval/issues/634
+    few_shots_id_column: str | None = None
+    few_shots_id_list: ListLike[str] | None = None
+
     # Generation args
     generation_size: int | None = None
     generation_grammar: TextGenerationInputGrammarType | None = None
@@ -235,6 +240,8 @@ def __init__(
             config.hf_avail_splits or []
         )
         self.fewshot_selection = config.few_shots_select
+        self.fewshot_id_column = config.few_shots_id_column
+        self.fewshot_id_list = config.few_shots_id_list
         self.must_remove_duplicate_docs = config.must_remove_duplicate_docs
 
         self.formatter = config.prompt_function

src/lighteval/tasks/prompt_manager.py

@@ -237,7 +237,10 @@ 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":
+            # ID-based selection takes priority when configured
+            if self.task.fewshot_id_list is not None:
+                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)
@@ -325,6 +328,52 @@ def _init_fewshot_sampling_balanced(
 
         self._fewshot_cache[variance_seed] = examples  # Store few shot examples
 
+
+    def _init_fewshot_sampling_by_id(self, variance_seed: int):
+        """Select few-shot examples by matching specific IDs from the dataset.
+
+        Uses ``task.fewshot_id_list`` and ``task.fewshot_id_column`` to pick exact
+        rows from the few-shot split. The order of examples follows the order
+        given in ``fewshot_id_list``, making selection fully deterministic and
+        reproducible.
+
+        See: https://github.com/huggingface/lighteval/issues/634
+        """
+        fewshot_id_list = [str(x) for x in self.task.fewshot_id_list]
+        fewshotpool = self.task.fewshot_docs()
+
+        # Build a lookup: id_value -> Doc
+        id_to_doc = {}
+        for doc in fewshotpool:
+            doc_id = doc.specific.get("__fewshot_id") if doc.specific else None
+            if doc_id is not None and doc_id in fewshot_id_list:
+                id_to_doc[doc_id] = doc
+
+        # Preserve the order specified in fewshot_id_list
+        selected = []
+        missing_ids = []
+        for target_id in fewshot_id_list:
+            if target_id in id_to_doc:
+                selected.append(id_to_doc[target_id])
+            else:
+                missing_ids.append(target_id)
+
+        if missing_ids:
+            logger.warning(
+                f"Task {self.task.name}: could not find few-shot examples for IDs: {missing_ids}. "
+                f"Check that `few_shots_id_column` ('{self.task.fewshot_id_column}') exists in the dataset "
+                f"and contains these values."
+            )
+
+        if not selected:
+            raise ValueError(
+                f"Task {self.task.name}: no few-shot examples matched the provided "
+                f"`few_shots_id_list` {fewshot_id_list}. Ensure `few_shots_id_column` "
+                f"('{self.task.fewshot_id_column}') is correct."
+            )
+
+        self._fewshot_cache[variance_seed] = selected
+
     def get_fewshot_seeds(self, few_shot_iterations: int = None) -> list[int]:
         """Return a list of seeds for sampling several times the few shots"""
         # todo @saylortwift: check which seed for bb

tests/tasks/test_fewshot_id_selection.py

@@ -0,0 +1,95 @@
+"""Tests for ID-based few-shot example selection (issue #634)."""
+
+import logging
+import random
+from unittest.mock import MagicMock
+
+import pytest
+
+from lighteval.tasks.lighteval_task import LightevalTask
+from lighteval.tasks.prompt_manager import FewShotSampler
+from lighteval.tasks.requests import Doc
+
+
+def _make_doc(doc_id: str, query: str, fewshot_id: str | None = None) -> Doc:
+    specific = {"__fewshot_id": fewshot_id} if fewshot_id is not None else None
+    return Doc(query=query, choices=["A", "B"], gold_index=0, id=doc_id, specific=specific)
+
+
+def _make_task_with_fewshot_docs(fewshot_docs, fewshot_id_column=None, fewshot_id_list=None):
+    task = MagicMock(spec=LightevalTask)
+    task.name = "test_task"
+    task.fewshot_docs.return_value = fewshot_docs
+    task.fewshot_selection = "balanced"
+    task.fewshot_split = "train"
+    task.fewshot_id_column = fewshot_id_column
+    task.fewshot_id_list = fewshot_id_list
+    return task
+
+
+class TestFewShotIdSelection:
+    def test_id_based_selection_returns_correct_docs(self):
+        docs = [
+            _make_doc("0", "What is 1+1?", fewshot_id="q_001"),
+            _make_doc("1", "What is 2+2?", fewshot_id="q_002"),
+            _make_doc("2", "What is 3+3?", fewshot_id="q_003"),
+            _make_doc("3", "What is 4+4?", fewshot_id="q_004"),
+        ]
+        task = _make_task_with_fewshot_docs(docs, fewshot_id_column="id", fewshot_id_list=["q_002", "q_004"])
+        sampler = FewShotSampler(task)
+        sampler._init_fewshot_pool(num_fewshot=2, variance_seed=0)
+        pool = sampler._fewshot_cache[0]
+        pool_ids = [d.specific["__fewshot_id"] for d in pool]
+        assert pool_ids == ["q_002", "q_004"]
+
+    def test_id_based_selection_preserves_order(self):
+        docs = [
+            _make_doc("0", "Q1", fewshot_id="a"),
+            _make_doc("1", "Q2", fewshot_id="b"),
+            _make_doc("2", "Q3", fewshot_id="c"),
+        ]
+        task = _make_task_with_fewshot_docs(docs, fewshot_id_column="id", fewshot_id_list=["c", "a"])
+        sampler = FewShotSampler(task)
+        sampler._init_fewshot_pool(num_fewshot=2, variance_seed=0)
+        pool = sampler._fewshot_cache[0]
+        pool_ids = [d.specific["__fewshot_id"] for d in pool]
+        assert pool_ids == ["c", "a"]
+
+    def test_id_based_selection_warns_on_missing_ids(self, caplog):
+        docs = [_make_doc("0", "Q1", fewshot_id="exists")]
+        task = _make_task_with_fewshot_docs(docs, fewshot_id_column="id", fewshot_id_list=["exists", "does_not_exist"])
+        sampler = FewShotSampler(task)
+        with caplog.at_level(logging.WARNING):
+            sampler._init_fewshot_pool(num_fewshot=2, variance_seed=0)
+        assert "does_not_exist" in caplog.text
+
+    def test_id_based_selection_raises_on_all_missing(self):
+        docs = [_make_doc("0", "Q1", fewshot_id="x")]
+        task = _make_task_with_fewshot_docs(docs, fewshot_id_column="id", fewshot_id_list=["nonexistent_1", "nonexistent_2"])
+        sampler = FewShotSampler(task)
+        with pytest.raises(ValueError, match="no few-shot examples matched"):
+            sampler._init_fewshot_pool(num_fewshot=2, variance_seed=0)
+
+    def test_no_id_list_falls_back_to_default(self):
+        docs = [_make_doc("0", "Q1", fewshot_id="a"), _make_doc("1", "Q2", fewshot_id="b")]
+        task = _make_task_with_fewshot_docs(docs, fewshot_id_column=None, fewshot_id_list=None)
+        sampler = FewShotSampler(task)
+        sampler._init_fewshot_pool(num_fewshot=2, variance_seed=0)
+        pool = sampler._fewshot_cache[0]
+        assert len(pool) == 2
+
+    def test_sample_fewshot_examples_with_id_selection(self):
+        docs = [
+            _make_doc("0", "Q1", fewshot_id="id_A"),
+            _make_doc("1", "Q2", fewshot_id="id_B"),
+            _make_doc("2", "Q3", fewshot_id="id_C"),
+            _make_doc("3", "Q4", fewshot_id="id_D"),
+        ]
+        eval_doc = _make_doc("99", "Eval question", fewshot_id="id_EVAL")
+        task = _make_task_with_fewshot_docs(docs, fewshot_id_column="id", fewshot_id_list=["id_B", "id_D"])
+        sampler = FewShotSampler(task)
+        rnd = random.Random(42)
+        result = sampler.sample_fewshot_examples(num_fewshot=2, variance_seed=0, formatted_doc=eval_doc, sampler=rnd)
+        result_ids = [d.specific["__fewshot_id"] for d in result]
+        assert result_ids == ["id_B", "id_D"]
+        assert len(result) == 2