add export type settings and enhance metadata structure (#181)

* fix(session): enhance database connection settings with pool pre-ping and recycle options * feat(metadata): add export type settings and enhance metadata structure * fix(base_op): improve sample handling by introducing target_type key and consolidating text/data retrieval logic * feat(metadata): add export type settings and enhance metadata structure * feat(metadata): add export type settings and enhance metadata structure
2025-12-19 11:54:08 +08:00
parent d70a3eda0d
commit ab4523b556
6 changed files with 57 additions and 11 deletions
--- a/runtime/datamate-python/app/db/session.py
+++ b/runtime/datamate-python/app/db/session.py
@@ -10,7 +10,9 @@ logger = get_logger(__name__)
 engine = create_async_engine(
    settings.database_url,
    echo=False,  # 关闭SQL调试日志以减少输出
-    future=True
+    future=True,
+    pool_pre_ping=True,
+    pool_recycle=3600
 )

 # 创建会话工厂
--- a/runtime/ops/examples/test_operator/metadata.yml
+++ b/runtime/ops/examples/test_operator/metadata.yml
@@ -4,9 +4,9 @@ language: 'python'
 vendor: 'huawei'
 raw_id: 'TestMapper'
 version: '1.0.0'
-modal: 'text'
-inputs: 'text'
-outputs: 'text'
+modal: 'text'     # text/image/audio/video/multimodal
+inputs: 'text'    # text/image/audio/video/multimodal
+outputs: 'text'   # text/image/audio/video/multimodal
 settings:
  sliderTest:
    name: '滑窗测试'
--- a/runtime/ops/formatter/mineru_formatter/metadata.yml
+++ b/runtime/ops/formatter/mineru_formatter/metadata.yml
@@ -14,3 +14,15 @@ effect:
  after: ''
 inputs: 'text'
 outputs: 'text'
+settings:
+  exportType:
+    name: '导出类型'
+    description: '指定清洗结果文件类型。若指定为md且后续存在其他清洗算子，可能导致文件格式错乱。'
+    type: 'select'
+    defaultVal: 'markdown'
+    required: false
+    options:
+      - label: 'markdown'
+        value: 'md'
+      - label: 'txt'
+        value: 'txt'
--- a/runtime/ops/formatter/mineru_formatter/process.py
+++ b/runtime/ops/formatter/mineru_formatter/process.py
@@ -6,12 +6,16 @@ Description: MinerU PDF文本抽取
 Create: 2025/10/29 17:24
 """
 import asyncio
+import glob
 import os
 import shutil
 import time
+import uuid
+from pathlib import Path
 from typing import Dict, Any

-from datamate.core.base_op import Mapper
+from datamate.core.base_op import Mapper, FileExporter
+from datamate.sql_manager.persistence_atction import TaskInfoPersistence
 from loguru import logger
 from mineru.cli.common import aio_do_parse, read_fn
 from mineru.cli.fast_api import get_infer_result
@@ -27,6 +31,7 @@ class MineruFormatter(Mapper):
        self.backend = "vlm-http-client"
        self.output_dir = "/dataset/outputs"
        self.max_retries = 3
+        self.target_type = kwargs.get("exportType", "md")

    def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
        start = time.time()
@@ -35,6 +40,7 @@ class MineruFormatter(Mapper):
            return sample
        try:
            sample[self.text_key] = asyncio.run(self.async_process_file(sample))
+            sample[self.target_type_key] = self.target_type
            logger.info(
                f"fileName: {filename}, method: MineruFormatter costs {(time.time() - start):6f} s")
        except Exception as e:
@@ -77,5 +83,23 @@ class MineruFormatter(Mapper):
                        raise  # 耗尽次数后抛出异常，交给上层 execute 处理
            if os.path.exists(parse_dir):
                content += get_infer_result(".md", filename_without_ext, parse_dir)
+                self.save_images(parse_dir, sample["dataset_id"], os.path.abspath(sample[self.export_path_key]) + "/images")
                shutil.rmtree(parse_dir)
        return content
+
+    def save_images(self, parse_dir, dataset_id, export_path):
+        Path(export_path).mkdir(parents=True, exist_ok=True)
+
+        images_dir = os.path.join(parse_dir, "images")
+        image_paths = glob.glob(os.path.join(glob.escape(images_dir), "*.jpg"))
+        for image_path in image_paths:
+            shutil.copy(image_path, export_path)
+            image_sample = {}
+            image = Path(image_path)
+            image_name = image.name
+            image_sample[self.filename_key] = image_name
+            image_sample[self.filetype_key] = "jpg"
+            image_sample[self.filesize_key] = image.stat().st_size
+            image_sample["dataset_id"] = dataset_id
+            image_sample[self.filepath_key] = export_path + "/" + image_name
+            TaskInfoPersistence().update_file_result(image_sample, str(uuid.uuid4()))
--- a/runtime/python-executor/datamate/core/base_op.py
+++ b/runtime/python-executor/datamate/core/base_op.py
@@ -72,6 +72,7 @@ class BaseOp:
        self.filesize_key = kwargs.get('fileSize_key', "fileSize")
        self.export_path_key = kwargs.get('export_path_key', "export_path")
        self.ext_params_key = kwargs.get('ext_params_key', "ext_params")
+        self.target_type_key = kwargs.get('target_type_key', "target_type")

    @property
    def name(self):
@@ -437,7 +438,7 @@ class FileExporter(BaseOp):
            elif file_type in self.medical_support_ext:
                sample, save_path = self.get_medicalfile_handler(sample)
            else:
-                raise TypeError(f"{file_type} is unsupported! please check support_ext in FileExporter Ops")
+                return False

            if sample[self.text_key] == '' and sample[self.data_key] == b'':
                sample[self.filesize_key] = "0"
@@ -483,11 +484,11 @@ class FileExporter(BaseOp):

        # target_type存在则保存为扫描件, docx格式
        if target_type:
-            sample = self._get_from_data(sample)
+            sample = self._get_from_text_or_data(sample)
            save_path = self.get_save_path(sample, target_type)
        # 不存在则保存为txt文件，正常文本清洗
        else:
-            sample = self._get_from_text(sample)
+            sample = self._get_from_text_or_data(sample)
            save_path = self.get_save_path(sample, 'txt')
        return sample, save_path

@@ -496,11 +497,11 @@ class FileExporter(BaseOp):

        # target_type存在, 图转文保存为target_type，markdown格式
        if target_type:
-            sample = self._get_from_text(sample)
+            sample = self._get_from_text_or_data(sample)
            save_path = self.get_save_path(sample, target_type)
        # 不存在则保存为原本图片文件格式，正常图片清洗
        else:
-            sample = self._get_from_data(sample)
+            sample = self._get_from_text_or_data(sample)
            save_path = self.get_save_path(sample, sample[self.filetype_key])
        return sample, save_path

@@ -533,6 +534,13 @@ class FileExporter(BaseOp):
        sample[self.text_key] = str(sample[self.text_key])
        return sample

+    def _get_from_text_or_data(self, sample: Dict[str, Any]) -> Dict[str, Any]:
+        if sample[self.data_key] is not None and sample[self.data_key] != b'':
+            return self._get_from_data(sample)
+        else:
+            return self._get_from_text(sample)
+
+
    @staticmethod
    def _get_uuid():
        return str(uuid.uuid4())
--- a/scripts/images/runtime/Dockerfile
+++ b/scripts/images/runtime/Dockerfile
@@ -3,7 +3,7 @@ FROM ghcr.io/astral-sh/uv:python3.11-bookworm
 RUN --mount=type=cache,target=/var/cache/apt \
    --mount=type=cache,target=/var/lib/apt \
    apt update \
-    && apt install -y libgl1 libglib2.0-0 vim libmagic1 libreoffice dos2unix swig
+    && apt install -y libgl1 libglib2.0-0 vim libmagic1 libreoffice dos2unix swig poppler-utils tesseract-ocr

 RUN mkdir -p /home/models \
    && wget https://paddleocr.bj.bcebos.com/dygraph_v2.0/ch/ch_ppocr_mobile_v2.0_cls_infer.tar \