Balanced few-shot sampling truncates on a falsy class label and uses the global random module

[ErenAta16]

Bug

FewShotSampler._init_fewshot_sampling_balanced (src/lighteval/tasks/prompt_manager.py) has two issues in the same function.

1. A falsy class label silently truncates the balanced sample

The selection loop guards the label cycle with if not next_label:

labels_iterable = cycle(sorted_labels)
while num_instances_to_sample > 0:
    next_label = next(labels_iterable, None)
    if not next_label:
        break
    ...

labels_iterable is a cycle, which never yields None for a non-empty pool, so the intent of the guard is "stop when there are no labels". But if not next_label also fires for a label that is present but falsy: the integer 0, an empty string, False. Labels come from instance.fewshot_sorting_class or as_list(instance.get_golds())[0], so any task whose class label is a falsy value hits this, e.g. a classification task with an integer 0 label, or a gold that is the empty string. When the cycle reaches that label, the loop breaks and the balanced selection is cut short, down to zero examples if the falsy label sorts first.

2. Balanced sampling uses the global random module

Unlike the sibling _init_fewshot_sampling_random (which uses a local random.Random(variance_seed)), the balanced path calls random.seed(variance_seed), random.shuffle(...), and random.randrange(...) on the global module. That makes the selection non-reproducible if any other code touches the global RNG between calls, and it leaks state into the rest of the program.

Reproduction

Using the real FewShotSampler (only a duck-typed task is stubbed, exposing fewshot_docs() / fewshot_selection / fewshot_split):

from lighteval.tasks.prompt_manager import FewShotSampler
from lighteval.tasks.requests import Doc
import random

class FakeTask:
    fewshot_selection = "balanced"
    fewshot_split = "test"
    def __init__(self, docs): self._docs = docs
    def fewshot_docs(self): return self._docs

# one class has the empty-string gold -> falsy label
docs  = [Doc(query=f"q{i}", choices=["", "x"], gold_index=0) for i in range(20)]
docs += [Doc(query=f"p{i}", choices=["", "x"], gold_index=1) for i in range(20)]
selected = FewShotSampler(FakeTask(docs)).sample_fewshot_examples(10, variance_seed=0)
print(len(selected))   # 0  (expected 10)

# global RNG pollution
random.seed(12345); before = [random.random() for _ in range(3)]
random.seed(12345)
FewShotSampler(FakeTask([Doc(query=f"q{i}", choices=["a","b"], gold_index=i%2) for i in range(20)])).sample_fewshot_examples(5, variance_seed=999)
after = [random.random() for _ in range(3)]
print(before == after)  # False  (the global stream was disturbed)

Output on current main: 0 balanced examples for the falsy-label case, and False for the RNG check.

Suggested fix

1. Change if not next_label to if next_label is None (a cycle only yields None when sorted_labels is empty, which is the case the guard was meant for; falsy-but-valid labels no longer truncate).
2. Use a local rnd = random.Random(variance_seed) and its .shuffle / .randrange, mirroring _init_fewshot_sampling_random.

Happy to open a PR with both plus regression tests (a falsy-label task gets the full num_fewshot, and a balanced call leaves the global RNG stream untouched).

Environment

Reproduced against huggingface/lighteval main (current). Pure Python, no model needed.

[atirna]

I opened #1334 for both sampler issues. Balanced sampling now keeps falsy labels and uses a local seeded RNG. uv run --no-sync pytest tests/unit/prompt/test_prompt_manager.py -q passes (6). Thanks for the detailed repro!