CodePLAN/generate_code_with_plan.py at main · sssszh/CodePLAN · 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
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146

import json
import os
import pprint
import torch
import pdb
import glob
from tqdm import tqdm
import pickle as pkl
import numpy as np
from collections import Counter
from transformers import RobertaTokenizer, T5ForConditionalGeneration
import Datasets_codeT5.utils as dsutils

os.environ["CUDA_VISIBLE_DEVICES"] = '0'

def generate_prompt(args, test_case_path, prompt_path, pro_path, plan_str, solutions_path, tokenizer,
                    starter_path=None):

    if args.is_plan:
        _input = "[GEN_PLAN]"
    else:
        _input = "[GEN_CODE]"
    _input += "\nQUESTION:\n"
    with open(pro_path, "r") as f:
        data = f.readlines()
        data = "".join(data)
    _input += data

    if starter_path != None:
        with open(starter_path, "r") as f:
            data = f.readlines()
            data = "".join(data)
            data = "\n" + data
        _input += data

    if os.path.exists(test_case_path):
        with open(test_case_path, "r") as f:
            data = json.load(f)
        if not data.get("fn_name"):
            _input += "\nUse Standard Input format"
        else:
            _input += "\nUse Call-Based format"
    elif starter_path is not None and os.path.exists(starter_path):
        _input += "\nUse Call-Based format"
    else:
        _input += "\nUse Standard Input format"
    _input += "\nLet's think step by step:\n"
    _input += plan_str
    _input += "\nANSWER:\n"

    return _input

def main(args):

    argsdict = vars(args)
    print(pprint.pformat(argsdict))

    original_problems = glob.glob(args.test_path + '/*')
    problems = sorted(original_problems)

    if not os.path.exists(args.output_path):
        os.makedirs(args.output_path, exist_ok=True)
    print("Saving results to {}".format(args.output_path))

    if args.start > len(problems) or args.start < 0:
        print(f"start index {args.start} > number of problems {len(problems)}")
        return
    start = args.start
    if args.end is None or args.end > len(problems):
        end = len(problems)
    else:
        end = args.end
    problems = problems[start:end]

    # Set up model
    tokenizer = RobertaTokenizer.from_pretrained('Salesforce/codet5-large')
    print("Loading model from {}...".format(args.model_path))
    model = T5ForConditionalGeneration.from_pretrained(args.model_path)
    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    model.to(device)
    sols = []

    for index, problem in tqdm(enumerate(problems), ncols=0, total=len(problems)):

        prob_path = os.path.join(problem)
        print(f"problem path = {prob_path}")

        problem_id = int(problem.split('/')[-1])
        with open(os.path.join("outputs\plans", f"{problem_id}.json"), 'r') as f:
            plans_list = json.load(f)[f"{problem_id}"]["codes"]

        for idx in tqdm(range(20), ncols=0, total=20):
            test_case_path = os.path.join(prob_path, "input_output.json")
            prompt_path = os.path.join(prob_path, f"{problem_id}_{idx}.txt")
            pro_path = os.path.join(prob_path, "question.txt")
            starter_path = os.path.join(prob_path, "starter_code.py")

            solutions_path = os.path.join(prob_path, "solutions.json")
            plan_str = plans_list[idx]
            if os.path.exists(os.path.join(args.output_path, f"{problem_id}_{idx}.json")):
                continue

            if not os.path.exists(starter_path):
                starter_path = None
            try:
                input_text = generate_prompt(args, test_case_path, prompt_path, pro_path, plan_str, solutions_path,
                                            tokenizer, starter_path)
            except:
                continue

            with torch.no_grad():

                input_ids = torch.LongTensor(tokenizer.encode(input_text,
                                                                verbose=False,
                                                                max_length=args.source_len)).unsqueeze(0).cuda()

                num_loops = int(args.num_seqs / args.num_seqs_per_iter)
                output_programs = []
                for i in tqdm(range(num_loops), ncols=0, total=num_loops, leave=False):
                    output_ids = model.generate(
                            input_ids,
                            do_sample=True,
                            temperature=args.temperature,
                            max_length=args.max_len,
                            num_return_sequences=args.num_seqs_per_iter,
                            top_p=0.95)

                    for output_id in output_ids:
                            output_programs.append(tokenizer.decode(output_id, skip_special_tokens=True))

                save_codes = {
                    f"{problem_id}": {'plans': output_programs}
                }
                save_loc = os.path.join(args.output_path, f"{problem_id}_{idx}.json")
                with open(save_loc, 'w') as f:
                    json.dump(save_codes, f)


if __name__ == "__main__":

    from configs.generate_code_with_plan_configs import *

    main(args)