(llm,file_path,skip_cross_page_merge=False,max_page_retries=0)
| 70 | return "\n\n".join(document_text_list) |
| 71 | |
| 72 | def parse(llm,file_path,skip_cross_page_merge=False,max_page_retries=0): |
| 73 | sampling_params = SamplingParams(temperature=0.0,max_tokens=8192) |
| 74 | if file_path.lower().endswith(".pdf"): |
| 75 | try: |
| 76 | reader = PdfReader(file_path) |
| 77 | num_pages = reader.get_num_pages() |
| 78 | except: |
| 79 | return None |
| 80 | else: |
| 81 | num_pages = 1 |
| 82 | |
| 83 | try: |
| 84 | # Stage 1: Page to Markdown |
| 85 | page_to_markdown_query_list = [build_page_to_markdown_query(file_path,page_num) for page_num in range(1, num_pages + 1)] |
| 86 | responses = llm.generate(page_to_markdown_query_list, sampling_params=sampling_params) |
| 87 | results = [response.outputs[0].text for response in responses] |
| 88 | page_to_markdown_result = {} |
| 89 | retry_list = [] |
| 90 | for i,result in enumerate(results): |
| 91 | try: |
| 92 | json_data = json.loads(result) |
| 93 | page_response = PageResponse(**json_data) |
| 94 | natural_text = page_response.natural_text |
| 95 | markdown_element_list = [] |
| 96 | for text in natural_text.split('\n\n'): |
| 97 | if text.startswith("<Image>") and text.endswith("</Image>"): |
| 98 | pass |
| 99 | elif text.startswith("<table>") and text.endswith("</table>"): |
| 100 | try: |
| 101 | new_text = table_matrix2html(text) |
| 102 | except: |
| 103 | new_text = text.replace("<t>","").replace("<l>","").replace("<lt>","") |
| 104 | markdown_element_list.append(new_text) |
| 105 | else: |
| 106 | markdown_element_list.append(text) |
| 107 | page_to_markdown_result[i+1] = markdown_element_list |
| 108 | except: |
| 109 | retry_list.append(i) |
| 110 | |
| 111 | attempt = 0 |
| 112 | while len(retry_list) > 0 and attempt < max_page_retries: |
| 113 | retry_page_to_markdown_query_list = [build_page_to_markdown_query(file_path,i+1) for i in retry_list] |
| 114 | retry_sampling_params = SamplingParams(temperature=0.1*attempt, max_tokens=8192) |
| 115 | responses = llm.generate(retry_page_to_markdown_query_list, sampling_params=retry_sampling_params) |
| 116 | results = [response.outputs[0].text for response in responses] |
| 117 | next_retry_list = [] |
| 118 | for i,result in zip(retry_list,results): |
| 119 | try: |
| 120 | json_data = json.loads(result) |
| 121 | page_response = PageResponse(**json_data) |
| 122 | natural_text = page_response.natural_text |
| 123 | markdown_element_list = [] |
| 124 | for text in natural_text.split('\n\n'): |
| 125 | if text.startswith("<Image>") and text.endswith("</Image>"): |
| 126 | pass |
| 127 | elif text.startswith("<table>") and text.endswith("</table>"): |
| 128 | try: |
| 129 | new_text = table_matrix2html(text) |
no test coverage detected