minor optimization for preventing oob errors

bigcode-project · Jan 3, 2024 · 0754793 · 0754793
1 parent df8ffc9
commit 0754793
Showing 1 changed file with 2 additions and 1 deletion.
diff --git a/bigcode_eval/evaluator.py b/bigcode_eval/evaluator.py
@@ -44,7 +44,8 @@ def generate_text(self, task_name, intermediate_generations=None):
         task = tasks.get_task(task_name, self.args)
         dataset = task.get_dataset()
         # if args.limit is None, use all samples
-        n_tasks = self.args.limit if self.args.limit else len(dataset)
+        # if args.limit is used, make sure args.limit_start + args.limit <= len(dataset)
+        n_tasks = min(self.args.limit, len(dataset) - self.args.limit_start) if self.args.limit else len(dataset)
         # when args.limit is None
         # adjust n_tasks by args.limit_start to prevent out of bounds issues 
         if not self.args.limit: