[Fix Playground] Fix errors of filepath and import path (#389)

* [Fix Playground] Fix errors of filepath and import path

* [Fix Bug] fix bug of prompt template and import error

* [Fix Bug] Fix parameters name in run()

* [Fix Bug] Fix parameters name in run()

* [Fix Bug] Fix parameters name in run()

Author: HeRunming

dataflow/operators/knowledge_cleaning/__init__.py

@@ -8,7 +8,7 @@
     from .generate.file_or_url_to_markdown_converter_batch import FileOrURLToMarkdownConverterBatch
     from .generate.kbc_text_cleaner import KBCTextCleaner
     from .generate.kbc_text_cleaner_batch import KBCTextCleanerBatch
-    # from .generate.mathbook_question_extract import MathBookQuestionExtract
+    from .generate.mathbook_question_extract import MathBookQuestionExtract
     # from .generate.kbc_multihop_qa_generator import KBCMultiHopQAGenerator
     from .generate.kbc_multihop_qa_generator_batch import KBCMultiHopQAGeneratorBatch
     from .generate.qa_extract import QAExtractor

dataflow/operators/knowledge_cleaning/generate/mathbook_question_extract.py

@@ -11,22 +11,32 @@
 import re
 from openai import OpenAI
 import base64
-from typing import List, Literal
-from concurrent.futures import ThreadPoolExecutor, as_completed
+from typing import Literal, Union
 from dataflow.core import LLMServingABC
 from dataflow.serving import APIVLMServing_openai
-
-
+from dataflow.core.prompt import DIYPromptABC
+from dataflow.utils.storage import DataFlowStorage
 
 @OPERATOR_REGISTRY.register()
 class MathBookQuestionExtract(OperatorABC):
-    def __init__(self, llm_serving: APIVLMServing_openai, prompt_template = None):
+    def __init__(self, 
+                 llm_serving: APIVLMServing_openai, 
+                 prompt_template: Union[MathbookQuestionExtractPrompt, DIYPromptABC] = MathbookQuestionExtractPrompt(),
+                 mineru_backend: str = "vlm-vllm-engine",
+                 dpi: int = 300,
+                 key_name_of_api_key: str = "DF_API_KEY",
+                 model_name: str = "o4-mini",
+                 max_workers: int = 20
+                ):
         self.logger = get_logger()
         self.llm_serving = llm_serving
-        if prompt_template:
-            self.prompt_template = prompt_template
-        else:
-            self.prompt_template = MathbookQuestionExtractPrompt()
+        self.prompt_template = prompt_template
+        
+        self.mineru_backend = mineru_backend
+        self.dpi = dpi
+        self.key_name_of_api_key = key_name_of_api_key
+        self.model_name = model_name
+        self.max_workers = max_workers # 注意：这个参数在原逻辑中并未被使用，但仍按要求移入init
 
     @staticmethod   
     def get_desc(lang: str = "zh"):
@@ -80,7 +90,6 @@ def mineru2_runner(self,
                         # pipeline|vlm-transformers|vlm-vllm-engine|vlm-http-client
                         mineru_backend: Literal["pipeline", "vlm-transformers", "vlm-vllm-engine", "vlm-http-client"] = "pipeline"
                         ):
-
         try:
             import mineru
         except ImportError:
@@ -227,6 +236,7 @@ def process_input(self,
         return full_input_image_list,full_input_label_list
 
     def analyze_and_save(self,result_list,save_folder,img_folder,output_file_name):
+        # ... (analyze_and_save 方法保持不变)
         # make save_folder if not exist
         if not os.path.exists(save_folder):
             os.makedirs(save_folder)
@@ -281,42 +291,38 @@ def analyze_and_save(self,result_list,save_folder,img_folder,output_file_name):
 
     def run(
         self,
-        pdf_file_path: str,
+        storage: DataFlowStorage,
+        input_pdf_file_path: str,
         output_file_name: str,
         output_folder: str,
-        MinerU_Backend: str = "vlm-sglang-engine",
-        dpi: int = 300,
-        api_url: str = "http://123.129.219.111:3000/v1",
-        key_name_of_api_key: str = "DF_API_KEY",
-        model_name: str = "o4-mini",
-        max_workers: int = 20
     ):
-        api_key = os.environ.get(key_name_of_api_key)
+        # get the configuration parameters from self
+        api_key = os.environ.get(self.key_name_of_api_key)
         if not api_key:
-            raise ValueError(f"API key not found in environment variable {key_name_of_api_key}")
+            raise ValueError(f"API key not found in environment variable {self.key_name_of_api_key}")
 
         # 1. convert pdf to images
-        pdf2images_folder_name = output_folder+"/pdfimages"
-        self.pdf2images(pdf_file_path, pdf2images_folder_name, dpi)
+        pdf2images_folder_name = os.path.join(output_folder, "pdfimages")
+        self.pdf2images(input_pdf_file_path, pdf2images_folder_name, self.dpi)
 
         # 2. use mineru to extract content and pics
-        json_content_file, pic_folder = self.mineru2_runner(pdf_file_path, output_folder, MinerU_Backend)
+        json_content_file, pic_folder = self.mineru2_runner(input_pdf_file_path, output_folder, self.mineru_backend)
 
         # 3. organize_pics
-        output_image_folder = output_folder+"/organized_images"
-        output_json_file = output_folder+"/organized_images/organized_info.json"
-        self.organize_pics(json_content_file, pic_folder,output_json_file, output_image_folder)
+        output_image_folder = os.path.join(output_folder, "organized_images")
+        output_json_file = os.path.join(output_image_folder, "organized_info.json")
+        self.organize_pics(json_content_file, pic_folder, output_json_file, output_image_folder)
 
         # 4. process input
-        full_input_image_list,full_input_label_list = self.process_input(pdf2images_folder_name, output_json_file)
+        full_input_image_list, full_input_label_list = self.process_input(pdf2images_folder_name, output_json_file)
 
         # 5. init server and generate
         system_prompt = self.prompt_template.build_prompt()
         result_text_list = self.llm_serving.generate_from_input_multi_images(
             list_of_image_paths=full_input_image_list,
             list_of_image_labels=full_input_label_list,
             system_prompt=system_prompt,
-            model=model_name,
+            model=self.model_name,
             timeout=1800
         )
 

dataflow/prompts/kbcleaning.py

@@ -1,6 +1,6 @@
 from dataflow.utils.registry import PROMPT_REGISTRY
 from dataflow.core.prompt import PromptABC
-
+import re
 @PROMPT_REGISTRY.register()
 class KnowledgeCleanerPrompt(PromptABC):
    '''
@@ -220,7 +220,7 @@ def _post_process(self, cleaned_text: str) -> str:
 
 
 @PROMPT_REGISTRY.register()
-class MathbookQuestionExtractPrompt:
+class MathbookQuestionExtractPrompt(PromptABC):
    def __init__(self):
       pass
 

dataflow/statics/playground/playground/mathbook_extract.py

@@ -1,41 +1,43 @@
-from dataflow.operators.knowledge_cleaning.generate.mathbook_question_extract import MathBookQuestionExtract
-from dataflow.serving.api_vlm_serving_openai import APIVLMServing_openai
+from dataflow.operators.knowledge_cleaning import MathBookQuestionExtract
+from dataflow.serving import APIVLMServing_openai
 
 class QuestionExtractPipeline:
-    def __init__(self, llm_serving: APIVLMServing_openai):
-        self.extractor = MathBookQuestionExtract(llm_serving)
-        self.test_pdf = "../example_data/KBCleaningPipeline/questionextract_test.pdf" 
+    def __init__(self, 
+                 llm_serving: APIVLMServing_openai,
+                 api_url: str = "https://api.openai.com/v1", # end with /v1
+                 key_name_of_api_key: str = "DF_API_KEY", # set in environment first: export DF_API_KEY="your_openai_api_key"
+                 model_name: str = "o4-mini",
+                 max_workers: int = 20
+                 ):
+        self.extractor = MathBookQuestionExtract(
+            llm_serving=llm_serving,
+            key_name_of_api_key=key_name_of_api_key,
+            model_name=model_name,
+            max_workers=max_workers
+        )
+        self.test_pdf = "../example_data/PDF2VQAPipeline/questionextract_test.pdf" 
 
     def forward(
         self,
         pdf_path: str,
         output_name: str,
         output_dir: str,
-        api_url: str = "https://api.openai.com/v1/chat/completions",
-        key_name_of_api_key: str = "DF_API_KEY",
-        model_name: str = "o4-mini",
-        max_workers: int = 20
     ):
         self.extractor.run(
-            pdf_file_path=pdf_path,
+            storage=None,
+            input_pdf_file_path=pdf_path,
             output_file_name=output_name,
-            output_folder=output_dir,
-            api_url=api_url,
-            key_name_of_api_key=key_name_of_api_key,
-            model_name=model_name,
-            max_workers=max_workers
+            output_folder=output_dir
         )
 
 if __name__ == "__main__":
-    # 1. initialize LLM Serving
     llm_serving = APIVLMServing_openai(
-        api_url="https://api.openai.com/v1",  # end with /v1, DO NOT add /chat/completions
-        model_name="o4-mini",      # recommend using strong reasoning model
-        max_workers=20             # number of concurrent requests
+        api_url="https://api.openai.com/v1",
+        model_name="o4-mini",
+        max_workers=20
     )
 
-    # 2. construct and run pipeline
-    pipeline = QuestionExtractPipeline(llm_serving)
+    pipeline = QuestionExtractPipeline(llm_serving=llm_serving)
     pipeline.forward(
         pdf_path=pipeline.test_pdf,
         output_name="test_question_extract",

dataflow/statics/playground/playground/vqa.py

@@ -1,5 +1,5 @@
 from dataflow.operators.core_vision import PromptedVQAGenerator
-from dataflow.serving.APIVLMServing_openai import APIVLMServing_openai
+from dataflow.serving import APIVLMServing_openai
 from dataflow.utils.storage import FileStorage
 
 class VQA_generator():