reformat code

Author: loubnabnl

finetuning/APPS/apps_dataset.py

@@ -10,7 +10,9 @@ class APPSBaseDataset(torch.utils.data.Dataset):
     def __init__(self, dataset, max_tokens, tokenizer_path):
         self.dataset = dataset
         self.max_tokens = max_tokens
-        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_path, use_auth_token=True)
+        self.tokenizer = AutoTokenizer.from_pretrained(
+            tokenizer_path, use_auth_token=True
+        )
         self.samples = []  # Should be set in initialize()
 
         self.initialize(self.tokenizer)

finetuning/APPS/apps_train.py

@@ -4,17 +4,12 @@
 
 import argparse
 import os
-import torch
 
+import torch
 from apps_dataset import APPSBaseDataset
 from datasets import load_dataset
-from transformers import (
-    AutoModelForCausalLM,
-    Trainer,
-    TrainingArguments,
-    logging,
-    set_seed,
-)
+from transformers import (AutoModelForCausalLM, Trainer, TrainingArguments,
+                          logging, set_seed)
 
 
 def get_args():
@@ -59,22 +54,20 @@ def run_training(args, train_data, val_data):
     training_args = TrainingArguments(
         output_dir=args.output_dir,
         dataloader_drop_last=True,
-        evaluation_strategy = "steps",
+        evaluation_strategy="steps",
         num_train_epochs=args.num_epochs,
-        max_steps = args.max_steps,
-        eval_steps = args.eval_freq,
+        max_steps=args.max_steps,
+        eval_steps=args.eval_freq,
         save_steps=args.save_freq,
         logging_steps=args.log_freq,
-
         per_device_train_batch_size=args.batch_size,
         per_device_eval_batch_size=args.batch_size,
         learning_rate=args.learning_rate,
         lr_scheduler_type=args.lr_scheduler_type,
-        warmup_steps = args.num_warmup_steps,
+        warmup_steps=args.num_warmup_steps,
         gradient_accumulation_steps=args.gradient_accumulation_steps,
         weight_decay=args.weight_decay,
         fp16=args.fp16,
-
         run_name="apps-train",
         report_to="wandb",
     )
@@ -99,8 +92,14 @@ def main(args):
     dataset.shuffle(seed=args.seed)
     data = get_dataset(dataset, args)
     train_size = int(0.95 * len(data))
-    train_data, val_data = torch.utils.data.random_split(data, [train_size, len(data) - train_size], generator=torch.Generator().manual_seed(args.seed))
-    print(f"size of training data {len(train_data)}\nsize of validation data {len(val_data)}")
+    train_data, val_data = torch.utils.data.random_split(
+        data,
+        [train_size, len(data) - train_size],
+        generator=torch.Generator().manual_seed(args.seed),
+    )
+    print(
+        f"size of training data {len(train_data)}\nsize of validation data {len(val_data)}"
+    )
     run_training(args, train_data, val_data)
 
 

finetuning/Code-to-text/train.py

@@ -1,19 +1,15 @@
 import argparse
 
 from datasets import load_dataset
-
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    Trainer,
-    TrainingArguments,
-    set_seed,
-)
+from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
+                          Trainer, TrainingArguments, set_seed)
 
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_ckpt", type=str, default="microsoft/unixcoder-base-nine")
+    parser.add_argument(
+        "--model_ckpt", type=str, default="microsoft/unixcoder-base-nine"
+    )
     parser.add_argument("--language", type=str, default="Python")
     parser.add_argument("--max_length", type=int, default=1024)
     parser.add_argument("--num_epochs", type=int, default=5)
@@ -40,7 +36,9 @@ def main():
     print("Loading tokenizer and model")
     tokenizer = AutoTokenizer.from_pretrained(args.model_ckpt)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForSequenceClassification.from_pretrained(args.model_ckpt, num_labels=2)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_ckpt, num_labels=2
+    )
     model.config.pad_token_id = model.config.eos_token_id
 
     if args.freeze:
@@ -49,13 +47,20 @@ def main():
 
     def tokenize(example):
         if args.language == "Python":
-            #remove docstring from code
+            # remove docstring from code
             chunks = example["code"].split('"""')
             code = chunks[0].strip() + chunks[2]
         else:
             code = example["code"]
-        inputs = tokenizer(code, padding="max_length", truncation=True, max_length=args.max_length)  
-        labels = tokenizer(example["docstring"], padding="max_length", truncation=True, max_length=args.max_length).input_ids
+        inputs = tokenizer(
+            code, padding="max_length", truncation=True, max_length=args.max_length
+        )
+        labels = tokenizer(
+            example["docstring"],
+            padding="max_length",
+            truncation=True,
+            max_length=args.max_length,
+        ).input_ids
         labels_with_ignore_index = []
         for labels_example in labels:
             labels_example = [label if label != 0 else -100 for label in labels_example]
