(args, actor_config)
| 108 | return total_input_token, total_output_token |
| 109 | |
| 110 | def stylize_poster(args, actor_config): |
| 111 | total_input_token, total_output_token = 0, 0 |
| 112 | poster_content = json.load(open(f'contents/{args.model_name}_{args.poster_name}_poster_content_{args.index}.json', 'r')) |
| 113 | agent_name = 'style_agent' |
| 114 | |
| 115 | with open(f"prompt_templates/{agent_name}.yaml", "r") as f: |
| 116 | style_config = yaml.safe_load(f) |
| 117 | |
| 118 | actor_model = ModelFactory.create( |
| 119 | model_platform=actor_config['model_platform'], |
| 120 | model_type=actor_config['model_type'], |
| 121 | model_config_dict=actor_config['model_config'], |
| 122 | ) |
| 123 | |
| 124 | actor_sys_msg = style_config['system_prompt'] |
| 125 | |
| 126 | actor_agent = ChatAgent( |
| 127 | system_message=actor_sys_msg, |
| 128 | model=actor_model, |
| 129 | message_window_size=10, |
| 130 | ) |
| 131 | |
| 132 | ckpt = pkl.load(open(f'checkpoints/{args.model_name}_{args.poster_name}_content_ckpt_{args.index}.pkl', 'rb')) |
| 133 | content_logs = ckpt['content_logs'] |
| 134 | outline = ckpt['outline'] |
| 135 | |
| 136 | sections = list(outline.keys()) |
| 137 | sections = [s for s in sections if s != 'meta'] |
| 138 | |
| 139 | jinja_env = Environment(undefined=StrictUndefined) |
| 140 | |
| 141 | template = jinja_env.from_string(style_config["template"]) |
| 142 | style_logs = {} |
| 143 | |
| 144 | for section_index in range(len(sections)): |
| 145 | section_name = sections[section_index] |
| 146 | section_outline = json.dumps(outline[section_name]) |
| 147 | section_code = content_logs[section_name][-1]['code'] |
| 148 | |
| 149 | print(f'Stylizing for {section_name}') |
| 150 | |
| 151 | img_ratio_json = get_img_ratio_in_section(poster_content[section_name]) |
| 152 | |
| 153 | jinja_args = { |
| 154 | 'content_json': poster_content[section_name], |
| 155 | 'function_docs': documentation, |
| 156 | 'existing_code': section_code, |
| 157 | 'image_ratio': img_ratio_json, |
| 158 | } |
| 159 | |
| 160 | prompt = template.render(**jinja_args) |
| 161 | if section_index == 0: |
| 162 | existing_code = '' |
| 163 | else: |
| 164 | existing_code = style_logs[sections[section_index - 1]][-1]['concatenated_code'] |
| 165 | style_logs[section_name] = stylize( |
| 166 | actor_agent, |
| 167 | prompt, |
no test coverage detected