From cd782a1977d3e849c9c8d83d69d2dbe521eefc58 Mon Sep 17 00:00:00 2001 From: Pyry Takala Date: Thu, 20 Nov 2025 04:18:42 +0000 Subject: [PATCH 1/6] Fix: Validate stop parameter against dataset size Add validation in Task.__len__() to ensure stop parameter does not exceed the actual dataset size. This prevents IndexError crashes during training when invalid stop values are provided. The validation is centralized in the base Task class and preserves the original lazy evaluation behavior - num_examples() is only called when needed (for validation when stop is provided, or for default value when stop is None). Fixes issue where training would crash with IndexError when iterating over Task instances with stop > dataset_size. --- tasks/common.py | 11 ++++++++++- 1 file changed, 10 insertions(+), 1 deletion(-) diff --git a/tasks/common.py b/tasks/common.py index dcd2e91..afa47cc 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -34,7 +34,16 @@ class Task: def __len__(self): start = self.start - stop = self.num_examples() if self.stop is None else self.stop + if self.stop is not None: + num_ex = self.num_examples() + if self.stop > num_ex: + raise ValueError( + f"Stop parameter ({self.stop}) exceeds dataset size ({num_ex}). " + f"Please use stop <= {num_ex} or remove the stop parameter to use the full dataset." + ) + stop = self.stop + else: + stop = self.num_examples() step = self.step span = stop - start num = (span + step - 1) // step # ceil_div(span, step) From 85e49943ed9eb072b8a55b32375be35b2d97cef6 Mon Sep 17 00:00:00 2001 From: Pyry Takala Date: Fri, 21 Nov 2025 20:04:33 +0000 Subject: [PATCH 2/6] Gracefully handle stop > dataset_size with warning --- tasks/common.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/tasks/common.py b/tasks/common.py index afa47cc..a63cb7a 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -36,12 +36,14 @@ class Task: start = self.start if self.stop is not None: num_ex = self.num_examples() + stop = min(self.stop, num_ex) # Gracefully cap at dataset size if self.stop > num_ex: - raise ValueError( + import warnings + warnings.warn( f"Stop parameter ({self.stop}) exceeds dataset size ({num_ex}). " - f"Please use stop <= {num_ex} or remove the stop parameter to use the full dataset." + f"Using {num_ex} examples instead.", + UserWarning ) - stop = self.stop else: stop = self.num_examples() step = self.step From a33d04dca175304b9fb1614b7b3d028eca0f4780 Mon Sep 17 00:00:00 2001 From: Pyry Takala Date: Fri, 21 Nov 2025 20:51:46 +0000 Subject: [PATCH 3/6] Cap stop parameter and warn once when it exceeds dataset size --- tasks/common.py | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) diff --git a/tasks/common.py b/tasks/common.py index a63cb7a..540ff6b 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -6,6 +6,9 @@ Example tasks: MMLU, ARC-Easy, ARC-Challenge, GSM8K, HumanEval, SmolTalk. """ import random +import logging + +logger = logging.getLogger(__name__) class Task: """ @@ -36,14 +39,14 @@ class Task: start = self.start if self.stop is not None: num_ex = self.num_examples() - stop = min(self.stop, num_ex) # Gracefully cap at dataset size if self.stop > num_ex: - import warnings - warnings.warn( + # Warn once, then cap stop + logger.warning( f"Stop parameter ({self.stop}) exceeds dataset size ({num_ex}). " - f"Using {num_ex} examples instead.", - UserWarning + f"Using {num_ex} examples instead." ) + self.stop = num_ex + stop = self.stop else: stop = self.num_examples() step = self.step From df9a644e249227a4f49376cd8672cc1da36a206e Mon Sep 17 00:00:00 2001 From: Sofie Van Landeghem Date: Sat, 22 Nov 2025 22:48:55 +0100 Subject: [PATCH 4/6] make code bit more succinct --- tasks/common.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/tasks/common.py b/tasks/common.py index 540ff6b..e202508 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -46,9 +46,7 @@ class Task: f"Using {num_ex} examples instead." ) self.stop = num_ex - stop = self.stop - else: - stop = self.num_examples() + stop = self.num_examples() if self.stop is None else self.stop step = self.step span = stop - start num = (span + step - 1) // step # ceil_div(span, step) From c09b8976013f0a346c92f5d4185e9b59b5e4b4a9 Mon Sep 17 00:00:00 2001 From: Sofie Van Landeghem Date: Sat, 22 Nov 2025 22:50:34 +0100 Subject: [PATCH 5/6] further cleanup --- tasks/common.py | 7 +++---- 1 file changed, 3 insertions(+), 4 deletions(-) diff --git a/tasks/common.py b/tasks/common.py index e202508..1fd2c23 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -37,16 +37,15 @@ class Task: def __len__(self): start = self.start - if self.stop is not None: - num_ex = self.num_examples() - if self.stop > num_ex: + num_ex = self.num_examples() + if self.stop is not None and self.stop > num_ex: # Warn once, then cap stop logger.warning( f"Stop parameter ({self.stop}) exceeds dataset size ({num_ex}). " f"Using {num_ex} examples instead." ) self.stop = num_ex - stop = self.num_examples() if self.stop is None else self.stop + stop = num_ex if self.stop is None else self.stop step = self.step span = stop - start num = (span + step - 1) // step # ceil_div(span, step) From 26b0941f75608df8dfa5d6c6dc0bcf1020f8989d Mon Sep 17 00:00:00 2001 From: Sofie Van Landeghem Date: Sat, 22 Nov 2025 22:51:42 +0100 Subject: [PATCH 6/6] fix --- tasks/common.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tasks/common.py b/tasks/common.py index 1fd2c23..f2228eb 100644 --- a/tasks/common.py +++ b/tasks/common.py @@ -39,12 +39,12 @@ class Task: start = self.start num_ex = self.num_examples() if self.stop is not None and self.stop > num_ex: - # Warn once, then cap stop - logger.warning( - f"Stop parameter ({self.stop}) exceeds dataset size ({num_ex}). " - f"Using {num_ex} examples instead." - ) - self.stop = num_ex + # Warn once, then cap stop + logger.warning( + f"Stop parameter ({self.stop}) exceeds dataset size ({num_ex}). " + f"Using {num_ex} examples instead." + ) + self.stop = num_ex stop = num_ex if self.stop is None else self.stop step = self.step span = stop - start