(args, actor_config, version=1)
| 44 | |
| 45 | @retry(stop=stop_after_attempt(5)) |
| 46 | def parse_raw(args, actor_config, version=1): |
| 47 | raw_source = args.poster_path |
| 48 | markdown_clean_pattern = re.compile(r"<!--[\s\S]*?-->") |
| 49 | |
| 50 | raw_result = doc_converter.convert(raw_source) |
| 51 | |
| 52 | raw_markdown = raw_result.document.export_to_markdown() |
| 53 | text_content = markdown_clean_pattern.sub("", raw_markdown) |
| 54 | |
| 55 | if len(text_content) < 500: |
| 56 | print('\nParsing with docling failed, using marker instead\n') |
| 57 | parser_model = create_model_dict(device='cuda', dtype=torch.float16) |
| 58 | text_content, rendered = parse_pdf(raw_source, model_lst=parser_model, save_file=False) |
| 59 | |
| 60 | if version == 1: |
| 61 | template = Template(open("utils/prompts/gen_poster_raw_content.txt").read()) |
| 62 | elif version == 2: |
| 63 | template = Template(open("utils/prompts/gen_poster_raw_content_v2.txt").read()) |
| 64 | |
| 65 | if args.model_name_t.startswith('vllm_qwen'): |
| 66 | actor_model = ModelFactory.create( |
| 67 | model_platform=actor_config['model_platform'], |
| 68 | model_type=actor_config['model_type'], |
| 69 | model_config_dict=actor_config['model_config'], |
| 70 | url=actor_config['url'], |
| 71 | ) |
| 72 | else: |
| 73 | actor_model = ModelFactory.create( |
| 74 | model_platform=actor_config['model_platform'], |
| 75 | model_type=actor_config['model_type'], |
| 76 | model_config_dict=actor_config['model_config'], |
| 77 | ) |
| 78 | |
| 79 | actor_sys_msg = 'You are the author of the paper, and you will create a poster for the paper.' |
| 80 | |
| 81 | actor_agent = ChatAgent( |
| 82 | system_message=actor_sys_msg, |
| 83 | model=actor_model, |
| 84 | message_window_size=10, |
| 85 | token_limit=actor_config.get('token_limit', None) |
| 86 | ) |
| 87 | |
| 88 | while True: |
| 89 | prompt = template.render( |
| 90 | markdown_document=text_content, |
| 91 | ) |
| 92 | actor_agent.reset() |
| 93 | response = actor_agent.step(prompt) |
| 94 | input_token, output_token = account_token(response) |
| 95 | |
| 96 | content_json = get_json_from_response(response.msgs[0].content) |
| 97 | |
| 98 | if len(content_json) > 0: |
| 99 | break |
| 100 | print('Error: Empty response, retrying...') |
| 101 | if args.model_name_t.startswith('vllm_qwen'): |
| 102 | text_content = text_content[:80000] |
| 103 |
no test coverage detected