Merge pull request #7 from YanxinLu/main

add background prompt template

eval/data/multistep_template.txt

@@ -0,0 +1,23 @@
+PROBLEM DESCRIPTION:
+You will be provided with problem steps along with background knowledge necessary for solving the problem. Your task will be to develop a Python solution focused on the next step of the problem-solving process.
+
+PROBLEM STEPS AND FUNCTION CODE:
+Here, you'll find the Python code for the initial steps of the problem-solving process. This code is integral to building the solution.
+
+{problem_steps_str}
+
+NEXT STEP - PROBLEM STEP AND FUNCTION HEADER:
+This part will describe the next step in the problem-solving process. A function header will be provided, and your task is to develop the Python code for this next step based on the provided description and function header.
+
+{next_step_str}
+
+DEPENDENCIES:
+Use only the following dependencies in your solution. Do not include these dependencies at the beginning of your code.
+
+{dependencies}
+
+RESPONSE GUIDELINES:
+Now, based on the instructions and information provided above, write the complete and executable Python program for the next step in a single block.
+Your response should focus exclusively on implementing the solution for the next step, adhering closely to the specified function header and the context provided by the initial steps.
+Your response should NOT include the dependencies and functions of all previous steps. If your next step function calls functions from previous steps, please make sure it uses the headers provided without modification.
+DO NOT generate EXAMPLE USAGE OR TEST CODE in your response. Please make sure your response python code in format of ```python```.

eval/scripts/gencode_json.py

@@ -10,7 +10,7 @@
 
 
 DEFAULT_PROMPT_TEMPLATE = Path("eval", "data", "background_comment_template.txt").read_text()
-
+BACKGOUND_PROMPT_TEMPLATE = Path("eval", "data", "multistep_template.txt").read_text()
 
 class Gencode:
     def __init__(self, model: str, output_dir: Path,
@@ -132,7 +132,6 @@ def process_problem_steps(self, problem_data: dict, num_steps: int):
     def generate_prompt_with_steps(self, prob_data: dict, num_steps: int,
                                    prompt_template=DEFAULT_PROMPT_TEMPLATE):
         # parse the input file and extract the content
-
         problem_steps_str, next_step_str, previous_code_str = self.process_problem_steps(prob_data,
                                                                                          num_steps)
         dependencies = prob_data["required_dependencies"]
@@ -194,6 +193,7 @@ def main(model: str,
         model=model, output_dir=output_dir,
         prompt_dir=prompt_dir,  with_background=with_background, temperature=temperature
     )
+    prompt_template = BACKGOUND_PROMPT_TEMPLATE if with_background else DEFAULT_PROMPT_TEMPLATE
     data = read_from_jsonl(input_path)
     for problem in data:
         prob_id = problem['problem_id']
@@ -203,7 +203,7 @@ def main(model: str,
             if (prob_id == "13" and i == 5) or (prob_id == "62" and i == 0)\
                     or (prob_id == "76" and i == 2):
                 continue
-            gcode.generate_response_with_steps(problem, i + 1, steps, model)
+            gcode.generate_response_with_steps(problem, i + 1, steps, model, prompt_template)
 
 
 if __name__ == "__main__":