Spaces:
Sleeping
Sleeping
File size: 2,121 Bytes
2fd8593 | 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 | import json
import re
import os
import argparse
import shutil
from utils import extract_planning, content_to_json, format_json_data
parser = argparse.ArgumentParser()
parser.add_argument('--paper_name',type=str)
parser.add_argument('--output_dir',type=str, default="")
args = parser.parse_args()
output_dir = args.output_dir
with open(f'{output_dir}/planning_trajectories.json', encoding='utf8') as f:
traj = json.load(f)
yaml_raw_content = ""
for turn_idx, turn in enumerate(traj):
if turn_idx == 8:
yaml_raw_content = turn['content']
if "</think>" in yaml_raw_content:
yaml_raw_content = yaml_raw_content.split("</think>")[-1]
match = re.search(r"```yaml\n(.*?)\n```", yaml_raw_content, re.DOTALL)
if match:
yaml_content = match.group(1)
with open(f'{output_dir}/planning_config.yaml', 'w', encoding='utf8') as f:
f.write(yaml_content)
else:
# print("No YAML content found.")
match2 = re.search(r"```yaml\\n(.*?)\\n```", yaml_raw_content, re.DOTALL)
if match2:
yaml_content = match2.group(1)
with open(f'{output_dir}/planning_config.yaml', 'w', encoding='utf8') as f:
f.write(yaml_content)
else:
print("No YAML content found.")
# ---------------------------------------
artifact_output_dir=f"{output_dir}/planning_artifacts"
os.makedirs(artifact_output_dir, exist_ok=True)
context_lst = extract_planning(f'{output_dir}/planning_trajectories.json')
arch_design = content_to_json(context_lst[1])
logic_design = content_to_json(context_lst[2])
formatted_arch_design = format_json_data(arch_design)
formatted_logic_design = format_json_data(logic_design)
with open(f"{artifact_output_dir}/1.1_overall_plan.txt", "w", encoding="utf-8") as f:
f.write(context_lst[0])
with open(f"{artifact_output_dir}/1.2_arch_design.txt", "w", encoding="utf-8") as f:
f.write(formatted_arch_design)
with open(f"{artifact_output_dir}/1.3_logic_design.txt", "w", encoding="utf-8") as f:
f.write(formatted_logic_design)
shutil.copy(f"{output_dir}/planning_config.yaml", f"{artifact_output_dir}/1.4_config.yaml")
|