add: adding precommit config configuration

.pre-commit-config.yaml

@@ -0,0 +1,13 @@
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v4.4.0
+    hooks:
+      - id: check-yaml
+      - id: end-of-file-fixer
+        types: [python]
+      - id: trailing-whitespace
+
+  - repo: https://github.com/psf/black
+    rev: 23.1.0
+    hooks:
+      - id: black

api/codegeex-api-example-python/generation_example.py

@@ -4,29 +4,31 @@
 
 import requests
 
-'''
+"""
 Code Generation
-'''
+"""
 API_KEY = ""  # Get from Tianqi console. 从控制台获取
 API_SECRET = ""  # Get from Tianqi console. 从控制台获取
-PROMPT = "from typing import List\n\ndef has_close_elements(numbers: List[float], threshold: float) -> bool:\n    " \
-         "\"\"\" Check if in given list of numbers, are any two numbers closer to each other than\n    given " \
-         "threshold.\n    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)\n    False\n    >>> has_close_elements(" \
-         "[1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)\n    True\n    \"\"\"\n"
+PROMPT = (
+    "from typing import List\n\ndef has_close_elements(numbers: List[float], threshold: float) -> bool:\n    "
+    '""" Check if in given list of numbers, are any two numbers closer to each other than\n    given '
+    "threshold.\n    >>> has_close_elements([1.0, 2.0, 3.0], 0.5)\n    False\n    >>> has_close_elements("
+    '[1.0, 2.8, 3.0, 4.0, 5.0, 2.0], 0.3)\n    True\n    """\n'
+)
 NUMBER = 3
 LANG = "Python"
 request_url = "https://tianqi.aminer.cn/api/v2/"
-api = 'multilingual_code_generate'
+api = "multilingual_code_generate"
 
 # Request is in json format. 指定请求参数格式为json
-headers = {'Content-Type': 'application/json'}
+headers = {"Content-Type": "application/json"}
 request_url = request_url + api
 data = {
     "apikey": API_KEY,
     "apisecret": API_SECRET,
     "prompt": PROMPT,
     "n": NUMBER,
-    "lang": LANG
+    "lang": LANG,
 }
 
 
@@ -36,5 +38,5 @@ def main():
         print(response.json())
 
 
-if __name__ == '__main__':
+if __name__ == "__main__":
     main()

codegeex/__init__.py

