update

Author: bdqnghi

README.md

@@ -87,7 +87,7 @@ conda activate codetf
 
 2. Install from [PyPI](https://pypi.org/project/salesforce-codetf/):
 ```bash
-pip install salesforce-codetf==1.0.2.1
+pip install salesforce-codetf
 ```
 
 3. Alternatively, build CodeTF from source:

codetf/common/__init__.py

@@ -7,8 +7,8 @@ causallm-codegen-350M-multi-pretrained:
     tokenizer_url: "Salesforce/codegen-350M-multi"
     max_prediction_length: 512
 causallm-codegen-350M-nl-pretrained:
-    huggingface_url: "Salesforce/codegen-350-nl"
-    tokenizer_url: "Salesforce/codegen-350-nl"
+    huggingface_url: "Salesforce/codegen-350M-nl"
+    tokenizer_url: "Salesforce/codegen-350M-nl"
     max_prediction_length: 512
 causallm-codegen-2B-mono-pretrained:
     huggingface_url: "Salesforce/codegen-2B-mono"

codetf/configs/inference/causal_lm.yaml

@@ -140,20 +140,20 @@ codet5-plus-770M-python-pretrained:
     max_prediction_length: 512
     beam_size: 5
     trust_remote_code: False
-    device_map: True
+    device_map: False
 codet5-plus-770M-pretrained:
     huggingface_url: "Salesforce/codet5p-770m"
     tokenizer_url: "Salesforce/codet5p-770m"
     max_source_length: 512
     max_prediction_length: 512
     beam_size: 5
     trust_remote_code: False
-    device_map: True
+    device_map: False
 codet5-plus-220M-pretrained:
     huggingface_url: "Salesforce/codet5p-220m"
     tokenizer_url: "Salesforce/codet5p-220m"
     max_source_length: 512
     max_prediction_length: 512
     beam_size: 5
     trust_remote_code: False
-    device_map: True
+    device_map: False

codetf/configs/inference/codet5.yaml

@@ -6,9 +6,6 @@ hyperparameters:
   num_train_epochs: 10
   auto_find_batch_size: True
   batch_size: 4
-  max_steps: 1000
-  eval_steps: 100
-  save_steps: 1000
   logging_steps: 100
   per_device_train_batch_size: 8
   per_device_eval_batch_size: 8
@@ -30,3 +27,4 @@ hyperparameters:
   beam_size: 5
   max_grad_norm: 5.0
   adam_epsilon : 1e-06
+  load_best_model_at_end: True

codetf/configs/training/causal_lm.yaml

@@ -6,9 +6,6 @@ hyperparameters:
   num_train_epochs: 1
   auto_find_batch_size: True
   batch_size: 4
-  max_steps: 1000
-  eval_steps: 100
-  save_steps: 1000
   logging_steps: 100
   per_device_train_batch_size: 8
   per_device_eval_batch_size: 8
@@ -30,6 +27,7 @@ hyperparameters:
   beam_size: 5
   max_grad_norm: 5.0
   adam_epsilon : 1e-06
+  load_best_model_at_end: True
 lora:
   r: 8
   lora_alpha: 32

codetf/configs/training/codet5.yaml

@@ -25,6 +25,7 @@ def load_codexglue_text_to_code_dataset(self, *args, **kwargs):
         dataset = load_dataset(dataset)
 
         train = dataset["train"]
+        train = train[:50]
         train_nl_tensors, _ = self.process_data(train["nl"])
         train_code_tensors, _ = self.process_data(train["code"])
 

codetf/data_utility/codexglue_dataset.py

@@ -9,6 +9,12 @@ class HumanEvalDataset(BaseDataset):
     def __init__(self, tokenizer, max_length=512):
         super().__init__(tokenizer, max_length)
 
+    def get_reference(self, task):
+        """Builds the reference solution for the doc (sample from the test dataset)."""
+        test_func = task["test"]
+        entry_point = f"check({task['entry_point']})"
+        return "\n" + test_func + "\n" + entry_point
+
     def load(self):
         dataset = self.dataset_config["openai_humaneval"]
 
@@ -22,9 +28,10 @@ def load(self):
             # without strip, the model generates commented codes ...
             prompts.append(self.tokenizer.eos_token + dataset[task_index]["prompt"].strip())
 
-            unit_test = dataset[task_index]["test"]
-            unit_test = re.sub(r'METADATA = {[^}]*}', '', unit_test, flags=re.MULTILINE)
-            references.append(unit_test)
+            # unit_test = dataset[task_index]["test"]
+            # unit_test = re.sub(r'METADATA = {[^}]*}', '', unit_test, flags=re.MULTILINE)
+            reference = self.get_reference(dataset[task_index])
+            references.append(reference)
 
         prompt_token_ids, prompt_attention_masks = self.process_data(prompts, padding="max_length")
 

codetf/data_utility/human_eval_dataset.py

@@ -51,6 +51,13 @@ def load_model_pipeline(model_name, model_type="base", task="sum",
 
     return model
 
+def load_model_from_path(checkpoint_path, tokenizer_path, model_name, is_eval=True, load_in_8bit=False, load_in_4bit=False):
+    model_cls = registry.get_model_class(model_name)
+    model = model_cls.from_custom(checkpoint_path=checkpoint_path, tokenizer_path=tokenizer_path, load_in_8bit=load_in_8bit, load_in_4bit=load_in_4bit)
+    if is_eval:
+        model.eval()
+
+    return model
 
 class ModelZoo:
     def __init__(self, config_files):

codetf/data_utility/stackexchange_instruction_dataset.py

@@ -45,7 +45,17 @@ def from_pretrained(model_class, model_card, load_in_8bit=False, load_in_4bit=Fa
         Build a pretrained model from default configuration file, specified by model_type.
         """
         model_config = OmegaConf.load(get_abs_path(model_class.MODEL_DICT))[model_card]
-        model_cls = model_class.load_model_from_config(model_config=model_config, load_in_8bit=load_in_8bit, load_in_4bit=load_in_4bit, weight_sharding=weight_sharding)
+        model_cls = model_class.load_huggingface_model_from_config(model_config=model_config, load_in_8bit=load_in_8bit, load_in_4bit=load_in_4bit, weight_sharding=weight_sharding)
+
+        return model_cls
+    
+
+    @classmethod
+    def from_custom(model_class, checkpoint_path, tokenizer_path, load_in_8bit=False, load_in_4bit=False):
+        """
+        Build a pretrained model from default configuration file, specified by model_type.
+        """
+        model_cls = model_class.load_custom_model(checkpoint_path, tokenizer_path, load_in_8bit=load_in_8bit, load_in_4bit=load_in_4bit)
 
         return model_cls
 

codetf/models/__init__.py

@@ -29,7 +29,7 @@ def init_tokenizer(cls, model):
         return tokenizer
 
     @classmethod
-    def load_model_from_config(model_class, model_config, load_in_8bit=False, load_in_4bit=False, weight_sharding=False):
+    def load_huggingface_model_from_config(model_class, model_config, load_in_8bit=False, load_in_4bit=False, weight_sharding=False):
         checkpoint = model_config["huggingface_url"]
 
         if load_in_8bit and load_in_4bit:
@@ -79,6 +79,35 @@ def load_model_from_config(model_class, model_config, load_in_8bit=False, load_i
             model_config=model_config,
             tokenizer=tokenizer
         )
+
+    @classmethod
+    def load_custom_model(model_class, checkpoint_path, tokenizer_path, load_in_8bit=False, load_in_4bit=False):
+
+        if load_in_8bit and load_in_4bit:
+            raise ValueError("Only one of load_in_8bit or load_in_4bit can be True. Please choose one.")
+        
+        if load_in_8bit:
+            model = AutoModelForCausalLM.from_pretrained(checkpoint_path, 
+                                        load_in_8bit=load_in_8bit, 
+                                        low_cpu_mem_usage=True,
+                                        device_map="auto")
+        elif load_in_4bit:
+            model = AutoModelForCausalLM.from_pretrained(checkpoint_path, 
+                                        load_in_4bit=load_in_4bit, 
+                                        low_cpu_mem_usage=True,
+                                        device_map="auto")
+        else:
+            model = AutoModelForCausalLM.from_pretrained(checkpoint_path, 
+                                        low_cpu_mem_usage=True,
+                                        device_map="auto")
+
+        tokenizer = model_class.init_tokenizer(tokenizer_path)
+        
+        return model_class(
+            model=model,
+            model_config=model_config,
+            tokenizer=tokenizer
+        )
 
     def forward(self, sources, max_length=512):
         encoding = self.tokenizer(sources, return_tensors='pt').to(self.device)

codetf/models/base_model.py

@@ -30,15 +30,15 @@ def init_tokenizer(cls, model):
 
 
     @classmethod
-    def load_model_from_config(model_class, model_config, load_in_8bit=False, load_in_4bit=False, weight_sharding=False):
+    def load_huggingface_model_from_config(model_class, model_config, load_in_8bit=False, load_in_4bit=False, weight_sharding=False):
 
         checkpoint = model_config["huggingface_url"]
 
         if load_in_8bit and load_in_4bit:
             raise ValueError("Only one of load_in_8bit or load_in_4bit can be True. Please choose one.")
 
         # This "device" is for the case of CodeT5plus, will be removed in the future
-        device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+        device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
         if weight_sharding:
             try:
                 # Try to download and load the json index file
@@ -85,12 +85,10 @@ def load_model_from_config(model_class, model_config, load_in_8bit=False, load_i
             else:
                 if model_config["device_map"]:
                     model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, 
-                                                load_in_4bit=load_in_4bit, 
                                                 low_cpu_mem_usage=True,
                                                 device_map=model_config["device_map"], trust_remote_code=model_config["trust_remote_code"])
                 else:
                     model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint, 
-                                                load_in_4bit=load_in_4bit, 
                                                 low_cpu_mem_usage=True,
                                                 trust_remote_code=model_config["trust_remote_code"]).to(device)
 
@@ -103,6 +101,35 @@ def load_model_from_config(model_class, model_config, load_in_8bit=False, load_i
             tokenizer=tokenizer
         )
 
+    @classmethod
+    def load_custom_model(model_class, checkpoint_path, tokenizer_path, load_in_8bit=False, load_in_4bit=False):
+
+        if load_in_8bit and load_in_4bit:
+            raise ValueError("Only one of load_in_8bit or load_in_4bit can be True. Please choose one.")
+        
+        if load_in_8bit:
+            model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint_path, 
+                                        load_in_8bit=load_in_8bit, 
+                                        low_cpu_mem_usage=True,
+                                        device_map="auto")
+        elif load_in_4bit:
+            model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint_path, 
+                                        load_in_4bit=load_in_4bit, 
+                                        low_cpu_mem_usage=True,
+                                        device_map="auto")
+        else:
+            model = AutoModelForSeq2SeqLM.from_pretrained(checkpoint_path, 
+                                        low_cpu_mem_usage=True,
+                                        device_map="auto")
+
+        tokenizer = model_class.init_tokenizer(tokenizer_path)
+        
+        return model_class(
+            model=model,
+            model_config=model_config,
+            tokenizer=tokenizer
+        )
+
 
     def forward(self, sources, max_length=512, beam_size=5):
         encoding = self.tokenizer(sources, return_tensors='pt').to(self.model.device)

codetf/models/causal_lm_models/__init__.py

@@ -23,16 +23,17 @@ 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 = 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 +55,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 +66,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/models/seq2seq_models/__init__.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/performance/model_evaluator.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/base_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()

codetf/trainer/causal_lm_trainer.py

@@ -20,4 +20,5 @@ torchvision==0.15.2
 tqdm==4.63.0
 transformers==4.30.2
 tree-sitter==0.20.1
-bitsandbytes==0.39.1
+bitsandbytes==0.39.1
+evaluate==0.4.0

codetf/trainer/codet5_trainer.py

@@ -22,7 +22,8 @@
   "tqdm==4.63.0",
   "transformers==4.30.2",
   "tree-sitter==0.20.1",
-  "bitsandbytes==0.39.1"
+  "bitsandbytes==0.39.1",
+  "evaluate==0.4.0"
 ]
 
 DEPENDENCY_LINKS = []
@@ -33,7 +34,7 @@
 
 setup(
   name = 'salesforce-codetf',
-  version = "1.0.2",
+  version = "1.0.2.2",
   py_modules = ['codetf'],
   description = 'CodeTF: A Transformer-based Library for Code Intelligence',
   author = 'Nghi D. Q. Bui',

requirements.txt

@@ -11,7 +11,7 @@
 os.environ["TOKENIZERS_PARALLELISM"] = "true"
 
 model_class = load_model_pipeline(model_name="causallm", task="pretrained",
-            model_type="codegen-350M-mono", is_eval=True,
+            model_type="codegen-350M-multi", is_eval=True,
             load_in_8bit=True, weight_sharding=False)
 
 dataset = HumanEvalDataset(tokenizer=model_class.get_tokenizer())
@@ -20,6 +20,6 @@
 problems = TensorDataset(prompt_token_ids, prompt_attention_masks)
 
 evaluator = ModelEvaluator(model_class)
-avg_pass_at_k = evaluator.evaluate_pass_k(problems=problems, unit_tests=references)
+avg_pass_at_k = evaluator.evaluate_pass_k(problems=problems, unit_tests=references, sequences_per_chunk=200, num_workers=5)
 print("Pass@k: ", avg_pass_at_k)