diff --git a/README.md b/README.md index 99f959dd421dd07a0e31b6b7fdb9798c043e9b55..bd57ba27f22a8cc48596a1e3ec35d0a55d60c8ec 100644 --- a/README.md +++ b/README.md @@ -19,3 +19,5 @@ ## 同类评测项目 1. 斯坦福大学的评测：AlpacaEval Logo Leaderboard +2. https://github.com/the-crypt-keeper/can-ai-code +3. https://github.com/THUDM/CodeGeeX/tree/main/codegeex/benchmark \ No newline at end of file diff --git a/eval_set/humaneval-x/cpp/data/humaneval_cpp.jsonl.gz b/eval_set/humaneval-x/cpp/data/humaneval_cpp.jsonl.gz new file mode 100644 index 0000000000000000000000000000000000000000..e3409c83d19c3071fb5904926e6134e5d8abee18 Binary files /dev/null and b/eval_set/humaneval-x/cpp/data/humaneval_cpp.jsonl.gz differ diff --git a/eval_set/humaneval-x/cpp/evaluation/test.cpp b/eval_set/humaneval-x/cpp/evaluation/test.cpp new file mode 100644 index 0000000000000000000000000000000000000000..54cbc230dedca001c5fbea712be04aeb9a901878 --- /dev/null +++ b/eval_set/humaneval-x/cpp/evaluation/test.cpp @@ -0,0 +1,56 @@ +/* +Input to this function is a string containing multiple groups of nested parentheses. Your goal is to +separate those group into separate strings and return the vector of those. +Separate groups are balanced (each open brace is properly closed) and not nested within each other +Ignore any spaces in the input string. +>>> separate_paren_groups("( ) (( )) (( )( ))") +{"()", "(())", "(()())"} +*/ +#include +#include +#include +using namespace std; +vector separate_paren_groups(string paren_string){ + + vector all_parens; + string current_paren; + int level=0; + char chr; + int i; + for (i=0;i +bool issame(vector a,vectorb){ + if (a.size()!=b.size()) return false; + for (int i=0;i", + "markdown" : "", + "html" : "", + "php" : "// language: PHP", + "js" : "// language: JavaScript", + "javascript" : "// language: JavaScript", + "typescript" : "// language: TypeScript", + "go" : "// language: Go", + "shell" : "# language: Shell", + "rust" : "// language: Rust", + "sql" : "-- language: SQL", + "kotlin" : "// language: Kotlin", + "vb" : "' language: Visual Basic", + "ruby" : "# language: Ruby", + "pascal" : "// language: Pascal", + "r" : "# language: R", + "fortran" : "!language: Fortran", + "lean" : "-- language: Lean", + "matlab" : f"% language: Matlab", + "delphi" : "{language: Delphi}", + "scheme" : "; language: Scheme", + "basic" : "' language: Basic", + "assembly" : "; language: Assembly", + "groovy" : "// language: Groovy", + "abap" : "* language: Abap", + "gdscript" : "# language: GDScript", + "haskell" : "-- language: Haskell", + "julia" : "# language: Julia", + "elixir" : "# language: Elixir", + "excel" : "' language: Excel", + "clojure" : "; language: Clojure", + "actionscript" : "// language: ActionScript", + "solidity" : "// language: Solidity", + "powershell" : "# language: PowerShell", + "erlang" : f"% language: Erlang", + "cobol" : "// language: Cobol", +} + + +def stream_jsonl(filename: str) -> Iterable[Dict]: + """ + Parses each jsonl line and yields it as a dictionary + """ + if filename.endswith(".gz"): + with open(filename, "rb") as gzfp: + with gzip.open(gzfp, "rt") as fp: + for line in fp: + if any(not x.isspace() for x in line): + yield json.loads(line) + else: + with open(filename, "r") as fp: + for line in fp: + if any(not x.isspace() for x in line): + yield json.loads(line) + + +def write_jsonl(filename: str, data: Iterable[Dict], append: bool = False): + """ + Writes an iterable of dictionaries to jsonl + """ + if append: + mode = "ab" + else: + mode = "wb" + filename = os.path.expanduser(filename) + if filename.endswith(".gz"): + with open(filename, mode) as fp: + with gzip.GzipFile(fileobj=fp, mode="wb") as gzfp: + for x in data: + gzfp.write((json.dumps(x) + "\n").encode("utf-8")) + else: + with open(filename, mode) as fp: + for x in data: + fp.write((json.dumps(x) + "\n").encode("utf-8")) + + +def sliding_window( + prompt_tokens: list, + code_tokens: list, + seq_len: int, + sliding_stride: int, + minimum_code_len: int = 1, +) -> Iterable[Tuple[list, list]]: + """ + Generate a series of (prompt, code) pairs by sliding the window over the code. + """ + prompt_len = len(prompt_tokens) + code_len = len(code_tokens) + total_len = prompt_len + code_len + + start_idx = max(0, prompt_len - seq_len + minimum_code_len) # at least `minimum_code_len` code token should be in the window + end_idx = max(0, total_len - seq_len) + start_idx = min(start_idx, end_idx) + + for i in range(start_idx, end_idx + 1, sliding_stride): + current_prompt = prompt_tokens[i:i + seq_len] + current_code = code_tokens[max(i - prompt_len, 0):i - prompt_len + seq_len] + yield current_prompt, current_code + + if (end_idx - start_idx) % sliding_stride != 0: + current_prompt = prompt_tokens[end_idx:end_idx + seq_len] + current_code = code_tokens[max(end_idx - prompt_len, 0):end_idx - prompt_len + seq_len] + yield current_prompt, current_code diff --git a/src/evaluate_humaneval_x.py b/src/evaluate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..4aa2424652af1934ba2d19afbc4b21430dc9507f --- /dev/null +++ b/src/evaluate_humaneval_x.py @@ -0,0 +1,239 @@ +import os +import sys +import fire +import json +import gzip +import regex +import numpy as np + +from typing import * +from tqdm.auto import tqdm +from collections import defaultdict +from concurrent.futures import ThreadPoolExecutor, as_completed + +from utils import read_dataset, IMPORT_HELPER +from metric import estimate_pass_at_k +from execution import check_correctness + +LANGUAGE_NAME = { + "cpp" : "CPP", + "go" : "Go", + "java" : "Java", + "js" : "JavaScript", + "python": "Python", +} + + +def process_humaneval_test(sample, problems, example_test=False): + task_id = sample["task_id"] + 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"] != "": + test = problems[task_id]["example_test"] + else: + test = problems[task_id]["test"] + code = sample["generation"] + + # Pre-process for different languages + if language == "python": + code_ = [] + for line in code.split("\n"): + if (len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t'): + break + code_.append(line) + code = "\n".join(code_) + test_setup = "\n".join(IMPORT_HELPER["python"]) + "\n" + test_string = test_setup + prompt + code + "\n" + test + "\n" + elif language == "cpp": + test_set_up = "" + for s in IMPORT_HELPER["cpp"]: + if s not in prompt: + test_set_up += s + "\n" + test_string = test_set_up + "\n" + prompt + code + "\n" + test + elif language == "java": + test_string = prompt + code + "\n" + test + elif language == "js" or language == "javascript": + test_string = prompt + code + "\n" + test + elif language == "go": + import_string = problems[task_id]["import"] + prompt = prompt.replace(import_string, "") + if example_test and "example_test" in problems[task_id]: + test = problems[task_id]["example_test"] + else: + test = problems[task_id]["test"] + test_setup = problems[task_id]["test_setup"] + other_pkgs = [] + for pkg in IMPORT_HELPER["go"]: + if pkg not in test_setup: + p = pkg.split("/")[-1] + if p + "." in code: + 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 + else: + test_string = test_setup + "\n" + prompt + code + "\n" + test + elif language == "rust": + main = "\nfn main(){ \n } \n" + declaration = problems[task_id]["declaration"] + test_string = main + declaration + prompt + code + test + + return test_string + + +def stream_jsonl_all(filename: str) -> Iterable[Dict]: + results = [] + if filename.endswith(".gz"): + fp = gzip.open(open(filename, "rb"), "rt") + else: + fp = open(filename, "r") + for line in fp: + if any(not x.isspace() for x in line): + results.append(json.loads(line)) + fp.close() + + return results + + +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, +): + if example_test: + print("Example test...") + + problems = read_dataset(problem_file, + dataset_type="humaneval") + sample_jsonl = stream_jsonl_all(input_file) + + if example_test: + suffix = "_example_test.jsonl" + else: + suffix = "_results.jsonl" + 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)) + else: + out_file = os.path.join(input_file.replace(".jsonl", suffix)) + + if "/codegeex/benchmark/humaneval-x/" in input_file: + test_groundtruth = True + + if "-to-" in input_file: + translation_mode = True + else: + translation_mode = False + + with ThreadPoolExecutor(max_workers=n_workers) as executor: + + futures = [] + completion_id = Counter() + n_samples = 0 + results = defaultdict(list) + + if test_groundtruth: + print("Testing ground truth...") + for sample in tqdm(problems.values()): + task_id = sample["task_id"] + lang = task_id.split("/")[0].lower() + if lang == "javascript": + 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) + if sample["test_code"] is None: + continue + 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 + n_samples += 1 + else: + print("Reading samples...") + for sample in tqdm(sample_jsonl): + task_id = sample["task_id"] + 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("-") + for l in LANGUAGE_NAME: + if l in lang: + lang = l + break + task_id = f"{LANGUAGE_NAME[lang]}/{task_id}" + if lang == "javascript": + 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) + if sample["test_code"] is None: + continue + if "completion_id" in sample: + completion_id_ = sample["completion_id"] + else: + completion_id_ = completion_id[task_id] + args = (task_id, sample, lang, timeout, tmp_dir_, completion_id_) + future = executor.submit(check_correctness, *args) + futures.append(future) + completion_id[task_id] += 1 + n_samples += 1 + + print(completion_id) + if len(completion_id) == len(problems): + evaluate_pass_at_k = True + else: + evaluate_pass_at_k = False + + print("Running test suites...") + for future in tqdm(as_completed(futures), total=len(futures)): + result = future.result() + results[result["task_id"]].append((result["completion_id"], result)) + + # Calculate pass@k. + total, correct = [], [] + for result in results.values(): + passed = [r[1]["passed"] for r in result] + total.append(len(passed)) + correct.append(sum(passed)) + total = np.array(total) + 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()} + print(pass_at_k) + else: + print("Total:", np.sum(total)) + print("Correct:", np.sum(correct)) + + print("Writing to: ", out_file) + if out_file.endswith(".gz"): + fp = gzip.GzipFile(fileobj=open(out_file, "wb"), mode="wb") + for res in results.values(): + for r in res: + fp.write((json.dumps(r[1]) + "\n").encode("utf-8")) + else: + fp = open(out_file, 'w') + for res in results.values(): + for r in res: + fp.write(json.dumps(r[1]) + "\n") + fp.close() + + print("Evaluation finished.") + + +def main(): + fire.Fire(evaluate_functional_correctness) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/execution.py b/src/execution.py new file mode 100644 index 0000000000000000000000000000000000000000..cbdf14f4954d9f458d1e2c8e5fb8016865d6244a --- /dev/null +++ b/src/execution.py @@ -0,0 +1,549 @@ +import contextlib +import faulthandler +import io +import multiprocessing +import os +import platform +import signal +import random +import subprocess +import tempfile +import gzip +import json +from typing import * + +def dicts_to_jsonl(data_list: list, filename: str, compress: bool = True) -> None: + """ + Method saves list of dicts into jsonl file. + :param data: (list) list of dicts to be stored, + :param filename: (str) path to the output file. If suffix .jsonl is not given then methods appends + .jsonl suffix into the file. + :param compress: (bool) should file be compressed into a gzip archive? + """ + sjsonl = '.jsonl' + sgz = '.gz' + # Check filename + if not filename.endswith(sjsonl): + filename = filename + sjsonl + # Save data + + if compress: + filename = filename + sgz + with gzip.open(filename, 'w') as compressed: + for ddict in data_list: + jout = json.dumps(ddict) + '\n' + jout = jout.encode('utf-8') + compressed.write(jout) + else: + with open(filename, 'w') as out: + for ddict in data_list: + jout = json.dumps(ddict) + '\n' + out.write(jout) + + +def check_correctness( + task_id: str, + sample: dict, + language_type: str, + timeout: float = 3.0, + tmp_dir: str = None, + completion_id: Optional[int] = None, +) -> Dict: + """ + Evaluates the functional correctness of a completion by running the test + suite provided in the problem. + """ + + def unsafe_execute(tmp_dir): + random_id = random.uniform(1, 1000) + if "python" in language_type.lower(): + with create_tempdir(): + + # These system calls are needed when cleaning up tempdir. + import os + import shutil + rmtree = shutil.rmtree + rmdir = os.rmdir + chdir = os.chdir + + # Disable functionalities that can make destructive changes to the test. + reliability_guard() + + try: + exec_globals = {} + with swallow_io(): + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec(sample["test_code"], exec_globals) + result.append("passed") + except TimeoutException: + result.append("timed out") + except AssertionError as e: + result.append(f"failed: AssertionError") + except BaseException as e: + result.append(f"failed: {e}") + + # Needed for cleaning up. + shutil.rmtree = rmtree + os.rmdir = rmdir + os.chdir = chdir + + elif "go" in language_type.lower(): + assert tmp_dir is not None, "Go should be evaluated in a dir where necessary module files installed." + + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"main_test.go", 'w').write(sample["test_code"]) + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["go", "test", f"-timeout={timeout}s", "main_test.go"], timeout=timeout, capture_output=True) + + if exec_result.returncode == 0: + result.append("passed") + else: + if exec_result.stderr: + try: + err = exec_result.stderr.decode() + except: + err = exec_result.stderr + else: + try: + err = exec_result.stdout.decode() + except: + err = exec_result.stdout + result.append(f"failed: {err}") + + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "js" in language_type.lower(): + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"test.js", 'w').write(sample["test_code"]) + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["node", "test.js"], timeout=timeout, capture_output=True) + + if exec_result.stderr.decode(): + err = exec_result.stderr.decode() + result.append(f"failed: {err}") + elif exec_result.stdout.decode(): + err = exec_result.stdout.decode() + result.append(f"failed: {err}") + else: + result.append("passed") + + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "cpp" in language_type.lower(): + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(f"test.cpp", 'w').write(sample["test_code"]) + if "162" in task_id: + compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp", "-lcrypto", "-lssl"], + timeout=timeout, + capture_output=True) + else: + compilation_result = subprocess.run(["/usr/bin/g++", "-std=c++11", "test.cpp"], timeout=timeout, + capture_output=True) + if compilation_result.returncode != 0: + if compilation_result.stderr: + err = compilation_result.stderr.decode() + else: + err = compilation_result.stdout.decode() + result.append(f"failed: compilation error: {err}") + else: + try: + exec_result = None + with time_limit(timeout): + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + exec_result = subprocess.run(["./a.out"], timeout=timeout, capture_output=True) + + if exec_result.returncode == 0: + result.append("passed") + else: + if exec_result.stderr: + try: + err = exec_result.stderr.decode() + except: + err = exec_result.stderr + else: + try: + err = exec_result.stdout.decode() + except: + err = exec_result.stdout + result.append(f"failed: {err}") + except TimeoutException: + result.append("timed out") + + shutil.rmtree(tmp_dir) + elif "rust" in language_type.lower(): + import os + + WD: str = os.path.dirname(os.path.abspath(__file__)) + RUST_DIR: str = os.path.join(WD, "rust") + RUST_SRC: str = os.path.join(RUST_DIR, "src") + RUST_BIN: str = os.path.join(RUST_SRC, "bin") + RUST_TMP_DIR: str = os.path.join(RUST_DIR, "tmp") + RUST_LOGS: str = os.path.join(RUST_TMP_DIR, "logs") + RUST_EXT: str = ".rs" + + # Create mandatory tmp directories + os.makedirs(RUST_TMP_DIR, exist_ok=True) + os.makedirs(RUST_LOGS, exist_ok=True) + os.makedirs(RUST_SRC, exist_ok=True) + os.makedirs(RUST_BIN, exist_ok=True) + + with tempfile.NamedTemporaryFile(dir = RUST_BIN, delete=False) as f: + #temporal file name + file_prefix = sample["task_id"].lower().replace("/", "_") + file_name:str = file_prefix +RUST_EXT + + os.rename(f.name, os.path.join(RUST_BIN, file_name)) + + # Sample to pure Rust function + rust_code: str = sample["test_code"] + + # dump the rust source code in the target temporal file + f.write(rust_code.encode('utf-8')) + + # Proceed towards Rust binaries compilation. Therefore move to Rust module root dir. + os.chdir(RUST_DIR) + + # Two possible outcomes + # Pass OR Fail compilation + log_filename: str = file_prefix + ".jsonl" + log_path: str = os.path.join(RUST_LOGS, log_filename) + cargo_check: str = "cargo check --bin " + file_prefix + " --message-format json >> " + log_path + # Compilation build status + returned_val_compilation: int + + # Overwrite file content + if os.path.exists(log_path): + if(file_size := os.path.getsize(log_path)) >= 0: + os.remove(log_path) + returned_val_compilation = os.system(cargo_check) + + else: + returned_val_compilation = os.system(cargo_check) + + # 0 means success + if returned_val_compilation == 0: + + #Execution pipeline + cargo_test: str = "cargo test --bin " +file_prefix+ " --message-format json >> " + log_path + returned_val_execution = os.system(cargo_test) + + if returned_val_execution == 0: + result.append("passed") + else: + result.append(f"failed: execution error") + + else: + result.append(f"failed: compilation error") + + + elif "java" in language_type.lower(): + assert tmp_dir is not None, "Java should be evaluated in a temporary dir." + + import os + import shutil + + if "tmp" not in tmp_dir: + tmp_dir = os.path.join(tmp_dir, "tmp") + tmp_dir = os.path.join(tmp_dir, f"{task_id.replace('/', '-')}-{random_id}") + if not os.path.exists(tmp_dir): + os.makedirs(tmp_dir) + + os.chdir(tmp_dir) + open(os.path.join(tmp_dir, "Main.java"), 'w').write(sample["test_code"]) + res = "failed: unknown error" + compile_returncode = -1 + for _ in range(5): + try: + compilation_result = subprocess.run(['javac', os.path.join(tmp_dir, "Main.java")], timeout=5, + capture_output=True) + compile_returncode = compilation_result.returncode + break + except subprocess.TimeoutExpired as e: + continue + if compile_returncode != 0: + res = "failed: compilation error" + else: + exec_result = None + try: + # WARNING + # This program exists to execute untrusted model-generated code. Although + # it is highly unlikely that model-generated code will do something overtly + # malicious in response to this test suite, model-generated code may act + # destructively due to a lack of model capability or alignment. + # Users are strongly encouraged to sandbox this evaluation suite so that it + # does not perform destructive actions on their host or network. + # Once you have read this disclaimer and taken appropriate precautions, + # uncomment the following line and proceed at your own risk: + # exec_result = subprocess.run([f'java', '-cp', tmp_dir, 'Main'], timeout=timeout, capture_output=True) + if exec_result.returncode == 0: + res = "passed" + elif exec_result.returncode == 1: + if "AssertionError" in exec_result.stderr.decode('unicode-escape'): + res = "failed: wrong answer" + else: + res = f"failed: {exec_result.stderr.decode()}" + except subprocess.TimeoutExpired as e: + res = "time out" + except BaseException as e: + res = f"failed: {e}" + result.append(res) + + shutil.rmtree(tmp_dir) + + manager = multiprocessing.Manager() + result = manager.list() + + p = multiprocessing.Process(target=unsafe_execute, args=(tmp_dir,)) + p.start() + p.join(timeout=timeout + 1) + if p.is_alive(): + p.kill() + + if not result: + result.append("timed out") + + return { + "task_id" : task_id, + "completion_id": completion_id, + "test_code" : sample["test_code"], + "prompt" : sample["prompt"], + "generation" : sample["generation"], + "result" : result[0], + "passed" : result[0] == "passed", + "finish" : -1 if "finish" not in sample else sample["finish"], + "file" : "" if "file" not in sample else sample["file"], + "output" : [] if "output" not in sample else sample["output"], + } + +# Copyright (c) OpenAI (https://openai.com) + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# ============================================================================ +@contextlib.contextmanager +def time_limit(seconds: float): + def signal_handler(signum, frame): + raise TimeoutException("Timed out!") + + signal.setitimer(signal.ITIMER_REAL, seconds) + signal.signal(signal.SIGALRM, signal_handler) + try: + yield + finally: + signal.setitimer(signal.ITIMER_REAL, 0) + + +@contextlib.contextmanager +def swallow_io(): + stream = WriteOnlyStringIO() + with contextlib.redirect_stdout(stream): + with contextlib.redirect_stderr(stream): + with redirect_stdin(stream): + yield + + +@contextlib.contextmanager +def create_tempdir(): + with tempfile.TemporaryDirectory() as dirname: + with chdir(dirname): + yield dirname + + +class TimeoutException(Exception): + pass + + +class WriteOnlyStringIO(io.StringIO): + """ StringIO that throws an exception when it's read from """ + + def read(self, *args, **kwargs): + raise IOError + + def readline(self, *args, **kwargs): + raise IOError + + def readlines(self, *args, **kwargs): + raise IOError + + def readable(self, *args, **kwargs): + """ Returns True if the IO object can be read. """ + return False + + +class redirect_stdin(contextlib._RedirectStream): # type: ignore + _stream = 'stdin' + + +@contextlib.contextmanager +def chdir(root): + if root == ".": + yield + return + cwd = os.getcwd() + os.chdir(root) + try: + yield + except BaseException as exc: + raise exc + finally: + os.chdir(cwd) + + +def reliability_guard(maximum_memory_bytes: Optional[int] = None): + """ + This disables various destructive functions and prevents the generated code + from interfering with the test (e.g. fork bomb, killing other processes, + removing filesystem files, etc.) + + WARNING + This function is NOT a security sandbox. Untrusted code, including, model- + generated code, should not be blindly executed outside of one. See the + Codex paper for more information about OpenAI's code sandbox, and proceed + with caution. + """ + + if maximum_memory_bytes is not None: + import resource + resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes)) + resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes)) + if not platform.uname().system == 'Darwin': + resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes)) + + faulthandler.disable() + + import builtins + builtins.exit = None + builtins.quit = None + + import os + os.environ['OMP_NUM_THREADS'] = '1' + + os.kill = None + os.system = None + os.putenv = None + os.remove = None + os.removedirs = None + os.rmdir = None + os.fchdir = None + os.setuid = None + os.fork = None + os.forkpty = None + os.killpg = None + os.rename = None + os.renames = None + os.truncate = None + os.replace = None + os.unlink = None + os.fchmod = None + os.fchown = None + os.chmod = None + os.chown = None + os.chroot = None + os.fchdir = None + os.lchflags = None + os.lchmod = None + os.lchown = None + os.getcwd = None + os.chdir = None + + import shutil + shutil.rmtree = None + shutil.move = None + shutil.chown = None + + import subprocess + subprocess.Popen = None # type: ignore + + __builtins__['help'] = None + + import sys + sys.modules['ipdb'] = None + sys.modules['joblib'] = None + sys.modules['resource'] = None + sys.modules['psutil'] = None + sys.modules['tkinter'] = None diff --git a/src/gather_output.py b/src/gather_output.py new file mode 100644 index 0000000000000000000000000000000000000000..db63986f3c078447924a935861f6a8daa2e6269f --- /dev/null +++ b/src/gather_output.py @@ -0,0 +1,56 @@ +import os +import sys +import fire +import glob + + +def gather_output( + output_dir: str = "./output", + output_prefix: str = None, + if_remove_rank_files: int = 0, +): + if output_prefix is None: + output_list = glob.glob(output_dir + "/*") + else: + output_list = glob.glob(os.path.join(output_dir, output_prefix + "*")) + + for output_file in output_list: + if "rank0" in output_file: + output_prefix_ = output_file.split("_rank0.jsonl")[0] + rank_files = glob.glob(output_prefix_ + "_rank*") + with open(output_prefix_ + ".jsonl", "w") as f_out: + for rank_file in rank_files: + with open(rank_file, "r") as f_in: + for line in f_in: + f_out.write(line) + if if_remove_rank_files: + os.remove(rank_file) + print(f"Removing {rank_file}...") + + if output_prefix is None: + output_list = glob.glob(output_dir + "/*") + else: + output_list = glob.glob(os.path.join(output_dir, output_prefix + "*")) + + for output_file in output_list: + if "rank" in output_file or "_unfinished" in output_file or "all" in output_file or "_result" in output_file: + continue + if "_finished" not in output_file: + continue + output_prefix_ = output_file.split("_finished.jsonl")[0] + files = [output_file, output_prefix_ + "_unfinished.jsonl"] + with open(output_prefix_ + "_all.jsonl", "w") as f_out: + for f in files: + with open(f, "r") as f_in: + for line in f_in: + f_out.write(line) + + print("Gathering finished. Saved in {}".format(output_prefix_ + "_all.jsonl")) + + +def main(): + fire.Fire(gather_output) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/generate_humaneval_x.py b/src/generate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..056ed440ea89e5e4b0a3c648bcd97eea187291b5 --- /dev/null +++ b/src/generate_humaneval_x.py @@ -0,0 +1,89 @@ +import argparse +import logging +import os +import random +import socket +import time +import json +import torch + +from utils import read_dataset, process_extra_prompt +from utils import cleanup_code + +logging.getLogger("torch").setLevel(logging.WARNING) + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") + parser = argparse.ArgumentParser() + parser.add_argument( + "--input_path", + type=str, + default="eval_set/humaneval-x" + ) + parser.add_argument( + "--language_type", + type=str, + default="python" + ) + parser.add_argument( + "--model_name", + type=str, + default="chatgpt", + help='supported model in [chatgpt,chatglm2].' + ) + parser.add_argument( + "--output_prefix", + type=str, + default="chatgpt" + ) + args = parser.parse_known_args()[0] + + output_file_path = args.output_prefix + f"_finished_rank{args.gen_rank}.jsonl" + + entries = read_dataset(args.input_path, args.language_type, dataset_type="humaneval") + for entry in entries.values(): + entry["prompt"] = process_extra_prompt(entry["prompt"], args.language_type) + + model_inference = None + if args.model_name == "chatgpt": + from inference.chatgpt_inference import ChatGPTInference + model_inference = ChatGPTInference() + elif args.model_name == "chatglm2": + from inference.chatglm2_inference import GLMInference + model_inference = GLMInference() + else: + print(f"{args.model_name} not supported.") + + start_time = time.perf_counter() + num_finished = 0 + with open(output_file_path, "w") as f: + for entry in entries.values(): + prompt = entry["prompt"] + generated_tokens = model_inference.inference(prompt) + generated_code = cleanup_code(generated_code, language_type=args.language_type) + f.write( + json.dumps( + { + "task_id": entry['task_id'], + "prompt": prompt, + "generation": generated_code, + "scores": 0.0, + "finish": 2, + "output": generated_tokens, + } + ) + + "\n" + ) + f.flush() + num_finished += 1 + + time_elapsed = time.perf_counter() - start_time + time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished + print( + f"finished {num_finished}, " + f"elapsed {time_elapsed:.4f}", + f"speed {time_per_sampple:.4f}s/sample", + f"remaining {len(entries) - num_finished:.4f}", + flush=True, + ) diff --git a/src/inference/__init__.py b/src/inference/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/src/inference/chatglm2_inference.py b/src/inference/chatglm2_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..c7669b7db33fb8a93887ea7eb1b97f9a355b4bcc --- /dev/null +++ b/src/inference/chatglm2_inference.py @@ -0,0 +1,63 @@ +import os +import json +import torch +import logging +from transformers import AutoModel, AutoTokenizer +from inference import Inference + +logger = logging.getLogger(__name__) + + +class GLMInference(Inference): + + def __int__(self): + self.params_url = "llm_set/params/chatgpt.json" + self.paras_dict = self.get_params() + self.paras_dict.update(self.paras_base_dict) + + self.temperature = self.paras_dict.get("temperature") + self.quantize = self.paras_dict.get("quantize") + self.device = self.paras_dict.get("device") + self.model_path = self.paras_dict.get("model_path") + if self.device == "cpu": + self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True, + ignore_mismatched_sizes=True).float() + self.DEV = torch.device('cpu') + else: + os.environ["CUDA_VISIBLE_DEVICES"] = self.device.split(":")[1] + self.DEV = torch.device(self.device) + if self.quantize: + self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).quantize(8).cuda() + else: + self.model = AutoModel.from_pretrained(self.model_path, trust_remote_code=True).cuda() + + self.tokenizer = AutoTokenizer.from_pretrained(self.model_path, trust_remote_code=True) + self.model = self.model.eval() + self.max_length = self.paras_dict.get("max_length") + self.min_length = self.paras_dict.get("min_length") + self.top_p = self.paras_dict.get("top_p") + self.top_k = self.paras_dict.get("top_k") + + def get_params(self): + if not os.path.exists(self.params_url): + logger.error(f"params_url:{self.params_url} is not exists.") + content = open(self.params_url).read() + return json.loads(content) + + def inference(self, message): + input_ids = self.tokenizer.encode(message, return_tensors="pt").to(self.DEV) + + sentence = None + with torch.no_grad(): + generation_output = self.model.generate( + input_ids, + do_sample=True, + min_length=self.min_length, + max_length=self.max_length, + top_p=self.top_p, + top_k=self.top_k, + temperature=self.temperature + ) + output = self.tokenizer.decode([el.item() for el in generation_output[0]]) + sentence = output.split("ChatCSDN:")[1].strip() + return sentence diff --git a/src/inference/chatgpt_inference.py b/src/inference/chatgpt_inference.py new file mode 100644 index 0000000000000000000000000000000000000000..87928f4c01484a75f46df1b4674cb6db3cf3066d --- /dev/null +++ b/src/inference/chatgpt_inference.py @@ -0,0 +1,45 @@ + +import os +import json +import logging +import quote +import requests +import traceback +from inference import Inference + + +logger = logging.getLogger(__name__) + + +class ChatGPTInference(Inference): + def __int__(self): + self.params_url = "llm_set/params/chatgpt.json" + self.paras_dict = self.get_params() + self.paras_dict.update(self.paras_base_dict) + self.gpt_url = self.paras_dict.get("url") + self.id = self.paras_dict.get("id") + self.stream = self.paras_dict.get("stream") + self.temperature = self.paras_dict.get("temperature") + self.timeout = self.paras_dict.get("timeout") + + def get_params(self): + if not os.path.exists(self.params_url): + logger.error(f"params_url:{self.params_url} is not exists.") + content = open(self.params_url).read() + return json.loads(content) + + def inference(self, query_text): + query_text = quote(query_text) + param_str = f"?id={self.gpt_url}&stream={self.stream}&temperature={self.temperature}&q={query_text}" + try: + resp = requests.get(self.gpt_url + param_str, timeout=self.timeout) + except Exception as e: + logger.error(traceback.format_exc()) + logger.error(e) + return None + + if resp.status_code == 200: + resp_json = json.loads(resp.text) + return "\n".join(resp_json["message"]["content"]["parts"]) + else: + return None diff --git a/src/inference/inference.py b/src/inference/inference.py new file mode 100644 index 0000000000000000000000000000000000000000..c27d3b6bf1ef529ca2b14701f5a943b2d5ec72dd --- /dev/null +++ b/src/inference/inference.py @@ -0,0 +1,19 @@ + +import os +import json +import logging + + +logger = logging.getLogger(__name__) + + +class Inference(object): + def __int__(self): + self.params_base_url = "llm_set/params/default.json" + self.paras_base_dict = self.get_paras_base() + + def get_paras_base(self): + if not os.path.exists(self.params_base_url): + logger.error(f"params_base_url:{self.params_base_url} is not exists.") + content = open(self.params_base_url).read() + return json.loads(content) diff --git a/src/inspect_result.py b/src/inspect_result.py new file mode 100644 index 0000000000000000000000000000000000000000..2fb6069c50cb75e7469069c7ea65108d9ac84da4 --- /dev/null +++ b/src/inspect_result.py @@ -0,0 +1,286 @@ +import os +import sys +import fire +import json +import glob +import numpy as np +import pandas as pd + +from collections import defaultdict +from metric import estimate_pass_at_k + +error_types = { + "python" : [ + "accepted", + "assertion error", + "undefined error", + "runtime error", + "syntax error", + "timeout error", + "type error", + ], + "java" : [ + "accepted", + "compilation error", + "assertion error", + "timeout error", + "index error", + "class cast error", + "stack overflow", + "null pointer error", + "unsupported operation", + "number format error", + "no such element", + "illegal argument", + "out of memory", + "arithmetic error", + "others", + ], + "cpp" : [ + "accepted", + "compilation error", + "assertion error", + "range error", + "invalid argument", + "pointer error", + "out of memory", + "package error", + "others", + ], + "javascript": [ + "accepted", + "assertion error", + "undefined error", + "runtime error", + "syntax error", + "timeout error", + "range error", + "type error", + ], + "go" : [ + "accepted", + "assertion error", + "undefined error", + "runtime error", + "syntax error", + "timeout error", + "type error", + "notused error", + ], +} + + +def inspect_result( + input_dir: str = None, + input_file: str = None, + output_dir: str = None, + pass_at_k_outpath: str = None, +): + if input_dir is not None: + input_files = glob.glob(input_dir + "/*_results.jsonl") + else: + input_files = [input_file] + + if output_dir is not None: + if not os.path.exists(output_dir): + os.makedirs(output_dir) + else: + output_dir = "/" + + pass_at_k_outs = [] + for input_file in input_files: + result_stats = defaultdict(dict) + df = None + incompleted = False + with open(input_file, "r") as f: + for i, line in enumerate(f): + obj = json.loads(line) + task_id = obj["task_id"] + language_type = task_id.split("/")[0].lower() + if language_type not in error_types: + if language_type == "humaneval": + language_type = "python" + elif language_type == "C++": + language_type = "cpp" + else: + incompleted = True + break + if task_id not in result_stats.keys(): + default_stats = {} + default_stats["task_id"] = task_id + default_stats["n_sample"] = 0 + for k in error_types[language_type]: + default_stats[k] = 0 + + result_stats[task_id] = default_stats.copy() + if df is None: + df = pd.DataFrame(columns=error_types[language_type]) + result_stats[task_id]["n_sample"] += 1 + + if "failed" in obj["result"]: + error = obj["result"] + if language_type == "python": + if "assertionerror" in error.lower(): + result_stats[task_id]["assertion error"] += 1 + elif "syntax" in error.lower() or "indent" in error.lower() or "literal" in error.lower(): + result_stats[task_id]["syntax error"] += 1 + elif "not defined" in error.lower(): + result_stats[task_id]["undefined error"] += 1 + elif "timeout" in error.lower(): + result_stats[task_id]["timeout error"] += 1 + elif "type" in error.lower(): + result_stats[task_id]["type error"] += 1 + else: + result_stats[task_id]["runtime error"] += 1 + + elif language_type == "java": + if "wrong answer" in error: + result_stats[task_id]["assertion error"] += 1 + elif "compilation error" in error: + result_stats[task_id]["compilation error"] += 1 + elif "time out" in error: + result_stats[task_id]["timeout error"] += 1 + elif "IndexOutOfBounds" in error: + result_stats[task_id]["index error"] += 1 + elif "UnsupportedOperation" in error: + result_stats[task_id]["unsupported operation"] += 1 + elif "ClassCast" in error: + result_stats[task_id]["class cast error"] += 1 + elif "NullPointer" in error: + result_stats[task_id]["null pointer error"] += 1 + elif "NumberFormat" in error: + result_stats[task_id]["number format error"] += 1 + elif "NoSuchElement" in error: + result_stats[task_id]["no such element"] += 1 + elif "StackOverflow" in error: + result_stats[task_id]["stack overflow"] += 1 + elif "Arithmetic" in error: + result_stats[task_id]["arithmetic error"] += 1 + elif "OutOfMemory" in error: + result_stats[task_id]["out of memory"] += 1 + elif "IllegalArgument" in error: + result_stats[task_id]["illegal argument"] += 1 + else: + result_stats[task_id]["others"] += 1 + + elif language_type == "cpp": + if "compilation error" in error.lower(): + result_stats[task_id]["compilation error"] += 1 + elif "int main(): assertion" in error.lower(): + result_stats[task_id]["assertion error"] += 1 + elif "out_of_range" in error.lower(): + result_stats[task_id]['range error'] += 1 + elif "corrupted top size" in error.lower(): + result_stats[task_id]['range error'] += 1 + elif "length_error" in error.lower(): + result_stats[task_id]['range error'] += 1 + elif "invalid_argument" in error.lower(): + result_stats[task_id]['invalid argument'] += 1 + elif "invalid pointer" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "double free" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "free()" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "logic_error" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "sysmalloc: assertion" in error.lower(): + result_stats[task_id]['pointer error'] += 1 + elif "stack smashing" in error.lower(): + result_stats[task_id]['out of memory'] += 1 + elif "bad_alloc" in error.lower(): + result_stats[task_id]['out of memory'] += 1 + elif "terminate called after throwing an instance of" in error.lower(): + result_stats[task_id]['package error'] += 1 + else: + result_stats[task_id]["others"] += 1 + + elif language_type == "javascript": + if "Assertion failed" in error: + result_stats[task_id]["assertion error"] += 1 + elif "SyntaxError" in error: + result_stats[task_id]["syntax error"] += 1 + elif "ReferenceError" in error: + result_stats[task_id]["undefined error"] += 1 + elif "timed out" in error: + result_stats[task_id]["timeout error"] += 1 + elif "TypeError" in error: + result_stats[task_id]["type error"] += 1 + elif "RangeError" in error: + result_stats[task_id]["range error"] += 1 + else: + result_stats[task_id]["runtime error"] += 1 + + elif language_type == "go": + if "Error: \tNot equal:" in error: + result_stats[task_id]["assertion error"] += 1 + elif "undefined" in error: + result_stats[task_id]["undefined error"] += 1 + elif "expected" in error and "found" in error: + result_stats[task_id]["syntax error"] += 1 + elif "illegal" in error: + result_stats[task_id]["syntax error"] += 1 + elif "unexpected" in error: + result_stats[task_id]["syntax error"] += 1 + elif "FAIL" in error: + result_stats[task_id]["runtime error"] += 1 + elif "timed out" in error: + result_stats[task_id]["timeout error"] += 1 + elif "not used" in error: + result_stats[task_id]['notused error'] += 1 + elif "type" in error: + result_stats[task_id]['type error'] += 1 + else: + incompleted = True + break + else: + if obj["passed"]: + result_stats[task_id]["accepted"] += 1 + + if incompleted: + print(f"Language not supported, aborted. {input_file}") + else: + try: + total, correct = [], [] + for k, res in result_stats.items(): + total.append(res["n_sample"]) + correct.append(res["accepted"]) + df_res = pd.DataFrame(res, index=[int(k.split("/")[-1])]) + df = pd.concat([df, df_res], axis=0) + + total = np.array(total) + correct = np.array(correct) + + ks = [1, 10, 100, 1000] + 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) + pass_at_k["file"] = input_file + pass_at_k["n"] = res["n_sample"] + pass_at_k_outs.append(pass_at_k) + + output_prefix = input_file.split("/")[-1].split(".jsonl")[0] + output_file = os.path.join(output_dir, output_prefix + "_stats.xlsx") + df = df.sort_index(ascending=True) + df.to_excel(output_file) + + print(f"Stats saved in {output_file}") + except Exception as e: + print(e) + print(f"Data incompleted, aborted. {input_file}") + + if pass_at_k_outpath is not None: + jsonl_path = os.path.join(output_dir, pass_at_k_outpath) + with open(jsonl_path, "w") as f_out: + for p in pass_at_k_outs: + f_out.write(json.dumps(p) + "\n") + print(f"Pass at k saved in {jsonl_path}") + + +def main(): + fire.Fire(inspect_result) + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/src/metric.py b/src/metric.py new file mode 100644 index 0000000000000000000000000000000000000000..79d6db3763d0d38e0bb206544fa21ec0d0b6d8c9 --- /dev/null +++ b/src/metric.py @@ -0,0 +1,50 @@ +# Copyright (c) OpenAI (https://openai.com) + +# Permission is hereby granted, free of charge, to any person obtaining a copy +# of this software and associated documentation files (the "Software"), to deal +# in the Software without restriction, including without limitation the rights +# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +# copies of the Software, and to permit persons to whom the Software is +# furnished to do so, subject to the following conditions: + +# The above copyright notice and this permission notice shall be included in +# all copies or substantial portions of the Software. + +# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN +# THE SOFTWARE. +# ============================================================================ +import itertools +import numpy as np + +from typing import * + + +def estimate_pass_at_k( + num_samples: Union[int, List[int], np.ndarray], + num_correct: Union[List[int], np.ndarray], + k: int +) -> np.ndarray: + """ + Estimates pass@k of each problem and returns them in an array. + """ + + def estimator(n: int, c: int, k: int) -> float: + """ + Calculates 1 - comb(n - c, k) / comb(n, k). + """ + if n - c < k: + return 1.0 + return 1.0 - np.prod(1.0 - k / np.arange(n - c + 1, n + 1)) + + if isinstance(num_samples, int): + num_samples_it = itertools.repeat(num_samples, len(num_correct)) + else: + assert len(num_samples) == len(num_correct) + num_samples_it = iter(num_samples) + + return np.array([estimator(int(n), int(c), k) for n, c in zip(num_samples_it, num_correct)]) diff --git a/src/translate_humaneval_x.py b/src/translate_humaneval_x.py new file mode 100644 index 0000000000000000000000000000000000000000..dcc39d73ea050ceb2d97c7ee33fb20c704dfc158 --- /dev/null +++ b/src/translate_humaneval_x.py @@ -0,0 +1,480 @@ +import os +import zmq +import time +import torch +import random +import socket +import logging +import argparse + +from typing import * +from utils import read_translation_dataset +from codegeex.megatron import get_args +from codegeex.megatron.inference import run_generation_distributed, model_provider +from codegeex.megatron.initialize import initialize_megatron + +logging.getLogger("torch").setLevel(logging.WARNING) + + +def add_code_generate_args(parser): + """Code generation arguments.""" + group = parser.add_argument_group(title="code generation") + + group.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + group.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + group.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + group.add_argument( + "--master-port", + type=int, + default=5666, + ) + group.add_argument( + "--temperature", + type=float, + default=1.0, + help="Sampling temperature.", + ) + group.add_argument( + "--greedy", + action="store_true", + default=False, + help="Use greedy sampling.", + ) + group.add_argument( + "--top-p", + type=float, + default=0.0, + help="Top p sampling.", + ) + group.add_argument( + "--top-k", + type=int, + default=0, + help="Top k sampling.", + ) + group.add_argument( + "--out-seq-length", + type=int, + default=1024, + help="Size of the output generated text.", + ) + group.add_argument( + "--input-path", + type=str, + default="./benchmark/humaneval/HumanEval.jsonl", + help="Get input path", + ) + group.add_argument( + "--num-samples", + type=int, + default=0, + help="Number of samples to generate", + ) + group.add_argument( + "--recompute", + action="store_true", + help="During generation recompute all attention " + "instead of using previously computed keys/values.", + ) + group.add_argument( + "--load-deepspeed", + action="store_true", + help="Load DeepSpeed checkpoint", + ) + group.add_argument( + "--ws-encoding-start-id", + type=int, + default=None, + help="Start id for whitespace encoding", + ) + group.add_argument( + "--ws-encoding-length", + type=int, + default=None, + help="Length of whitespace encoding", + ) + group.add_argument( + "--dataset", + type=str, + default="humaneval", + ) + group.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + group.add_argument( + "--output-prefix", + type=str, + default="./output/humaneval", + help="Prefix for output files", + ) + group.add_argument( + "--gen-node-rank", + type=int, + default=None, + ) + group.add_argument( + "--gen-node-world-size", + type=int, + default=None, + ) + group.add_argument( + "--gen-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + group.add_argument( + "--gen-rank", + type=int, + default=0, + help="Machine rank for human eval generation", + ) + group.add_argument( + "--extra-prompt", + type=str, + default=None, + help="Extra prompt to use for human eval generation", + ) + group.add_argument( + "--verbose-interval", + type=int, + default=100, + ) + group.add_argument( + "--problem-split", + type=str, + default="test", + ) + group.add_argument( + "--prompt-type", + type=str, + default="notag", + ) + group.add_argument( + "--num-devices-per-node", + type=int, + default=None, + ) + group.add_argument( + "--return-scores", + action="store_true", + ) + group.add_argument( + "--free-guidance", + action="store_true", + ) + group.add_argument( + "--guide-temp", + type=float, + default=1.5, + ) + group.add_argument( + "--attention-upweight", + type=float, + default=None, + ) + group.add_argument( + '--bad-ids', + nargs="*", + type=int, + default=None, + help='Identify the type of programming language to generate', + ) + group.add_argument( + "--src-path", + type=str, + default="", + help="Get source path", + ) + group.add_argument( + "--tgt-path", + type=str, + default="", + help="Get target path", + ) + group.add_argument( + '--language-src-type', + type=str, + default=None, + help='Identify the type of programming language', + ) + group.add_argument( + '--language-tgt-type', + type=str, + default=None, + help='Identify the type of programming language to translate', + ) + + return parser + + +def main(node_rank: int, local_rank: int, master_port: int, num_devices: int): + """Main program.""" + os.environ["WORLD_SIZE"] = str(num_devices) + os.environ["RANK"] = str(local_rank) + os.environ["MASTER_ADDR"] = "0.0.0.0" + os.environ["MASTER_PORT"] = f"{master_port}" + + initialize_megatron( + extra_args_provider=add_code_generate_args, + args_defaults={ + "tokenizer_type": "GPT2BPETokenizer", + "no_load_rng" : True, + "no_load_optim" : True, + }, + ) + + # set_random_seed(node_rank * num_devices + local_rank) + args = get_args() + if args.num_layers_per_virtual_pipeline_stage is not None: + print("Interleaved pipeline schedule is not yet supported for text generation.") + exit() + + world_size = args.gen_node_world_size * num_devices + args.gen_rank = num_devices * node_rank + local_rank + args.gen_world_size = world_size + print(f"Generating on rank {args.gen_rank} of {args.gen_world_size}") + + # Set up model and load checkpoint. + state_dict = torch.load(args.load, map_location="cpu") + state_dict = state_dict["module"] + + print("Building CodeGeeX model ...") + model = model_provider() + model.load_state_dict(state_dict) + model.eval() + if args.fp16 and args.ln_fp16: + model.half() + model.cuda() + + # Generate samples. + run_generation_distributed(model) + + print(f"(gen_rank={args.gen_rank}, rank={local_rank}) finished, waiting ...") + torch.distributed.barrier() + + +def server(): + print(f"[ server ] starting ...", flush=True) + parser = argparse.ArgumentParser() + parser.add_argument( + "--channel-ip", + type=str, + default=None, + help="IP for ZeroMQ channel", + ) + parser.add_argument( + "--channel-port", + type=int, + default=5555, + help="Port for ZeroMQ channel", + ) + parser.add_argument( + "--master-port", + type=int, + default=6666, + help="Port for torch distributed", + ) + parser.add_argument( + "--samples-per-problem", + type=int, + default=200, + help="Number of samples to generate for each problem", + ) + parser.add_argument( + "--gen-node-world-size", + type=int, + default=1, + help="Number of machines to use for generation", + ) + parser.add_argument( + "--src-path", + type=str, + default="", + help="Get source path", + ) + parser.add_argument( + "--tgt-path", + type=str, + default="", + help="Get target path", + ) + parser.add_argument( + "--problem-split", + type=str, + default="test", + ) + parser.add_argument( + "--micro-batch-size", + type=int, + default=1, + ) + parser.add_argument( + '--language-src-type', + type=str, + default=None, + help='Identify the type of programming language', + ) + parser.add_argument( + '--language-tgt-type', + type=str, + default=None, + help='Identify the type of programming language to translate', + ) + + args = parser.parse_known_args()[0] + + entries = read_translation_dataset(args.src_path, + args.tgt_path, + lang_src=args.language_src_type, + lang_tgt=args.language_tgt_type, + dataset_type="humaneval") + + assert args.samples_per_problem % args.micro_batch_size == 0, "samples_per_problem should be divisible by micro_batch_size" + + res = [] + for entry in entries.values(): + res.extend([entry] * (args.samples_per_problem // args.micro_batch_size)) + random.shuffle(res) + all_entries = res + + # setup zeromq channel + print(f"[ server ] starting up on port {args.channel_port}", flush=True) + context = zmq.Context() + print(f"[ server ] creating socket", flush=True) + socket = context.socket(zmq.REP) + print(f"[ server ] binding to port {args.channel_port}", flush=True) + socket.bind(f"tcp://*:{args.channel_port}") + + print( + f"[ server ] loaded {len(entries)} entries, generated {len(all_entries)} samples", + flush=True, + ) + + remaining_entries = all_entries.copy() + running_workers = args.gen_node_world_size * torch.cuda.device_count() + num_finished = 0 + + print(f"[ server ] listening for requests ...", flush=True) + start_time = time.perf_counter() + while True: + # Wait for next request from client + msg = socket.recv_json() + rank = msg["rank"] + action = msg["action"] + + if action == "pull": + if len(remaining_entries) == 0: + print(f"[ server ] Shutting down worker {rank}", flush=True) + socket.send_json({"task_id": None}) + running_workers -= 1 + if running_workers == 0: + print(f"[ server ] All workers finished", flush=True) + break + else: + entry = remaining_entries.pop() + time_elapsed = time.perf_counter() - start_time + print(f"[ server ] Sending entry {entry['task_id']} to worker {rank}", flush=True) + remaining = ( + len(remaining_entries) + / (len(all_entries) - len(remaining_entries)) + * time_elapsed + ) + time_per_sampple = 0.0 if num_finished == 0 else time_elapsed / num_finished / args.micro_batch_size + print( + f"[ server ] total {len(all_entries)}, assigned {len(all_entries) - len(remaining_entries)}, " + f"finished {num_finished}, " + f"elapsed {time_elapsed:.4f}", + f"speed {time_per_sampple:.4f}s/sample", + f"remaining {remaining:.4f}", + flush=True, + ) + socket.send_json({"task_id": entry}) + else: + if action == "success": + print(f"[ server ] {msg['task_id']} is finished", flush=True) + socket.send_json({"pong": 1}) + else: + print(f"[ server ] {msg['task_id']} is not finished", flush=True) + remaining_entries.append(msg['task_id']) + socket.send_json({"pong": 1}) + break + + num_finished += 1 + + +if __name__ == "__main__": + torch.multiprocessing.set_start_method("spawn") + parser = argparse.ArgumentParser() + parser.add_argument( + "--hostfile", + type=str, + default="./hostfile", + ) + parser.add_argument( + "--master-port", + type=int, + default=5666, + ) + parser.add_argument( + "--seed", + type=int, + default=42, + ) + args = parser.parse_known_args()[0] + + print("start method: " + torch.multiprocessing.get_start_method()) + + processes = [] + num_devices = torch.cuda.device_count() + + hosts = open(args.hostfile, "r").readlines() + hosts = [host.strip() for host in hosts] + master_port = args.master_port + + node_rank = None + for i in range(len(hosts)): + if hosts[i] == socket.gethostbyname(socket.gethostname()): + node_rank = i + break + assert ( + node_rank is not None + ), f"Could not find hostname ({socket.gethostbyname(socket.gethostname())}) in hostlist" + + # launch server + if socket.gethostbyname(socket.gethostname()) == hosts[0]: + server_process = torch.multiprocessing.Process(target=server) + print(f"Launching server ...") + server_process.start() + processes.append(server_process) + + for i in range(num_devices): + local_rank = i + print(f"launching local rank {i}") + + p = torch.multiprocessing.Process( + target=main, + args=(node_rank, local_rank, master_port, num_devices), + ) + p.start() + processes.append(p) + + for p in processes: + p.join() diff --git a/src/utils.py b/src/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..998cefac490c556fbc6210ef3d7b7ac9899671d5 --- /dev/null +++ b/src/utils.py @@ -0,0 +1,187 @@ +import os +from typing import * +from data_utils import stream_jsonl, LANGUAGE_TAG + + +IMPORT_HELPER = { + "python": [ + "import math", + "import re", + "import sys", + "import copy", + "import datetime", + "import itertools", + "import collections", + "import heapq", + "import statistics", + "import functools", + "import hashlib", + "import numpy", + "import numpy as np", + "import string", + "from typing import *", + "from collections import *", + ], + "go" : [ + "math", + "strings", + "fmt", + "strconv", + "time", + "bytes", + "regexp", + "sort", + "math/rand", + "crypto/md5", + ], + "cpp" : [ + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + "#include", + ], +} + + +def read_dataset( + data_file: str = None, + dataset_type: str = "humaneval", + language_type: str = "python", + num_shot=None, +) -> Dict: + if num_shot is not None: + print(f"{num_shot}-shot setting...") + if "humaneval" in dataset_type.lower(): + data_file = os.path.join(data_file, language_type, "data", f"humaneval_{language_type}.jsonl.gz") + dataset = {task["task_id"]: task for task in stream_jsonl(data_file)} + else: + raise f"Dataset: {dataset_type} not supported." + + return dataset + + +def read_translation_dataset( + data_file_src: str = None, + data_file_tgt: str = None, + lang_src: str = None, + lang_tgt: str = None, + dataset_type: str = "humaneval", +) -> Dict: + if "humaneval" in dataset_type.lower(): + dataset_src = {task["task_id"]: task for task in stream_jsonl(data_file_src)} + dataset_tgt = {task["task_id"].split("/")[-1]: task for task in stream_jsonl(data_file_tgt)} + for k, sample in dataset_src.items(): + prompt = "code translation\n" + if lang_src == "cpp": + prompt += "C++:\n" + elif lang_src == "js": + prompt += "JavaScript:\n" + else: + prompt += f"{lang_src}:\n".capitalize() + prompt += dataset_src[k]["declaration"] + "\n" + dataset_src[k]["canonical_solution"].rstrip() + "\n" + if lang_tgt == "cpp": + prompt += "C++:\n" + elif lang_tgt == "js": + prompt += "JavaScript:\n" + else: + prompt += f"{lang_tgt}:\n".capitalize() + prompt += dataset_tgt[k.split("/")[-1]]["declaration"] + dataset_src[k]["prompt"] = prompt + else: + raise f"Dataset: {dataset_type} not supported." + + return dataset_src + + +def process_extra_prompt(prompt: str, language_type: str = None) -> str: + """ + Processes the extra prompt. + """ + language = language_type.lower() + if language in LANGUAGE_TAG: + extra_prompt = LANGUAGE_TAG[language] + "\n" + else: + extra_prompt = "" + + return extra_prompt + prompt + + +def is_code_generation_finished( + code: str, + language_type: str = None, + dataset: str = None, +): + """ + Checks whether the generated code is finished. + """ + if language_type is None or dataset is None: + return False + + if "humaneval" in dataset.lower(): + if language_type.lower() == "python": + for line in code.split("\n"): + if len(line.strip()) > 0 and line[0] != ' ' and line[0] != '\t': + return True + end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint"] + for w in end_words: + if w in code: + return True + elif language_type.lower() == "java": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "go": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "js": + if code.count("{") + 1 == code.count("}"): + return True + elif language_type.lower() == "cpp": + if code.count("{") + 1 == code.count("}"): + return True + + return False + + +def cleanup_code( + code: str, + language_type: str = None +): + """ + Cleans up the generated code. + """ + if language_type is None: + return code + + if language_type.lower() == "python": + end_words = ["\ndef", "\nclass", "\nif", "\n#", "\nprint", "\nassert"] + for w in end_words: + if w in code: + code = code[:code.rfind(w)] + elif language_type.lower() == "java": + main_pos = code.find("public static void main") + if main_pos != -1: + code = code[:main_pos] + '}' + if '}' in code: + code = code[:code.rfind('}')] + '}' + if code.count('{') + 1 == code.count('}'): + code += "\n}" + elif language_type.lower() == "go": + end_words = ["\n//", "\nfunc main("] + for w in end_words: + if w in code: + code = code[:code.rfind(w)] + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "cpp": + if '}' in code: + code = code[:code.rfind('}')] + '}' + elif language_type.lower() == "js": + if '}' in code: + code = code[:code.rfind('}')] + '}' + + return code