Skip to content

Commit ab4523b

Browse files
authored
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
1 parent d70a3ed commit ab4523b

File tree

6 files changed

+57
-11
lines changed

6 files changed

+57
-11
lines changed

runtime/datamate-python/app/db/session.py

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,7 +10,9 @@
1010
engine = create_async_engine(
1111
settings.database_url,
1212
echo=False, # 关闭SQL调试日志以减少输出
13-
future=True
13+
future=True,
14+
pool_pre_ping=True,
15+
pool_recycle=3600
1416
)
1517

1618
# 创建会话工厂

runtime/ops/examples/test_operator/metadata.yml

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -4,9 +4,9 @@ language: 'python'
44
vendor: 'huawei'
55
raw_id: 'TestMapper'
66
version: '1.0.0'
7-
modal: 'text'
8-
inputs: 'text'
9-
outputs: 'text'
7+
modal: 'text' # text/image/audio/video/multimodal
8+
inputs: 'text' # text/image/audio/video/multimodal
9+
outputs: 'text' # text/image/audio/video/multimodal
1010
settings:
1111
sliderTest:
1212
name: '滑窗测试'

runtime/ops/formatter/mineru_formatter/metadata.yml

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,3 +14,15 @@ effect:
1414
after: ''
1515
inputs: 'text'
1616
outputs: 'text'
17+
settings:
18+
exportType:
19+
name: '导出类型'
20+
description: '指定清洗结果文件类型。若指定为md且后续存在其他清洗算子,可能导致文件格式错乱。'
21+
type: 'select'
22+
defaultVal: 'markdown'
23+
required: false
24+
options:
25+
- label: 'markdown'
26+
value: 'md'
27+
- label: 'txt'
28+
value: 'txt'

runtime/ops/formatter/mineru_formatter/process.py

