feat: add ID-based few-shot example selection (#634) - #1263
Conversation
Add few_shots_id_column and few_shots_id_list to LightevalTaskConfig, allowing users to specify exact dataset rows as few-shot examples by their unique identifiers. This mirrors the id_sampler functionality in lm-evaluation-harness. When configured, the FewShotSampler filters the few-shot pool to only include docs matching the given IDs, preserving the order specified in few_shots_id_list. Existing behavior is unaffected (both fields default to None).
There was a problem hiding this comment.
Thanks for taking this on. I traced this exact head through LightevalTask._get_docs_from_split() → fewshot_docs() → FewShotSampler. The helper preserves order only for tests that pre-populate Doc.specific; the real dataset path never copies few_shots_id_column into any Doc, and the selected pool can still be sampled a second time. A configured task therefore either reports no matches or loses the fixed-ID ordering, depending on the selection mode.
There is also a same-split correctness boundary to make explicit. The example in #634 uses train for both evaluation and few-shot selection with exclude_from_task: true. This patch has no source-ID exclusion path, while the existing full-Doc equality check is not stable because formatters may produce different documents for __few_shots=True; that can leave the evaluated row in its own prompt. Please add an end-to-end in-memory DatasetDict test that exercises real LightevalTask formatting, plus same-split exclusion and missing-ID failure coverage. #1334 also changes this sampler's pool/RNG behavior, so coordinating with or rebasing after it would avoid reintroducing those state-isolation issues.
|
|
||
| # 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 |
There was a problem hiding this comment.
few_shots_id_column is stored but never consumed by _get_docs_from_split(): that method sets only __few_shots, __index, and Doc.id, so the production path never puts __fewshot_id in Doc.specific. The new sampler therefore reports no matches for a real dataset even though the MagicMock tests pass. Please carry the raw column value through the actual task path using a collision-safe source-ID contract, and cover it with an in-memory DatasetDict test.
| 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: |
There was a problem hiding this comment.
This changes only pool initialization; sample_fewshot_examples() still calls _sample_from_pool() afterward. With random_sampling_from_train, that can reorder the configured IDs, and with random_sampling it requests num_fewshot + 1, which raises when the fixed list contains exactly K items. Fixed-ID mode should bypass secondary sampling entirely so the configured order is the final order.
| # 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 |
There was a problem hiding this comment.
fewshot_id_list is normalized to strings, but doc_id remains in its source type. An integer-valued dataset ID therefore cannot match its configured string form. Please normalize the source ID at the dataset boundary too, while preserving one canonical representation for logs and cache hashing.
| else: | ||
| missing_ids.append(target_id) | ||
|
|
||
| if missing_ids: |
There was a problem hiding this comment.
Continuing after one missing ID silently returns fewer than num_fewshot examples, so a supposedly reproducible task changes behavior when a dataset row disappears. Please fail fast for any missing or duplicate configured ID (and for an absent ID column), rather than warning and accepting a partial selection.
2 participants
Add few_shots_id_column and few_shots_id_list to LightevalTaskConfig, allowing users to specify exact dataset rows as few-shot examples by their unique identifiers. This mirrors the id_sampler functionality in lm-evaluation-harness.
When configured, the FewShotSampler filters the few-shot pool to only include docs matching the given IDs, preserving the order specified in few_shots_id_list. Existing behavior is unaffected (both fields default to None).