QOL changes for generations

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:
@@ -75,6 +76,7 @@ def generate_text(self, task_name, intermediate_generations=None):
             args=self.args,
             curr_sample_idx=curr_sample_idx,  # curr_sample_idx will added to limit_start to fix indexing
             save_every_k_tasks=self.args.save_every_k_tasks,
+            intermediate_generations=generations,
             intermediate_save_generations_path=intermediate_save_generations_path,
         )
         generations.extend(new_generations)

bigcode_eval/generation.py

@@ -1,7 +1,7 @@
 import json
 from math import ceil
 
-from typing import Optional
+from typing import List, Optional
 
 from accelerate.utils import set_seed
 from torch.utils.data.dataloader import DataLoader
@@ -49,6 +49,7 @@ def parallel_generations(
         args,
         curr_sample_idx: int = 0,
         save_every_k_tasks: int = -1,
+        intermediate_generations: Optional[List[Optional[List[Optional[str]]]]] = None,
         intermediate_save_generations_path: Optional[str] = None,
 ):
     if args.load_generations_path:
@@ -151,6 +152,7 @@ def parallel_generations(
         postprocess=args.postprocess,
         is_wrapped=is_loaded_in_8bit or is_loaded_in_4bit,
         save_every_k_tasks=save_every_k_tasks,
+        intermediate_generations=intermediate_generations,
         intermediate_save_generations_path=intermediate_save_generations_path,
         **gen_kwargs,
     )

bigcode_eval/utils.py

@@ -235,6 +235,7 @@ def complete_code(
     postprocess=True,
     is_wrapped=False,
     save_every_k_tasks: int = -1,
+    intermediate_generations: Optional[List[Optional[List[Optional[str]]]]] = None,
     intermediate_save_generations_path: Optional[str] = None,
     **gen_kwargs,
 ):
@@ -246,6 +247,7 @@ def complete_code(
     # keep track of the list of generated codes
     # where len(code_gens) = n_tasks and len(code_gens[0]) = number of generated code samples
     code_gens: List[List[Optional[str]]] = [[] for _ in range(n_tasks)]
+    intermediate_generations = [] if not intermediate_generations else intermediate_generations
     gen_token_dict = defaultdict(list)  # dict of list of generated tokens
     for step, batch in tqdm(
         enumerate(dataloader),
@@ -332,7 +334,8 @@ def complete_code(
                     gen_token_dict,
                 )
                 with open(intermediate_save_generations_path, "w") as fp:
-                    json.dump(code_gens, fp)
+                    intermediate_generations.extend(code_gens)
+                    json.dump(intermediate_generations, fp)
                     print(
                         f"intermediate generations were saved at {intermediate_save_generations_path}"
                     )