add instruction model eval script

Evaluation/HumanEval/eval_instruct.py

@@ -0,0 +1,126 @@
+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
+from transformers import AutoTokenizer, AutoModelForCausalLM
+from human_eval.evaluation import evaluate_functional_correctness
+
+def build_deepseekcoder_instruction(languge: str, question: str):
+    return '''
+Please help me to complete the function. Use the given packages only and DO NOT refer any new package. Please return all completed function in a codeblock. 
+Here is the given code to do completion:
+```{}
+{}
+```
+'''.strip().format(languge.lower(), question)
+
+
+def generate_one(example, lang, tokenizer, model):
+    prompt = build_deepseekcoder_instruction(lang, example['prompt'])
+    inputs = tokenizer.apply_chat_template(
+        [{'role': 'user', 'content': prompt }],
+        return_tensors="pt"
+    ).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=512,
+        do_sample=False, 
+        top_p=0.95,
+        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)
+
+    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="cuda"
+    )
+    
+    model.eval()
+    problem_file = os.path.join(data_abs_dir, f"humaneval-{lang}.jsonl")
+    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, lang, 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