fix

src/lighteval/data.py

@@ -4,7 +4,7 @@
 from torch.utils.data import Dataset
 from torch.utils.data.distributed import DistributedSampler, T_co
 
-from lighteval.logging.hierarchical_logger import hlog_warn
+from lighteval.logging.hierarchical_logger import hlog, hlog_warn
 from lighteval.tasks.requests import Request
 
 
@@ -194,6 +194,22 @@ def _sorting_criteria(self, x) -> int:
         _, gen_length = meta_data[0], meta_data[1]
         return -(len(toks) + gen_length)
 
+class GenerativeTaskDatasetBrrr(DynamicBatchDataset):
+    def _sorting_criteria(self, x) -> int:
+        """
+        Collate function for generating batches.
+
+        Args:
+            x (Any): The input data.
+
+        Returns:
+            Any: The collated data.
+        """
+        toks = x[1][0]
+        meta_data = x[1][1]
+        _, gen_length = meta_data[0], meta_data[1]
+        return -(len(toks) + gen_length)
+
 
 class GenDistributedSampler(DistributedSampler):
     """A distributed sampler that copy the last element only when drop_last is False so we keep a small padding in the batches