@@ -99,10 +104,11 @@ def tokenize(example):
 
     print("Training...")
     trainer.train()
-    
+
     # push the model to the Hugging Face hub
     if args.push_to_hub:
         model.push_to_hub(args.model_hub_name)
 
+
 if __name__ == "__main__":
     main()

finetuning/CodeClone/train.py

@@ -3,22 +3,17 @@
 
 import numpy as np
 from datasets import ClassLabel, load_dataset
-
 from evaluate import load
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    DataCollatorWithPadding,
-    Trainer,
-    TrainerCallback,
-    TrainingArguments,
-    set_seed,
-)
+from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
+                          DataCollatorWithPadding, Trainer, TrainerCallback,
+                          TrainingArguments, set_seed)
 
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_ckpt", type=str, default="microsoft/unixcoder-base-nine")
+    parser.add_argument(
+        "--model_ckpt", type=str, default="microsoft/unixcoder-base-nine"
+    )
     parser.add_argument("--max_length", type=int, default=1024)
     parser.add_argument("--num_epochs", type=int, default=5)
     parser.add_argument("--batch_size", type=int, default=6)
@@ -52,7 +47,9 @@ def __init__(self, trainer) -> None:
     def on_epoch_end(self, args, state, control, **kwargs):
         if control.should_evaluate:
             control_copy = deepcopy(control)
-            self._trainer.evaluate(eval_dataset=self._trainer.train_dataset, metric_key_prefix="train")
+            self._trainer.evaluate(
+                eval_dataset=self._trainer.train_dataset, metric_key_prefix="train"
+            )
             return control_copy
 
 
@@ -61,21 +58,28 @@ def main():
     set_seed(args.seed)
 
     ds = load_dataset("code_x_glue_cc_clone_detection_big_clone_bench")
-    labels = ClassLabel(num_classes = 2, names=[True, False])
+    labels = ClassLabel(num_classes=2, names=[True, False])
     ds = ds.cast_column("label", labels)
 
     print("Loading tokenizer and model")
     tokenizer = AutoTokenizer.from_pretrained(args.model_ckpt)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForSequenceClassification.from_pretrained(args.model_ckpt, num_labels=2)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_ckpt, num_labels=2
+    )
     model.config.pad_token_id = model.config.eos_token_id
 
     if args.freeze:
         for param in model.roberta.parameters():
             param.requires_grad = False
 
     def tokenize(example):
-        inputs = tokenizer(example["func1"], example["func2"], truncation=True, max_length=args.max_length)  
+        inputs = tokenizer(
+            example["func1"],
+            example["func2"],
+            truncation=True,
+            max_length=args.max_length,
+        )
         return {
             "input_ids": inputs["input_ids"],
             "attention_mask": inputs["attention_mask"],
@@ -121,10 +125,11 @@ def tokenize(example):
 
     result = trainer.evaluate(eval_dataset=tokenized_datasets["test"])
     print(f"Evaluation accuracy on the test set: {result['eval_accuracy']}")
-    
+
     # push the model to the Hugging Face hub
     if args.push_to_hub:
         model.push_to_hub(args.model_hub_name)
 
+
 if __name__ == "__main__":
     main()

finetuning/CodeComplex/train.py

@@ -3,22 +3,17 @@
 
 import numpy as np
 from datasets import ClassLabel, DatasetDict, load_dataset
-
 from evaluate import load
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    DataCollatorWithPadding,
-    Trainer,
-    TrainerCallback,
-    TrainingArguments,
-    set_seed,
-)
+from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
+                          DataCollatorWithPadding, Trainer, TrainerCallback,
+                          TrainingArguments, set_seed)
 
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_ckpt", type=str, default="microsoft/unixcoder-base-nine")
+    parser.add_argument(
+        "--model_ckpt", type=str, default="microsoft/unixcoder-base-nine"
+    )
     parser.add_argument("--num_epochs", type=int, default=5)
     parser.add_argument("--batch_size", type=int, default=6)
     parser.add_argument("--gradient_accumulation_steps", type=int, default=1)
@@ -51,7 +46,9 @@ def __init__(self, trainer) -> None:
     def on_epoch_end(self, args, state, control, **kwargs):
         if control.should_evaluate:
             control_copy = deepcopy(control)
-            self._trainer.evaluate(eval_dataset=self._trainer.train_dataset, metric_key_prefix="train")
+            self._trainer.evaluate(
+                eval_dataset=self._trainer.train_dataset, metric_key_prefix="train"
+            )
             return control_copy
 
 
