[pdf2vqa] 适配api返回格式

Author: fatty-belly

dataflow/operators/pdf2vqa/generate/llm_output_parser.py

@@ -122,8 +122,12 @@ def run(self, storage: DataFlowStorage,
                     outfile.write('\n')
 
             # 复制图片
-            src_dir = os.path.join(self.intermediate_dir, 'mineru', Path(converted_json_path).stem).replace('_content_list_converted','')
+            src_dir = converted_json_path.rpartition('/')[0]
             src_images = os.path.join(src_dir, 'vlm', 'images')
+            if not os.path.exists(src_images):
+                src_images = os.path.join(src_dir, 'images')
+            if not os.path.exists(src_images):
+                raise ValueError("Images directory not found! There might be a change in Mineru API!")
             dst_images = os.path.join(self.output_dir, image_prefix)
 
             try:

dataflow/operators/pdf2vqa/generate/mineru_to_llm_input_operator.py

@@ -3,6 +3,8 @@
 from dataflow.utils.registry import OPERATOR_REGISTRY
 from dataflow.utils.storage import DataFlowStorage
 
+from pathlib import Path
+
 @OPERATOR_REGISTRY.register()
 class MinerU2LLMInputOperator(OperatorABC):
     def __init__(self):
@@ -57,8 +59,13 @@ def run(self, storage: DataFlowStorage,
         dataframe = storage.read("dataframe")
 
         for index, row in dataframe.iterrows():
-            input_json_path = row[input_markdown_path_key].replace('.md', '_content_list.json')
-            converted_path = input_json_path.replace('.json', '_converted.json')
+            md_path = Path(row[input_markdown_path_key])
+            try:
+                input_json_path = list(md_path.parent.glob("*_content_list.json"))[0]
+            except:
+                raise ValueError("No _content_list.json file found in the api result. There might be an error with the Mineru api.")
+            
+            converted_path = str(input_json_path).replace('.json', '_converted.json')
             self._convert_json(input_json_path, converted_path)
             dataframe.at[index, output_converted_layout_key] = converted_path
 

dataflow/statics/pipelines/api_pipelines/pdf_vqa_extract_pipeline.py

@@ -1,4 +1,4 @@
-from dataflow.operators.knowledge_cleaning import FileOrURLToMarkdownConverterFlash
+from dataflow.operators.knowledge_cleaning import FileOrURLToMarkdownConverterFlash, FileOrURLToMarkdownConverterLocal, FileOrURLToMarkdownConverterAPI
 
 from dataflow.serving import APILLMServing_request
 from dataflow.utils.storage import FileStorage
@@ -30,7 +30,7 @@ def __init__(self):
         self.vqa_extract_prompt = QAExtractPrompt()
 
         self.pdf_merger = PDF_Merger(output_dir="./cache")
-        self.mineru_executor = FileOrURLToMarkdownConverterBatch(intermediate_dir = "intermediate", mineru_backend="vlm-vllm-engine")
+        self.mineru_executor = FileOrURLToMarkdownConverterAPI(intermediate_dir = "intermediate")
         self.input_formatter = MinerU2LLMInputOperator()
         self.vqa_extractor = ChunkedPromptedGenerator(
             llm_serving=self.llm_serving,
@@ -42,9 +42,9 @@ def __init__(self):
     def forward(self):
         self.pdf_merger.run(
             storage=self.storage.step(),
-            input_pdf_path_key="input_pdf_paths",
+            input_pdf_list_key="input_pdf_paths",
             input_name_key="name",
-            output_merged_pdf_path_key="merged_pdf_path",
+            output_pdf_path_key="merged_pdf_path",
         )
         self.mineru_executor.run(
             storage=self.storage.step(),