@@ -13,9 +13,9 @@ def get_model(
 
 
 def generate(
-    model, 
-    tokenizer: CodeGeeXTokenizer, 
-    prompt: str, 
+    model,
+    tokenizer: CodeGeeXTokenizer,
+    prompt: str,
     out_seq_length: int,
     seq_length: int = 2048,
     top_k: int = 0,
@@ -32,7 +32,7 @@ def generate(
     if verbose:
         print(f"Current prompt:\n{prompt}")
         print("N_token_prompt:", n_token_prompt)
-    
+
     generated_codes = []
     if backend == "megatron":
         token_stream = get_token_stream(
@@ -53,17 +53,22 @@ def generate(
             for j in range(micro_batch_size):
                 if is_finished[j]:
                     continue
-
-                if generated_tokens[j].cpu().numpy()[-1] == tokenizer.eos_token_id or len(generated_tokens[j]) >= out_seq_length:
+
+                if (
+                    generated_tokens[j].cpu().numpy()[-1] == tokenizer.eos_token_id
+                    or len(generated_tokens[j]) >= out_seq_length
+                ):
                     is_finished[j] = True
                     generated_tokens_ = generated_tokens[j].cpu().numpy().tolist()
-                    generated_code = tokenizer.decode_code(generated_tokens_[n_token_prompt:])
+                    generated_code = tokenizer.decode_code(
+                        generated_tokens_[n_token_prompt:]
+                    )
                     generated_code = "".join(generated_code)
                     generated_codes.append(generated_code)
                     if verbose:
                         print(f"\nGenerated code {i}:\n{generated_code}")
-                    
+
                 if all(is_finished):
                     break
 
-    return generated_codes
+    return generated_codes

codegeex/benchmark/evaluate_humaneval_x.py

@@ -16,10 +16,10 @@
 from codegeex.benchmark.execution import check_correctness
 
 LANGUAGE_NAME = {
-    "cpp"   : "CPP",
-    "go"    : "Go",
-    "java"  : "Java",
-    "js"    : "JavaScript",
+    "cpp": "CPP",
+    "go": "Go",
+    "java": "Java",
+    "js": "JavaScript",
     "python": "Python",
 }
 
@@ -29,7 +29,11 @@ def process_humaneval_test(sample, problems, example_test=False):
     language = task_id.split("/")[0].lower()
 
     prompt = sample["prompt"]
-    if example_test and "example_test" in problems[task_id] and problems[task_id]["example_test"] != "":
+    if (
+        example_test
+        and "example_test" in problems[task_id]
+        and problems[task_id]["example_test"] != ""
+    ):
         test = problems[task_id]["example_test"]
     else:
         test = problems[task_id]["test"]
@@ -39,7 +43,7 @@ def process_humaneval_test(sample, problems, example_test=False):
     if language == "python":
         code_ = []
         for line in code.split("\n"):
-            if (len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t'):
+            if len(line.strip()) > 0 and line[0] != " " and line[0] != "\t":
                 break
             code_.append(line)
         code = "\n".join(code_)
@@ -68,10 +72,21 @@ def process_humaneval_test(sample, problems, example_test=False):
             if pkg not in test_setup:
                 p = pkg.split("/")[-1]
                 if p + "." in code:
-                    other_pkgs.append(f"\"{pkg}\"")
+                    other_pkgs.append(f'"{pkg}"')
         if other_pkgs:
-            import_other_pkgs = "import (\n" + "    ".join([p + "\n" for p in other_pkgs]) + ")"
-            test_string = test_setup + "\n" + import_other_pkgs + "\n" + prompt + code + "\n" + test
+            import_other_pkgs = (
+                "import (\n" + "    ".join([p + "\n" for p in other_pkgs]) + ")"
+            )
+            test_string = (
+                test_setup
+                + "\n"
+                + import_other_pkgs
+                + "\n"
+                + prompt
+                + code
+                + "\n"
+                + test
+            )
         else:
             test_string = test_setup + "\n" + prompt + code + "\n" + test
     elif language == "rust":
@@ -97,21 +112,20 @@ def stream_jsonl_all(filename: str) -> Iterable[Dict]:
 
 
 def evaluate_functional_correctness(
-        input_file: str = None,
-        tmp_dir: str = "./",
-        n_workers: int = 32,
-        timeout: float = 500.0,
-        problem_file: str = "../data/humaneval_python.jsonl.gz",
-        out_dir: str = None,
-        k: List[int] = [1, 10, 100],
-        test_groundtruth: bool = False,
-        example_test: bool = False,
+    input_file: str = None,
+    tmp_dir: str = "./",
+    n_workers: int = 32,
+    timeout: float = 500.0,
+    problem_file: str = "../data/humaneval_python.jsonl.gz",
+    out_dir: str = None,
+    k: List[int] = [1, 10, 100],
+    test_groundtruth: bool = False,
+    example_test: bool = False,
 ):
     if example_test:
         print("Example test...")
 
-    problems = read_dataset(problem_file,
-                            dataset_type="humaneval")
+    problems = read_dataset(problem_file, dataset_type="humaneval")
     sample_jsonl = stream_jsonl_all(input_file)
 
     if example_test:
@@ -121,7 +135,9 @@ def evaluate_functional_correctness(
     if out_dir is not None:
         if not os.path.exists(out_dir):
             os.makedirs(out_dir)
-        out_file = os.path.join(out_dir, input_file.split('/')[-1].replace(".jsonl", suffix))
+        out_file = os.path.join(
+            out_dir, input_file.split("/")[-1].replace(".jsonl", suffix)
+        )
     else:
         out_file = os.path.join(input_file.replace(".jsonl", suffix))
 
@@ -149,10 +165,19 @@ def evaluate_functional_correctness(
                     lang = "js"
                 tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation")
                 sample["generation"] = sample["canonical_solution"]
-                sample["test_code"] = process_humaneval_test(sample, problems, example_test)
+                sample["test_code"] = process_humaneval_test(
+                    sample, problems, example_test
+                )
                 if sample["test_code"] is None:
                     continue
-                args = (task_id, sample, lang, timeout, tmp_dir_, completion_id[task_id])
+                args = (
+                    task_id,
+                    sample,
+                    lang,
+                    timeout,
+                    tmp_dir_,
+                    completion_id[task_id],
+                )
                 future = executor.submit(check_correctness, *args)
                 futures.append(future)
                 completion_id[task_id] += 1
@@ -164,7 +189,11 @@ def evaluate_functional_correctness(
                 lang = task_id.split("/")[0].lower()
                 if translation_mode:
                     task_id = sample["task_id"].split("/")[-1]
-                    lang = regex.findall("-to-.*-", input_file)[0].split("-to-")[-1].rstrip("-")
+                    lang = (
+                        regex.findall("-to-.*-", input_file)[0]
+                        .split("-to-")[-1]
+                        .rstrip("-")
+                    )
                     for l in LANGUAGE_NAME:
                         if l in lang:
                             lang = l
@@ -174,7 +203,9 @@ def evaluate_functional_correctness(
                     lang = "js"
                 tmp_dir_ = os.path.join(tmp_dir, lang, "evaluation")
                 sample["task_id"] = task_id
-                sample["test_code"] = process_humaneval_test(sample, problems, example_test)
+                sample["test_code"] = process_humaneval_test(
+                    sample, problems, example_test
+                )
                 if sample["test_code"] is None:
                     continue
                 if "completion_id" in sample:
@@ -208,8 +239,11 @@ def evaluate_functional_correctness(
     correct = np.array(correct)
     if evaluate_pass_at_k:
         ks = k
-        pass_at_k = {f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
-                     for k in ks if (total >= k).all()}
+        pass_at_k = {
+            f"pass@{k}": estimate_pass_at_k(total, correct, k).mean()
+            for k in ks
+            if (total >= k).all()
+        }
         print(pass_at_k)
     else:
         print("Total:", np.sum(total))
@@ -222,7 +256,7 @@ def evaluate_functional_correctness(
             for r in res:
                 fp.write((json.dumps(r[1]) + "\n").encode("utf-8"))
     else:
-        fp = open(out_file, 'w')
+        fp = open(out_file, "w")
         for res in results.values():
             for r in res:
                 fp.write(json.dumps(r[1]) + "\n")