|
|
import argparse |
|
|
import json |
|
|
import os |
|
|
import torch |
|
|
from pathlib import Path |
|
|
from tqdm import tqdm |
|
|
|
|
|
data_abs_dir = Path(__file__).parent / "data" |
|
|
|
|
|
from utils.utils import extract_generation_code, languge_settings |
|
|
from transformers import AutoTokenizer, AutoModelForCausalLM |
|
|
from human_eval.evaluation import evaluate_functional_correctness |
|
|
|
|
|
def build_deepseekcoder_instruction(languge: str, question: str): |
|
|
return ''' |
|
|
Please continue to complete the function. You are not allowed to modify the given code and do the completion only. Please return all completed function in a codeblock. Here is the given code to do completion: |
|
|
```{} |
|
|
{} |
|
|
``` |
|
|
'''.strip().format(languge.lower(), question.strip()) |
|
|
|
|
|
def generate_one(example, lang, tokenizer, model): |
|
|
prompt = build_deepseekcoder_instruction(languge_settings[lang]['full_name'], example['prompt']) |
|
|
inputs = tokenizer.apply_chat_template( |
|
|
[{'role': 'user', 'content': prompt }], |
|
|
return_tensors="pt", |
|
|
add_generation_prompt=True |
|
|
).to(model.device) |
|
|
|
|
|
stop_id = tokenizer.convert_tokens_to_ids("<|EOT|>") |
|
|
assert isinstance(stop_id, int), "Invalid tokenizer, EOT id not found" |
|
|
|
|
|
outputs = model.generate( |
|
|
inputs, |
|
|
max_new_tokens=1024, |
|
|
do_sample=False, |
|
|
|
|
|
|
|
|
pad_token_id=stop_id, |
|
|
eos_token_id=stop_id |
|
|
) |
|
|
|
|
|
output = tokenizer.decode(outputs[0][len(inputs[0]):], skip_special_tokens=True) |
|
|
example['output'] = output |
|
|
|
|
|
return extract_generation_code(example, lang_code=lang) |
|
|
|
|
|
def generate_main(args): |
|
|
model_name_or_path = args.model |
|
|
lang = args.language |
|
|
saved_path = args.output_path |
|
|
temp_dir = args.temp_dir |
|
|
os.makedirs(temp_dir, exist_ok=True) |
|
|
problem_file = os.path.join(data_abs_dir, f"humaneval-{lang}.jsonl") |
|
|
|
|
|
print("model", model_name_or_path) |
|
|
tokenizer = AutoTokenizer.from_pretrained(model_name_or_path) |
|
|
print("load tokenizer {} from {} over.".format(tokenizer.__class__, model_name_or_path)) |
|
|
model = AutoModelForCausalLM.from_pretrained( |
|
|
model_name_or_path, |
|
|
torch_dtype=torch.bfloat16, |
|
|
device_map="auto", |
|
|
|
|
|
) |
|
|
model.eval() |
|
|
examples = [json.loads(x) for x in open(problem_file) if x.strip()] |
|
|
print("Read {} examples for evaluation over.".format(len(examples))) |
|
|
|
|
|
generated_examples = [] |
|
|
for ex in tqdm(examples, desc='Generating'): |
|
|
gen_example = generate_one(ex, args.language, tokenizer, model) |
|
|
generated_examples.append(gen_example) |
|
|
|
|
|
print("Generate all over!!!") |
|
|
with open(saved_path, 'w', encoding='utf-8') as fw: |
|
|
for ex in generated_examples: |
|
|
fw.write(json.dumps(ex) + '\n') |
|
|
print("Save {} processed examples into {} over!".format(len(generated_examples), saved_path)) |
|
|
|
|
|
result = evaluate_functional_correctness( |
|
|
input_file=saved_path, |
|
|
tmp_dir=temp_dir, |
|
|
n_workers=8, |
|
|
timeout=3.0, |
|
|
problem_file=problem_file, |
|
|
language=lang |
|
|
) |
|
|
print(lang, result, model_name_or_path) |
|
|
pass |
|
|
|
|
|
def evaluation_only(args): |
|
|
lang = args.language |
|
|
temp_dir = args.temp_dir |
|
|
assert os.path.exists(args.output_path), "Not fond output file: {}".format(args.output_path) |
|
|
os.makedirs(temp_dir, exist_ok=True) |
|
|
|
|
|
output_name = os.path.basename(args.output_path) |
|
|
output_examples = [json.loads(x) for x in open(args.output_path) if x.strip()] |
|
|
|
|
|
processed_examples = [extract_generation_code(ex, lang) for ex in tqdm(output_examples, "Processing")] |
|
|
processed_path = os.path.join(temp_dir, output_name) |
|
|
with open(processed_path, 'w', encoding='utf-8') as fw: |
|
|
for ex in processed_examples: |
|
|
fw.write(json.dumps(ex) + '\n') |
|
|
print("Save {} processed examples into {} over!".format(len(processed_examples), processed_path)) |
|
|
|
|
|
problem_file = os.path.join(data_abs_dir, f"humaneval-{lang}.jsonl") |
|
|
from human_eval.evaluation import evaluate_functional_correctness |
|
|
result = evaluate_functional_correctness( |
|
|
input_file=processed_path, |
|
|
tmp_dir=temp_dir, |
|
|
n_workers=8, |
|
|
timeout=3.0, |
|
|
problem_file=problem_file, |
|
|
language=lang |
|
|
) |
|
|
print(lang, result) |
|
|
|
|
|
if __name__ == '__main__': |
|
|
parser = argparse.ArgumentParser() |
|
|
parser.add_argument('--model', type=str, help="model name or path") |
|
|
parser.add_argument('--output_path', type=str, help="output path of your generation") |
|
|
parser.add_argument('--language', type=str, help="langauge") |
|
|
parser.add_argument('--temp_dir', type=str, help="temp dir for evaluation", default="tmp") |
|
|
args = parser.parse_args() |
|
|
|
|
|
os.environ["TOKENIZERS_PARALLELISM"] = "false" |
|
|
generate_main(args) |
|
|
pass |
|
|
|
