diff --git a/src/lighteval/tasks/lighteval_task.py b/src/lighteval/tasks/lighteval_task.py index 5e9bac215..65e32651b 100644 --- a/src/lighteval/tasks/lighteval_task.py +++ b/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 diff --git a/src/lighteval/tasks/prompt_manager.py b/src/lighteval/tasks/prompt_manager.py index f72b8050c..fcf6de210 100644 --- a/src/lighteval/tasks/prompt_manager.py +++ b/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 diff --git a/tests/tasks/test_fewshot_id_selection.py b/tests/tasks/test_fewshot_id_selection.py new file mode 100644 index 000000000..b465805f6 --- /dev/null +++ b/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