@@ -73,14 +70,18 @@ def main():
     print("Loading tokenizer and model")
     tokenizer = AutoTokenizer.from_pretrained(args.model_ckpt)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForSequenceClassification.from_pretrained(args.model_ckpt, num_labels=7)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_ckpt, num_labels=7
+    )
     model.config.pad_token_id = model.config.eos_token_id
 
     if args.freeze:
         for param in model.roberta.parameters():
             param.requires_grad = False
 
-    labels = ClassLabel(num_classes=7, names=list(set(train_test_validation["train"]["complexity"])))
+    labels = ClassLabel(
+        num_classes=7, names=list(set(train_test_validation["train"]["complexity"]))
+    )
 
     def tokenize(example):
         inputs = tokenizer(example["src"], truncation=True, max_length=1024)
@@ -131,10 +132,11 @@ def tokenize(example):
 
     result = trainer.evaluate(eval_dataset=tokenized_datasets["test"])
     print(f"Evaluation accuracy on the test set: {result['eval_accuracy']}")
-    
+
     # push the model to the Hugging Face hub
     if args.push_to_hub:
         model.push_to_hub(args.model_hub_name)
 
+
 if __name__ == "__main__":
     main()

finetuning/CodeDefect/train.py

@@ -3,22 +3,17 @@
 
 import numpy as np
 from datasets import ClassLabel, load_dataset
-
 from evaluate import load
-from transformers import (
-    AutoModelForSequenceClassification,
-    AutoTokenizer,
-    DataCollatorWithPadding,
-    Trainer,
-    TrainerCallback,
-    TrainingArguments,
-    set_seed,
-)
+from transformers import (AutoModelForSequenceClassification, AutoTokenizer,
+                          DataCollatorWithPadding, Trainer, TrainerCallback,
+                          TrainingArguments, set_seed)
 
 
 def get_args():
     parser = argparse.ArgumentParser()
-    parser.add_argument("--model_ckpt", type=str, default="microsoft/unixcoder-base-nine")
+    parser.add_argument(
+        "--model_ckpt", type=str, default="microsoft/unixcoder-base-nine"
+    )
     parser.add_argument("--max_length", type=int, default=1024)
     parser.add_argument("--num_epochs", type=int, default=5)
     parser.add_argument("--batch_size", type=int, default=6)
@@ -52,7 +47,9 @@ def __init__(self, trainer) -> None:
     def on_epoch_end(self, args, state, control, **kwargs):
         if control.should_evaluate:
             control_copy = deepcopy(control)
-            self._trainer.evaluate(eval_dataset=self._trainer.train_dataset, metric_key_prefix="train")
+            self._trainer.evaluate(
+                eval_dataset=self._trainer.train_dataset, metric_key_prefix="train"
+            )
             return control_copy
 
 
@@ -61,14 +58,16 @@ def main():
     set_seed(args.seed)
 
     ds = load_dataset("code_x_glue_cc_defect_detection")
-    labels = ClassLabel(num_classes = 2, names=[True, False])
+    labels = ClassLabel(num_classes=2, names=[True, False])
     ds = ds.cast_column("target", labels)
     ds = ds.rename_column("target", "label")
 
     print("Loading tokenizer and model")
     tokenizer = AutoTokenizer.from_pretrained(args.model_ckpt)
     tokenizer.pad_token = tokenizer.eos_token
-    model = AutoModelForSequenceClassification.from_pretrained(args.model_ckpt, num_labels=2)
+    model = AutoModelForSequenceClassification.from_pretrained(
+        args.model_ckpt, num_labels=2
+    )
     model.config.pad_token_id = model.config.eos_token_id
 
     if args.freeze:
@@ -128,5 +127,6 @@ def tokenize(example):
     if args.push_to_hub:
         model.push_to_hub(args.model_hub_name)
 
+
 if __name__ == "__main__":
     main()

lm_eval/base.py

@@ -1,7 +1,8 @@
-from abc import abstractmethod, ABC
-from datasets import load_dataset
+from abc import ABC, abstractmethod
 from warnings import warn
 
+from datasets import load_dataset
+
 
 class Task(ABC):
     """A task represents an entire benchmark including its dataset, problems,

lm_eval/evaluator.py

@@ -71,7 +71,9 @@ def evaluate(self, task_name):
                 if self.args.save_generations:
                     with open(self.args.save_generations_path, "w") as fp:
                         json.dump(generations, fp)
-                        print(f"generations were saved at {self.args.save_generations_path}")
+                        print(
+                            f"generations were saved at {self.args.save_generations_path}"
+                        )
                 if self.args.save_references:
                     with open("references.json", "w") as fp:
                         json.dump(references, fp)