generate_code.py

import argparse
from pathlib import Path
from LLMClient import *
from utils.dataset_loader import *
from utils.file_util import *
from utils.process_result import process_CoderEval, process_DevEval
from tqdm import tqdm
import multiprocessing as mp

def generate_code(dataset, language, model, req_weight, code_weight, reflection, expand, rule, ppl, mode):
    dataset = get_dataset_name(dataset, language)
    input_file = (
        PROMPT_DIR/dataset/
        f"{'reflection_' if reflection else ''}_sub_req_retrieve"/
        f"_req{req_weight}_code{code_weight}_expand{expand}"
        f"{'_rule' if rule else ''}"
        f"{'_ppl' if ppl else ''}.jsonl"
    )
    output_file = (
        OUT_CODE_DIR/dataset/
        f"{'reflection_' if reflection else ''}_sub_req_retrieve"/
        f"{model}_req{req_weight}_code{code_weight}_expand{expand}"
        f"{'_rule' if rule else ''}"
        f"{'_ppl' if ppl else ''}"
        f"_{mode}.jsonl"
    )    
    prompt = load_jsonl(input_file)
    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
    if model == "gpt":
        client = GPTClient()
    elif model == "deepseek":
        client = DeepSeekClient()
    elif model == "qwencoder":
        client = QwenCoderClient()
    elif model == "deepseekcoder":
        client = DeepSeekCoderClient()
    with open(output_file, 'a') as f:
        finished = len(load_jsonl(output_file))
        for i in tqdm(range(finished, len(prompt)), desc="Generating code"):
            p = prompt[i]
            code = client.generate(p["prompt"], mode)
            f.write(json.dumps({'namespace': p['namespace'], 'completion': code}) + '\n')
    
    # Use the same result processing method to reproduce the experimental results reported in the paper.
    # if dataset == "DevEval":
    #     process_DevEval(output_file)
    # elif dataset == "CoderEvalJava":
    #     process_CoderEval("java", output_file)
    # elif dataset == "CoderEvalPython":
    #     process_CoderEval("python", output_file)

def main():
    parser = argparse.ArgumentParser(description="Code generation runner")
    parser.add_argument("--divide", action="store_true", help="Whether to divide requirement")
    parser.add_argument("--reflection", action="store_true", help="Whether to enable reflection")
    parser.add_argument("--dataset", type=str, default="DevEval", help="Dataset name")
    parser.add_argument("--language", type=str, default="python", help="Dataset language")
    parser.add_argument("--model", type=str, default="gpt", help="Model name")
    parser.add_argument("--req_weight", type=float, default=0.5, help="Weight of requirement similarity")
    parser.add_argument("--code_weight", type=float, default=0.5, help="Weight of code similarity")
    parser.add_argument("--expand", type=int, help="Number of upstream functions to expand")
    parser.add_argument("--structure_rank", action="store_true", help="Whether to use rule-based ranking")
    parser.add_argument("--ppl_rank", action="store_true", help="Whether to use PPL-based ranking")
    parser.add_argument("--mode", type=str, default="greedy", choices=["greedy", "sample"], help="Generation mode")
    args = parser.parse_args()

    print("[INFO] Begin generate code")
    generate_code(
        args.dataset,
        args.language,
        args.model,
        args.req_weight,
        args.code_weight,
        args.reflection,
        args.expand,
        args.structure_rank,
        args.ppl_rank,
        args.mode,
    )

    print("All datasets processed.")


if __name__ == "__main__":
    main()