evaluation/general_benchmarks/HumanEval/eval_instruct_vllm.py

import argparse
import json
import os
import shutil
from pathlib import Path

import torch
import transformers
from human_eval.evaluation import evaluate_functional_correctness
from tqdm import tqdm
from transformers import AutoTokenizer
from utils.utils import extract_generation_code, languge_settings
from vllm import LLM, SamplingParams

data_abs_dir = Path(__file__).parent / "data"


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 create_dir(output_dir):
    if os.path.exists(output_dir):
        if not os.access(output_dir, os.W_OK):
            shutil.rmtree(output_dir)
            os.makedirs(output_dir)
            os.chmod(output_dir, 0o777)
            print("not write permission, makedir:", output_dir)
        else:
            print(f"{output_dir} exists!")
    else:
        os.makedirs(output_dir)
        os.chmod(output_dir, 0o777)
        print("makedir:", output_dir)


def get_client_res(messages, example, output_key, open_ai_key=False):
    try:
        if open_ai_key:
            from openai import AzureOpenAI, OpenAI
            try:
                api_key = os.environ["OPENAI_API_KEY"]
            except KeyError:
                print("环境变量 OPENAI_API_KEY 未设置")
                api_key = "default_value"

            client = AzureOpenAI(
                api_key=api_key,
                api_version="2024-07-01-preview",
                azure_endpoint="https://zju-tablegpt.openai.azure.com/",
            )
            chat_response = client.chat.completions.create(
                model="gpt-4o",
                # model="gpt-4o-mini",
                messages=messages,
                top_p=0.95,
                temperature=0,
                max_tokens=1024,
                timeout=40,
            )
        else:
            # Set OpenAI's API key and API base to use vLLM's API server.
            openai_api_key = "EMPTY"
            openai_api_base = "http://localhost:8080/v1"

            client = OpenAI(
                api_key=openai_api_key,
                base_url=openai_api_base,
            )
            chat_response = client.chat.completions.create(
                model="qwen2-7b-sft",
                messages=messages,
                top_p=0.3,
                temperature=0.1,
                max_tokens=1024,
            )
        example[output_key] = chat_response.choices[0].message.content
    except Exception as e:
        print(f"An unexpected error occurred: {e}")
        example[output_key] = None
    example["input"] = messages
    return example



def generate_main(args):
    model_name_or_path = args.model_path
    lang = args.language
    temp_dir = args.temp_dir
    create_dir(temp_dir)
    # os.makedirs(temp_dir, exist_ok=True)
    problem_file = os.path.join(data_abs_dir, f"humaneval-{lang}.jsonl")
    if not args.api:
        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
            )
        )
        llm_args = {
            "model": model_name_or_path,
            "gpu_memory_utilization": 0.95,
            "trust_remote_code": True,
            "tensor_parallel_size": args.gpus_num,
            "dtype": "half",
            "max_model_len": 8192,
            "enforce_eager": True,
        }

        llm = LLM(**llm_args)
        sampling_params = SamplingParams(
            temperature=0,
            max_tokens=1024,
            top_p=0.95,
            stop_token_ids=[tokenizer.eos_token_id],
        )

    examples = [json.loads(x) for x in open(problem_file) if x.strip()]
    print("Read {} examples for evaluation over.".format(len(examples)))
    messages_list = []
    for example in tqdm(examples, desc="Generating"):
        prompt = build_deepseekcoder_instruction(
            languge_settings[lang]["full_name"], example["prompt"]
        )
        message = [{"role": "user", "content": prompt}]
        if args.api:
            messages_list.append(message)
        else:
            messages_list.append(
                tokenizer.apply_chat_template(
                    message, tokenize=False, add_generation_prompt=True
                )
            )
    if args.api:
        from joblib import Parallel, delayed
        examples_ = Parallel(n_jobs=24)(
            delayed(get_client_res)(inp, examples[i], "output",open_ai_key=True)
            for i, inp in enumerate(tqdm(messages_list))
        )

        # 请求错误的重新请求
        examples = []
        for example in examples_:
            if example["output"] == None:
                example = get_client_res(
                    example["input"], example, "output", open_ai_key=True
                )
            del example["input"]
            examples.append(example)

        generated_examples = []
        for example in examples:
            example = extract_generation_code(example, lang_code=lang)
            generated_examples.append(example)
    else:
        outputs = llm.generate(messages_list, sampling_params=sampling_params)
        generated_examples = []
        for i, output in enumerate(tqdm(outputs)):
            output = output.outputs[0].text
            example = examples[i]
            example["output"] = output
            example = extract_generation_code(example, lang_code=lang)
            generated_examples.append(example)

    print("Generate all over!!!")
    # os.makedirs(args.save_dir, exist_ok=True)
    create_dir(args.save_dir)
    saved_path = os.path.join(args.save_dir, "results_humaneval.json")
    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,
        out_path=saved_path,
    )
    print(lang, result, model_name_or_path)


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--model_path",
        type=str,
        help="model name or path",
        default="/data4/sft_output/qwen2-instruct-0709/checkpoint-1400",
    )
    parser.add_argument(
        "--gpus_num", type=int, default=1, help="the number of GPUs you want to use."
    )
    parser.add_argument(
        "--save_dir",
        type=str,
        help="output path of your generation",
        default="output",
    )
    parser.add_argument("--api", action="store_true", help="infer api type")
    parser.add_argument("--language", type=str, help="langauge", default="python")
    parser.add_argument(
        "--temp_dir", type=str, help="temp dir for evaluation", default="output/tmp"
    )
    parser.add_argument("--seed", type=int, help="seed", default=42)
    args = parser.parse_args()

    os.environ["TOKENIZERS_PARALLELISM"] = "false"
    transformers.set_seed(args.seed)
    generate_main(args)