(text:str=None, trajectory:list=None, last_only=False)
| 67 | |
| 68 | |
| 69 | def extract_program(text:str=None, trajectory:list=None, last_only=False) -> str: |
| 70 | assert text is not None or trajectory is not None, "Either text or trajectory should be provided." |
| 71 | if trajectory is None: |
| 72 | try: |
| 73 | trajectory = text_to_trajectory(text) |
| 74 | except: |
| 75 | return "raise ValueError('Invalid trajectory')" |
| 76 | |
| 77 | program_list = [] |
| 78 | import_lines = [] |
| 79 | for i, item in enumerate(trajectory): |
| 80 | if item["role"] == "program": |
| 81 | cur_program = item["content"] |
| 82 | if i < len(trajectory) - 1: |
| 83 | assert trajectory[i+1]["role"] == "output" |
| 84 | output = trajectory[i+1]["content"].strip() |
| 85 | if is_execution_success(output): |
| 86 | program_list.append(cur_program) |
| 87 | else: |
| 88 | # extract import lines only |
| 89 | for line in cur_program.split("\n"): |
| 90 | if line.startswith("import") or line.startswith("from"): |
| 91 | import_lines.append(line) |
| 92 | else: |
| 93 | program_list.append(cur_program) |
| 94 | # add import lines to the first program |
| 95 | if len(program_list) == 0: |
| 96 | program_list.append("") |
| 97 | if len(import_lines) > 0: |
| 98 | program_list[0] = "\n".join(import_lines) + "\n" + program_list[0] |
| 99 | for i, program in enumerate(program_list[:-1]): |
| 100 | program_list[i] = "\n".join([line for line in program.split("\n") if not line.strip().startswith("print(")]) |
| 101 | |
| 102 | if last_only: |
| 103 | program = program_list[-1] |
| 104 | else: |
| 105 | program = "\n".join(program_list) |
| 106 | return program |
| 107 | |
| 108 | |
| 109 | def extract_program_output(pred_str, last_only=True): |
no test coverage detected