example_script.py

"""This is an example script to evaluate a code generation model on APPS, you can also use the APPS solutions as code generations
>>> python example_script.py --model_ckpt MODEL_NAME --num_tasks 10 --difficulty introductory --n_samples 1
>>> python example_script.py --use_solutions True --num_tasks 10 --difficulty introductory --n_samples 1"""

import json
import pprint
from tqdm import tqdm
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer, pipeline, set_seed
from tools.utils import compute_metrics

def generate_prompt(sample):
    starter_code = None if len(sample["starter_code"]) == 0 else sample["starter_code"] 
    try:
        input_outpout = json.loads(sample["input_output"])
        fn_name = None if not input_outpout.get("fn_name") else input_outpout["fn_name"] 
    except ValueError:
        fn_name = None 
    _input = "\nQUESTION:\n"
    _input += sample["question"]
    if starter_code:
        _input += starter_code
    if fn_name:
        _input += "\nUse Standard Input format"
    else:
        _input += "\nUse Call-Based format"
    
    _input += "\nANSWER:\n"
    return _input


def complete_code(pipe, prompt, num_completions=1, max_length=256, **gen_kwargs):
    """Complete prompt with text generation pipeline and return num_completions."""
    prompt = pipe.tokenizer.eos_token + prompt
    try:
        code_gens = pipe(prompt, num_return_sequences=num_completions, max_length=max_length, **gen_kwargs)
        return [code_gen["generated_text"][len(prompt):] for code_gen in code_gens]
    except IndexError:
        print("prompt is longer than the context size of the model, generation skipped")
        code_gens = ""
        return [""]


def make_generations(dataset, args, model, tokenizer):
    set_seed(args.seed)
    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, device=args.device_int)

    # Generation settings
    gen_kwargs = {
        "do_sample": args.do_sample,
        "temperature": args.temperature,
        "top_p": args.top_p,
        "top_k": args.top_k
    }

    # Generate completions for evaluation set
    n_tasks = args.num_tasks if args.num_tasks is not None else len(dataset)
    print(f"ntasks is {n_tasks}")
    generations = []
    for task in tqdm(range(n_tasks)):
        task_generations = []
        prompt = generate_prompt(dataset[task]).strip()
        task_generations.extend(complete_code(pipe, prompt, num_completions=args.n_samples, max_length=args.max_length, **gen_kwargs))
        generations.append([gen.replace(args.eos, "") for gen in task_generations])
    return generations


def main(args):
    DATA_PATH = "codeparrot/apps"
    argsdict = vars(args)
    print(pprint.pformat(argsdict))

    # setup
    print("Loading evaluation dataset...")
    dataset = load_dataset(DATA_PATH, split="test", difficulties=[args.difficulty])
    if args.use_solutions:
        print("Using data solutions as code generations")
        model = None
        tokenizer = None
        generations = []
        for index in range(args.num_tasks+1):
            try:
                sol = json.loads(dataset[index]["solutions"])
                generations.append(sol[:args.n_solutions])
            except ValueError:
                print(f"No solutions for task {index} or not enough to have {args.n_solutions} solutions")
                break
        
    else:
        print("Loading tokenizer and model...")
        tokenizer = AutoTokenizer.from_pretrained(args.tokenizer)
        model = AutoModelForCausalLM.from_pretrained(args.model_ckpt)
        generations = make_generations(dataset, args, model, tokenizer)
    
    metrics = compute_metrics(generations, level=args.difficulty, k_list=args.k_list, count_errors=args.count_errors, debug=args.debug)
    print(metrics)
    with open(args.output_file, "w") as fp:
        json.dump(metrics, fp)


if __name__ == "__main__":
    import argparse

    parser = argparse.ArgumentParser(description="Testing a Language Model on APPS Python Code dataset")
    #model and tokenizer arguments
    parser.add_argument("--model_ckpt", default="loubnabnl/apps-1.5B-model", type=str, help="path to model checkpoint.")
    parser.add_argument("--tokenizer", default="gpt2", type=str, help="tokenizer to use.")
    parser.add_argument("--eos", default="<|endoftext|>", type=str, help="end of sentence token.")
    # generation arguments
    parser.add_argument("--do_sample", default=True, type=bool, help="do sampling in generation")
    parser.add_argument("--temperature", default=0.2, type=float, help="temperature for sampling")
    parser.add_argument("--top_p", default=0.95, type=float, help="top p for sampling")
    parser.add_argument("--top_k", default=0, type=float, help="top k for sampling")
    parser.add_argument("--max_length", default=1024, type=int, help="max length of generated code")
    # evaluation arguments
    parser.add_argument("--difficulty", default="all", type=str, help="difficulty level to select in the dataset from:\
     'all', 'introductory', 'interview'  and 'competition' ")
    parser.add_argument("--num_tasks", default=6, type=int, help="number of tasks to evaluate")
    parser.add_argument("--use_solutions", default=False, type=bool, help="use solutions instead of generating new code")
    parser.add_argument("--n_samples", default=1, type=int, help="number of samples to generate")
    parser.add_argument("--n_solutions", default=1, type=int, help="number of solutions to use")
    parser.add_argument("--k_list", default=[1, 2, 3], type=list, help="list of k values to evaluate pass@k")
    parser.add_argument("--count_errors", default=False, type=bool, help="count compilation and runtime errors for single generations")
    # configuration
    parser.add_argument("--seed", default=0, type=int, help="generation seed")
    parser.add_argument("--device_int", default=-1, type=int, help="device on which code generation is run, if positive use GPU")
    parser.add_argument("--debug", default=False, type=bool, help="debug mode")
    # save
    parser.add_argument("--output_file", default="apps_metrics.json", type=str, help="output file to save the results")

    args = parser.parse_args()
    main(args)