AutoAC/cp_framework_usage.py at main · Sriram-Sai-Ganesh/AutoAC · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
import os
from argparse import ArgumentParser, Namespace
from datetime import datetime

import utils
from agents import Agent, CPAgent
from dataset import CPDataset
from framework import CPFramework, TestRunnerFramework
from prompts import CPPromptTemplate
from test_runner import TestRunner

# sample usage:
# python3 cp_framework_usage.py --dataset_path ./data/cses_20_train.json --prompt_template_path ./prompts/templates/cp_test_prompt --agent_name Qwen2.5-0.5B-Instruct --max_new_tokens 1000 --expt_name test_outputs


def parse_args() -> Namespace:
    parser = ArgumentParser()
    parser.add_argument("--dataset_path", type=str, default="./data/cses_20_train.json")
    parser.add_argument(
        "--prompt_template_path", type=str, default="./prompts/templates/cp_test_prompt"
    )
    parser.add_argument("--agent_name", type=str, default="Qwen/Qwen2.5-0.5B-Instruct")
    parser.add_argument("--max_new_tokens", type=int, default=1000)
    parser.add_argument("--temperature", type=float, default=0.1)
    parser.add_argument(
        "--expt_name", type=str, default=datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
    )
    parser.add_argument("--stress_test_behavior", action="store_true", help="run iterative stress test behavior? true if present.")
    parser.add_argument("--generator_prompt_path", type=str, default="./prompts/templates/cp_generator_prompt")
    parser.add_argument("--checker_prompt_path", type=str, default="./prompts/templates/cp_checker_prompt")
    parser.add_argument("--dmrl", action="store_true", help="run on DMRL machine? changes dirs.")
    parser.add_argument("--only_rerun_tests", action="store_true", help="only rerun tests on existing generated code?")
    return parser.parse_args()


def main():
    args = parse_args()
    data_filename = os.path.basename(args.dataset_path)
    data_dir = "/fs/scratch/PAS2912/rnr/cses_data/"
    model_dir = "/fs/scratch/PAS2912/local_models"

    if args.dmrl:
        # if using the DMRL machine:
        data_dir = "./data"
        model_dir = "/local/storage/models"

    dataset_path = os.path.join(data_dir, data_filename)

    prompt_template_path = args.prompt_template_path
    generator_prompt_path = args.generator_prompt_path
    checker_prompt_path = args.checker_prompt_path

    # Agent obj wraps LLM.
    # options for local models:

    utils.log(f"Got model ID: {utils.get_model_id(args.agent_name)}", level="status")
    utils.log(f"Run stress test?{args.stress_test_behavior}", level="status")
    agent_kwargs = {
        "id": utils.get_model_id(args.agent_name),
        "max_new_tokens": args.max_new_tokens,
        'model_dir': model_dir,
        'temperature': args.temperature,
    }

    agent = None
    Framework_Class = None

    if args.only_rerun_tests:
        utils.log('Running tests on existing generated code...', level='status')
        agent_kwargs['skip_loading'] = True
        Framework_Class = TestRunnerFramework
    else:
        utils.log('Instantiating a new model...', level='status')
        Framework_Class = CPFramework

    agent = CPAgent(**agent_kwargs)

    # testcase runner obj.
    test_runner = TestRunner()
    # define prompt based on local file
    prompt_template_obj = CPPromptTemplate(prompt_template_path)
    generator_prompt = CPPromptTemplate(generator_prompt_path)
    checker_prompt = CPPromptTemplate(checker_prompt_path)
    # create desired framework obj, using agent & prompt
    framework = Framework_Class(
                expt_name=args.expt_name,
                agent=agent,
                test_runner=test_runner,
                prompt=prompt_template_obj,
                stress_test_behavior=args.stress_test_behavior,
                generator_prompt=generator_prompt,
                checker_prompt=checker_prompt
        )
    # run framework on a .jsonl dataset
    dataset = CPDataset(dataset_path, jsonl=True)

    framework.run(dataset)


if __name__ == "__main__":
    main()