From 479f8dcb47c9fbac22961e9c54051ee6383f0979 Mon Sep 17 00:00:00 2001 From: ErenAta16 Date: Mon, 6 Jul 2026 14:58:01 +0300 Subject: [PATCH] Add support for selecting few-shot examples by id Custom tasks currently choose few-shot examples through few_shots_select (sequential, random, or balanced), which always picks from the pool programmatically. There was no way to pin a task to a specific, reproducible set of few-shot examples the way the LM-Evaluation-Harness id_sampler does, which matters when comparing against a published harness run or when a paper specifies exact few-shot examples by id. LightevalTaskConfig gains two new fields: - few_shots_id_column: name of a dataset column to use as a stable identifier for few-shot rows. When set, this value becomes the id of the corresponding few-shot Doc instead of the row's position in the split. Evaluation docs are unaffected and keep using the row position, so this is opt-in and backward compatible. - few_shots_id_list: an 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. FewShotSampler picks up few_shots_id_list ahead of the existing sorting strategies and raises a clear error if any requested id is missing from the few-shot pool, rather than failing silently or picking something else instead. Fixes #634 --- src/lighteval/tasks/lighteval_task.py | 18 ++++++++++- src/lighteval/tasks/prompt_manager.py | 22 +++++++++++++- tests/unit/prompt/test_prompt_manager.py | 38 ++++++++++++++++++++++++ tests/unit/tasks/test_lighteval_task.py | 28 +++++++++++++++++ 4 files changed, 104 insertions(+), 2 deletions(-) diff --git a/src/lighteval/tasks/lighteval_task.py b/src/lighteval/tasks/lighteval_task.py index 5e9bac215..4280a2b33 100644 --- a/src/lighteval/tasks/lighteval_task.py +++ b/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 diff --git a/src/lighteval/tasks/prompt_manager.py b/src/lighteval/tasks/prompt_manager.py index f72b8050c..d9f0a955a 100644 --- a/src/lighteval/tasks/prompt_manager.py +++ b/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 diff --git a/tests/unit/prompt/test_prompt_manager.py b/tests/unit/prompt/test_prompt_manager.py index 6a0c6a326..458efa5e6 100644 --- a/tests/unit/prompt/test_prompt_manager.py +++ b/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) diff --git a/tests/unit/tasks/test_lighteval_task.py b/tests/unit/tasks/test_lighteval_task.py index 7cdb7b6f5..2418dcd51 100644 --- a/tests/unit/tasks/test_lighteval_task.py +++ b/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"]