Update split iteration for DynamicBatchingDataset (#684)

This PR aims to make iterating over splits a bit more intuitive, at least in my opinion. Open to feedback though! If the current behavior was intentional, feel free to close.

---------

Co-authored-by: Nathan Habib <[email protected]>

Author: qubvel

src/lighteval/data.py

@@ -22,11 +22,11 @@
 
 import logging
 import math
-from typing import Iterator, Tuple
+from typing import Iterator
 
 import torch
 from packaging import version
-from torch.utils.data import Dataset
+from torch.utils.data import Dataset, Subset
 
 
 if version.parse(torch.__version__) >= version.parse("2.5.0"):
@@ -82,8 +82,6 @@ def __init__(
 
         self.num_dataset_splits, self.splits = self.init_split_limits(num_dataset_splits)
 
-        self.split_start, self.split_end = self.splits[0]
-
     def init_split_limits(self, num_dataset_splits):
         if num_dataset_splits >= self.total_size:
             logger.warning(
@@ -121,48 +119,31 @@ def get_original_order(self, new_arr: list) -> list:
 
         return original_order
 
-    def get_split_start_end(self, split_id: int) -> Tuple[int, int]:
-        """
-        Get the start and end indices of a dataset split.
-
-        Args:
-            split_id (int): The ID of the split.
-
-        Returns:
-            tuple: A tuple containing the start and end indices of the split.
-        """
-        self.split_start, self.split_end = self.splits[split_id]
-        return self.split_start, self.split_end
-
-    def splits_start_end_iterator(self) -> Iterator[Tuple[int, int]]:
+    def splits_iterator(self) -> Iterator[Subset]:
         """
-        Iterator that yields the start and end indices of each dataset split.
-        Also updates the starting batch size for each split (trying to double
-        the batch every time we move to a new split).
+        Iterator that yields the dataset splits based on the split limits.
 
         Yields:
-            tuple: A tuple containing the start and end indices of a split.
+            Subset: A subset of the dataset.
         """
         split_range = self.num_dataset_splits
         if self.total_size == 0:
             split_range = 0
-        for split_id in range(split_range):
-            yield self.get_split_start_end(split_id)
+        for i in range(split_range):
+            split_start, split_end = self.splits[i]
+            yield Subset(self, range(split_start, split_end))
 
     def __getitem__(self, index) -> Request:
         """
-        Get an item from the dataset depending on the split we are currently in.
-        For instance, if we are in split 0, we will get the item at index 0, if
-        we are in split 1, we will get the item at index self.split_size, etc.
-        Used for dynamic batching.
+        Get an item from the dataset.
 
         Args:
             index (int): The index of the item.
 
         Returns:
             Any: The item at the specified index.
         """
-        return self.sorted_data[index + self.split_start]
+        return self.sorted_data[index]
 
     def __len__(self) -> int:
         """
@@ -173,7 +154,7 @@ def __len__(self) -> int:
         Returns:
             int: The length of the dataset.
         """
-        return self.split_end - self.split_start
+        return len(self.sorted_data)
 
     def __iter__(self) -> Iterator[Request]:
         """
@@ -186,7 +167,7 @@ def __iter__(self) -> Iterator[Request]:
         Yields:
             Any: The items of the dataset.
         """
-        for i in range(self.split_start, self.split_end):
+        for i in range(len(self)):
             yield self.sorted_data[i]
 
     def _sorting_criteria(self, request) -> int:

src/lighteval/models/endpoints/endpoint_model.py

@@ -463,14 +463,14 @@ def greedy_until(
         batch_size = override_bs if override_bs is not None else BATCH_SIZE
         results: List[str] = []
 
-        for _, _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,
         ):
-            dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=lambda batch: batch)
+            dataloader = DataLoader(split, batch_size=batch_size, collate_fn=lambda batch: batch)
 
             for batch in tqdm(
                 dataloader, desc="Greedy generation", position=1, leave=False, disable=self.disable_tqdm
@@ -512,14 +512,14 @@ def loglikelihood(
         batch_size = override_bs if override_bs is not None else BATCH_SIZE
         results: List[str] = []
 
-        for _, _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,
         ):
-            dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=lambda batch: batch)
+            dataloader = DataLoader(split, batch_size=batch_size, collate_fn=lambda batch: batch)
 
             for batch in tqdm(dataloader, desc="Loglikelihoods", position=1, leave=False, disable=self.disable_tqdm):
                 if self.use_async:
@@ -563,14 +563,14 @@ def loglikelihood_rolling(
         batch_size = override_bs if override_bs is not None else BATCH_SIZE
         results: List[str] = []
 
-        for _, _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,
         ):
-            dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=lambda batch: batch)
+            dataloader = DataLoader(split, batch_size=batch_size, collate_fn=lambda batch: batch)
 
             for batch in tqdm(
                 dataloader, desc="Loglikelihoods, rolling", position=1, leave=False, disable=self.disable_tqdm

src/lighteval/models/endpoints/inference_providers_model.py

@@ -210,15 +210,15 @@ def greedy_until(
         dataset = GenerativeTaskDataset(requests=requests, num_dataset_splits=self.DATASET_SPLITS)
         results = []
 
-        for _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=False,  # self.disable_tqdm,
         ):
-            contexts = [c.context for c in dataset]
-            num_samples = dataset[0].num_samples
+            contexts = [sample.context for sample in split]
+            num_samples = split[0].num_samples
 
             responses = asyncio.run(self.__call_api_parallel(contexts, num_samples))
 

src/lighteval/models/endpoints/openai_model.py

@@ -184,17 +184,17 @@ def greedy_until(
         dataset = GenerativeTaskDataset(requests=requests, num_dataset_splits=self.DATASET_SPLITS)
         results = []
 
-        for _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=False,  # self.disable_tqdm,
         ):
-            max_new_tokens = dataset[0].generation_size  # could be none
-            return_logits = dataset[0].use_logits
-            num_samples = dataset[0].num_samples
-            contexts = [c.context for c in dataset]
+            max_new_tokens = split[0].generation_size  # could be none
+            return_logits = split[0].use_logits
+            num_samples = split[0].num_samples
+            contexts = [sample.context for sample in split]
 
             responses = self.__call_api_parallel(contexts, return_logits, max_new_tokens, num_samples)
 
@@ -251,24 +251,22 @@ def _loglikelihood_tokens(
         dataset = LoglikelihoodDataset(requests=requests, num_dataset_splits=1)
         results = []
 
-        for _ in tqdm(dataset.splits_start_end_iterator()):
-            inputs = [dataset[i].context for i in range(len(dataset))]
-            logit_biass = []
-            max_new_tokens = [len(dataset[i].tokenized_continuation) for i in range(len(dataset))]
+        for split in tqdm(dataset.splits_iterator()):
+            inputs = [sample.context for sample in split]
+            max_new_tokens = [len(sample.tokenized_continuation) for sample in split]
 
             assert all(
                 new_tokens == 1 for new_tokens in max_new_tokens
             ), "Only single token continuations are supported when using openai API."
 
-            for i in range(len(dataset)):
-                logit_bias = {tok: 100 for tok in dataset[i].tokenized_continuation}
-                logit_biass.append(logit_bias)
+            logit_biases = [{tok: 100 for tok in sample.tokenized_continuation} for sample in split]
 
             outputs = self.__call_api_parallel(
-                inputs, return_logits=True, max_new_tokens=max_new_tokens, num_samples=1, logit_bias=logit_biass
+                inputs, return_logits=True, max_new_tokens=max_new_tokens, num_samples=1, logit_bias=logit_biases
             )
 
-            for output, input in zip(outputs, dataset):
+            for i, output in enumerate(outputs):
+                input = split[i]
                 continuation_logprobs = [content.logprob for content in output.choices[0].logprobs.content]
                 answer = LoglikelihoodResponse(
                     input_tokens=input.tokenized_context + input.tokenized_continuation,

src/lighteval/models/litellm_model.py

@@ -227,17 +227,17 @@ def greedy_until(
         dataset = GenerativeTaskDataset(requests=requests, num_dataset_splits=self.DATASET_SPLITS)
         results = []
 
-        for _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
-            disable=False,  # self.disable_tqdm,
+            disable=self.disable_tqdm,
         ):
-            contexts = [c.context for c in dataset]
-            max_new_tokens = dataset[0].generation_size  # could be none
-            return_logits = dataset[0].use_logits
-            num_samples = dataset[0].num_samples
+            contexts = [sample.context for sample in split]
+            max_new_tokens = split[0].generation_size  # could be none
+            return_logits = split[0].use_logits
+            num_samples = split[0].num_samples
             stop_sequence = requests[0].stop_sequence
 
             responses = self.__call_api_parallel(contexts, return_logits, max_new_tokens, num_samples, stop_sequence)

src/lighteval/models/sglang/sglang_model.py

@@ -177,8 +177,8 @@ def greedy_until(
         dataset = GenerativeTaskDataset(requests=requests, num_dataset_splits=self.DATASET_SPLITS)
         results = []
 
-        for _ in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
@@ -187,12 +187,12 @@ def greedy_until(
             if self.use_chat_template:
                 stop_tokens = []
             else:
-                stop_tokens = dataset[0].stop_sequence
+                stop_tokens = split[0].stop_sequence
 
-            max_new_tokens = dataset[0].generation_size  # could be none
-            num_samples = dataset[0].num_samples
+            max_new_tokens = split[0].generation_size  # could be none
+            num_samples = split[0].num_samples
 
-            context = [c.context for c in dataset]
+            context = [sample.context for sample in split]
             tokenized = self.tokenizer(context, add_special_tokens=self.add_special_tokens)
 
             # The main question for this step is the following:
@@ -298,14 +298,15 @@ def _loglikelihood_tokens(
         dataset = LoglikelihoodDataset(requests=requests, num_dataset_splits=1)
         res = []
 
-        for _ in tqdm(dataset.splits_start_end_iterator(), disable=False):
+        for split in tqdm(dataset.splits_iterator(), disable=False):
             # the last token is an eos token, so we don't need to add it
-            inputs = [dataset[i].tokenized_context + dataset[i].tokenized_continuation for i in range(len(dataset))]
+            inputs = [sample.tokenized_context + sample.tokenized_continuation for sample in split]
             # Left truncate the inputs to the maximum length
             inputs = [input[-self.max_length :] for input in inputs]
             outputs = self._generate(inputs, generate=False)
 
-            for output, input in zip(outputs, dataset):
+            for i, output in enumerate(outputs):
+                input = split[i]
                 continuation_logprobs = []
                 meta_info = output["meta_info"]
                 input_token_logprobs = meta_info["input_token_logprobs"][::-1]

src/lighteval/models/transformers/transformers_model.py

@@ -530,21 +530,19 @@ def greedy_until(
         starting_batch_size = STARTING_BATCH_SIZE
         results = []
 
-        for split_start, split_end in tqdm(
-            dataset.splits_start_end_iterator(),
+        for split in tqdm(
+            dataset.splits_iterator(),
             total=dataset.num_dataset_splits,
             desc="Splits",
             position=0,
             disable=self.disable_tqdm,
         ):
-            if dataset[0].generation_size is None:
+            if split[0].generation_size is None:
                 # No constraints on the generation size: max length allowed is the max model context
                 max_context_continuation_size_allowed = self.max_length
             else:
                 # Longest context in the current split is the first item (since we sort reversed)
-                longest_context_continuation_size_in_split = (
-                    len(dataset[0].tokenized_context) + dataset[0].generation_size
-                )
+                longest_context_continuation_size_in_split = len(split[0].tokenized_context) + split[0].generation_size
                 max_context_continuation_size_allowed = min(
                     longest_context_continuation_size_in_split, self.max_length
                 )
@@ -556,7 +554,7 @@ def greedy_until(
             # For next iteration, since the batch will be smaller, we'll test a bigger batch size
             starting_batch_size = batch_size * 2
 
-            dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=lambda batch: batch)
+            dataloader = DataLoader(split, batch_size=batch_size, collate_fn=lambda batch: batch)
             if self.accelerator:
                 dataloader = self.accelerator.prepare(dataloader)
 
@@ -765,9 +763,9 @@ def _loglikelihood_tokens(
         starting_batch_size = STARTING_BATCH_SIZE
         res = []
 
-        for split_start, split_end in tqdm(dataset.splits_start_end_iterator()):
-            context_enc = dataset[0].tokenized_context
-            continuation_enc = dataset[0].tokenized_continuation
+        for split in tqdm(dataset.splits_iterator()):
+            context_enc = split[0].tokenized_context
+            continuation_enc = split[0].tokenized_continuation
             if rolling:  # we take all the sequence in rolling mode
                 max_context_continuation_size_allowed = len(context_enc + continuation_enc)
             else:  # in normal mode, we left cut the context if needed
@@ -782,7 +780,7 @@ def _loglikelihood_tokens(
             )
             starting_batch_size = batch_size * 2
 
-            dataloader = DataLoader(dataset, batch_size=batch_size, collate_fn=lambda batch: batch)
+            dataloader = DataLoader(split, batch_size=batch_size, collate_fn=lambda batch: batch)
             if self.accelerator:
                 dataloader = self.accelerator.prepare(dataloader)
 
@@ -1009,13 +1007,13 @@ def _loglikelihood_single_token(
         starting_batch_size = STARTING_BATCH_SIZE
         res = []
 
-        for split_start, split_end in tqdm(dataset.splits_start_end_iterator()):
-            context_enc = dataset[0].tokenized_context
+        for split in tqdm(dataset.splits_iterator()):
+            context_enc = split[0].tokenized_context
             max_context = len(context_enc[-self.max_length :])
             batch_size = self._get_batch_size(override_bs=self.config.batch_size, max_input_length=max_context)
             starting_batch_size = batch_size * 2
 
-            dataloader = DataLoader(dataset, batch_size=starting_batch_size, collate_fn=lambda batch: batch)
+            dataloader = DataLoader(split, batch_size=starting_batch_size, collate_fn=lambda batch: batch)
             if self.accelerator is not None:
                 dataloader = self.accelerator.prepare(dataloader)