update evaluator and trainerg

Author: Nghi Bui

codetf/common/__init__.py

@@ -23,16 +23,18 @@ def __init__(self, model_class, num_workers=5):
 
 
     def evaluate_pass_k(self, problems, unit_tests, batch_size=1, max_length=600, 
-                        top_p=0.95, k=[1,10,100], 
+                        top_p=0.95, k=[1,10,100], temperature=1.2,
                         num_return_sequences=200, sequences_per_chunk=10, num_workers=1):
         # Load dataset
-        data_loader = Dat aLoader(problems, batch_size=batch_size)
+        # Please keep batch_size = 1 to avoid any unexpected error
+        data_loader = DataLoader(problems, batch_size=batch_size)
         data_loader = self.accelerator.prepare(data_loader)
-        
+        model_name = type(self.model_class).__name__
         # Initialize stopping criteria
         gen_kwargs = {
             "do_sample": True,
             "top_p": top_p,
+            "temperature": temperature,
             "stopping_criteria": StoppingCriteriaList([EndOfFunctionCriteria(0, EOF_STRINGS, self.model_class.get_tokenizer())]),
         }
 
@@ -54,7 +56,6 @@ def evaluate_pass_k(self, problems, unit_tests, batch_size=1, max_length=600,
                     input_ids = prompt_ids[0, :attention_masks[0].sum().item()]
 
                     input_data = self.model_class.get_tokenizer().decode(input_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
-
                     batch_generated_ids = self.model_class.get_model().generate(
                         input_ids=input_ids.unsqueeze(0),
                         attention_mask=attention_masks[0, :attention_masks[0].sum().item()].unsqueeze(0), 
@@ -66,14 +67,16 @@ def evaluate_pass_k(self, problems, unit_tests, batch_size=1, max_length=600,
                     gen_codes = self.model_class.get_tokenizer().batch_decode(batch_generated_ids, 
                                             skip_special_tokens=True, clean_up_tokenization_spaces=True)
 
-                    for item in gen_codes:
-                        cleaned =  remove_last_block(item)
-                        solutions_per_chunk.append(cleaned)
+                    for i,item in enumerate(gen_codes):
+                        result =  remove_last_block(item)
+                        if model_name == "Seq2SeqModel":
+                            result = f"{input_data} {result}"
+                        
+                        solutions_per_chunk.append(result)
 
             solutions.append(solutions_per_chunk)
             dataloader_pbar.set_description(f"Processing step {step+1}/{len(data_loader)}")
 
-
         pass_at_k, _ = self.code_eval.compute(
             references=unit_tests, predictions=solutions, k=k, num_workers=num_workers
         )

codetf/performance/model_evaluator.py

@@ -52,6 +52,7 @@ def init_trainer(self):
 
     def train(self):
         self.trainer.train()
+        # self.trainer.save_model(self.checkpoints_path)
 
     def evaluate(self, dataset=None):
         self.trainer.evaluate(dataset)
@@ -70,8 +71,10 @@ def get_default_codet5_hyperparameters(self):
             sharded_ddp=hyperparameters_config["sharded_ddp"],
             logging_steps=hyperparameters_config["logging_steps"],
             evaluation_strategy=hyperparameters_config["evaluation_strategy"],
+            save_strategy=hyperparameters_config["save_strategy"],
             gradient_checkpointing=hyperparameters_config["gradient_checkpointing"],
             auto_find_batch_size=hyperparameters_config["auto_find_batch_size"],
+            load_best_model_at_end=hyperparameters_config["load_best_model_at_end"],
             output_dir=self.checkpoints_path
         )
         # return hyperparameters_config
@@ -91,8 +94,10 @@ def get_default_causal_lm_hyperparameters(self):
             sharded_ddp=hyperparameters_config["sharded_ddp"],
             logging_steps=hyperparameters_config["logging_steps"],
             evaluation_strategy=hyperparameters_config["evaluation_strategy"],
+            save_strategy=hyperparameters_config["save_strategy"],
             gradient_checkpointing=hyperparameters_config["gradient_checkpointing"],
             auto_find_batch_size=hyperparameters_config["auto_find_batch_size"],
+            load_best_model_at_end=hyperparameters_config["load_best_model_at_end"],
             output_dir=self.checkpoints_path
         )
         # return hyperparameters_config

codetf/trainer/base_trainer.py

@@ -21,6 +21,7 @@ def __init__(self, train_dataset, validation_dataset=None, tokenizer=None,
         self.trainer = self.init_trainer()
 
         if peft:
+            self.peft = peft
             self.model = prepare_model_for_int8_training(self.model)
             if peft == "lora":
                 peft_config = self.get_default_lora_config_for_codet5()

codetf/trainer/causal_lm_trainer.py

@@ -38,6 +38,7 @@ def __init__(self, train_dataset, validation_dataset=None, tokenizer=None,
         self.trainer = self.init_trainer()
 
         if peft:
+            self.peft = peft
             self.model = prepare_model_for_int8_training(self.model)
             if peft == "lora":
                 peft_config = self.get_default_lora_config_for_codet5()