Lines changed: 25 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,12 +6,16 @@
66
Create: 2025/10/29 17:24
77
"""
88
import asyncio
9+
import glob
910
import os
1011
import shutil
1112
import time
13+
import uuid
14+
from pathlib import Path
1215
from typing import Dict, Any
1316

14-
from datamate.core.base_op import Mapper
17+
from datamate.core.base_op import Mapper, FileExporter
18+
from datamate.sql_manager.persistence_atction import TaskInfoPersistence
1519
from loguru import logger
1620
from mineru.cli.common import aio_do_parse, read_fn
1721
from mineru.cli.fast_api import get_infer_result
@@ -27,6 +31,7 @@ def __init__(self, *args, **kwargs):
2731
self.backend = "vlm-http-client"
2832
self.output_dir = "/dataset/outputs"
2933
self.max_retries = 3
34+
self.target_type = kwargs.get("exportType", "md")
3035

3136
def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
3237
start = time.time()
@@ -35,6 +40,7 @@ def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
3540
return sample
3641
try:
3742
sample[self.text_key] = asyncio.run(self.async_process_file(sample))
43+
sample[self.target_type_key] = self.target_type
3844
logger.info(
3945
f"fileName: {filename}, method: MineruFormatter costs {(time.time() - start):6f} s")
4046
except Exception as e:
@@ -77,5 +83,23 @@ async def async_process_file(self, sample):
7783
raise # 耗尽次数后抛出异常,交给上层 execute 处理
7884
if os.path.exists(parse_dir):
7985
content += get_infer_result(".md", filename_without_ext, parse_dir)
86+
self.save_images(parse_dir, sample["dataset_id"], os.path.abspath(sample[self.export_path_key]) + "/images")
8087
shutil.rmtree(parse_dir)
8188
return content
89+
90+
def save_images(self, parse_dir, dataset_id, export_path):
91+
Path(export_path).mkdir(parents=True, exist_ok=True)
92+
93+
images_dir = os.path.join(parse_dir, "images")
94+
image_paths = glob.glob(os.path.join(glob.escape(images_dir), "*.jpg"))
95+
for image_path in image_paths:
96+
shutil.copy(image_path, export_path)
97+
image_sample = {}
98+
image = Path(image_path)
99+
image_name = image.name
100+
image_sample[self.filename_key] = image_name
101+
image_sample[self.filetype_key] = "jpg"
102+
image_sample[self.filesize_key] = image.stat().st_size
103+
image_sample["dataset_id"] = dataset_id
104+
image_sample[self.filepath_key] = export_path + "/" + image_name
105+
TaskInfoPersistence().update_file_result(image_sample, str(uuid.uuid4()))

runtime/python-executor/datamate/core/base_op.py

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,7 @@ def __init__(self, *args, **kwargs):
7272
self.filesize_key = kwargs.get('fileSize_key', "fileSize")
7373
self.export_path_key = kwargs.get('export_path_key', "export_path")
7474
self.ext_params_key = kwargs.get('ext_params_key', "ext_params")
75+
self.target_type_key = kwargs.get('target_type_key', "target_type")
7576

7677
@property
7778
def name(self):
@@ -437,7 +438,7 @@ def execute(self, sample: Dict[str, Any]):
437438
elif file_type in self.medical_support_ext:
438439
sample, save_path = self.get_medicalfile_handler(sample)
439440
else:
440-
raise TypeError(f"{file_type} is unsupported! please check support_ext in FileExporter Ops")
441+
return False
441442

442443
if sample[self.text_key] == '' and sample[self.data_key] == b'':
443444
sample[self.filesize_key] = "0"
@@ -483,11 +484,11 @@ def get_textfile_handler(self, sample: Dict[str, Any]):
483484

484485
# target_type存在则保存为扫描件, docx格式
485486
if target_type:
486-
sample = self._get_from_data(sample)
487+
sample = self._get_from_text_or_data(sample)
487488
save_path = self.get_save_path(sample, target_type)
488489
# 不存在则保存为txt文件,正常文本清洗
489490
else:
490-
sample = self._get_from_text(sample)
491+
sample = self._get_from_text_or_data(sample)
491492
save_path = self.get_save_path(sample, 'txt')
492493
return sample, save_path
493494

@@ -496,11 +497,11 @@ def get_datafile_handler(self, sample: Dict[str, Any]):
496497

497498
# target_type存在, 图转文保存为target_type,markdown格式
498499
if target_type:
499-
sample = self._get_from_text(sample)
500+
sample = self._get_from_text_or_data(sample)
500501
save_path = self.get_save_path(sample, target_type)
501502
# 不存在则保存为原本图片文件格式,正常图片清洗
502503
else:
503-
sample = self._get_from_data(sample)
504+
sample = self._get_from_text_or_data(sample)
504505
save_path = self.get_save_path(sample, sample[self.filetype_key])
505506
return sample, save_path
506507

@@ -533,6 +534,13 @@ def _get_from_text(self, sample: Dict[str, Any]) -> Dict[str, Any]:
533534
sample[self.text_key] = str(sample[self.text_key])
534535
return sample
535536

537+
def _get_from_text_or_data(self, sample: Dict[str, Any]) -> Dict[str, Any]:
538+
if sample[self.data_key] is not None and sample[self.data_key] != b'':
539+
return self._get_from_data(sample)
540+
else:
541+
return self._get_from_text(sample)
542+
543+
536544
@staticmethod
537545
def _get_uuid():
538546
return str(uuid.uuid4())

scripts/images/runtime/Dockerfile

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,7 @@ FROM ghcr.io/astral-sh/uv:python3.11-bookworm
33
RUN --mount=type=cache,target=/var/cache/apt \
44
--mount=type=cache,target=/var/lib/apt \
55
apt update \
6-
&& apt install -y libgl1 libglib2.0-0 vim libmagic1 libreoffice dos2unix swig
6+
&& apt install -y libgl1 libglib2.0-0 vim libmagic1 libreoffice dos2unix swig poppler-utils tesseract-ocr
77

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

0 commit comments

Comments
 (0)