Function main

evaluation/evalplus/evalplus/perf/sas.py:101–162

(
    output: str,  # output file
    n: int = 16,  # sample size and batch size
    model: Optional[str] = "TheBloke/deepseek-coder-33B-instruct-AWQ",
    port: str = 8088,
)

def main(
    output: str,  # output file
    n: int = 16,  # sample size and batch size
    model: Optional[str] = "TheBloke/deepseek-coder-33B-instruct-AWQ",
    port: str = 8088,
):
    assert output.endswith(".jsonl"), "output must be a .jsonl file"

    base_url = f"http://localhost:{port}/v1"
    print(f"Trying to query vLLM model: {model} at {base_url}")
    print(f"Note: To use SaS, you need to first set up a vLLM server for {model}")
    print(f"For example:")
    print(
        f"""python -m vllm.entrypoints.openai.api_server \\
--model "{model}" \\
--port {port} \\
--tensor-parallel-size 2 \\
--max-num-seqs 16 \\
--gpu-memory-utilization 1.0"""
    )

    # "task_id" -> { "task_id", "entry_point", "ref_code", }
    tasks = {}
    for task_id, item in get_human_eval_plus().items():
        tasks[task_id] = {
            "task_id": task_id,
            "entry_point": item["entry_point"],
            "ref_code": item["prompt"] + item["canonical_solution"],
        }

    for task_id, item in get_mbpp_plus().items():
        tasks[task_id] = {
            "task_id": task_id,
            "entry_point": item["entry_point"],
            "ref_code": item["prompt"] + item["canonical_solution"],
        }

    # Using vLLM as a backend, please make sure that a vLLM server is available first.
    # vLLM document: https://docs.vllm.ai/en/latest/
    client = openai.OpenAI(api_key="none", base_url=base_url)

    with open(output, "w") as f:
        for task_id, item in tqdm(tasks.items(), total=len(tasks)):
            responses = fewshot_cot(
                task_id=task_id,
                client=client,
                entry_point=item["entry_point"],
                code=item["ref_code"],
                model=model,
                n=n,
            )
            f.write(
                json.dumps(
                    {
                        "task_id": task_id,
                        "ref_code": item["ref_code"],
                        "synthesizers": responses,
                    }