q792602257/DataMate
0
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix: optimize PDF parsing by implementing concurrent processing with … (#177)

Browse Source 
* fix: optimize PDF parsing by implementing concurrent processing with ThreadPoolExecutor

* Refactor to async processing for file extraction

Refactor the file processing to use asyncio for improved performance and concurrency.
This commit is contained in:
hhhhsc701
2025-12-18 15:28:30 +08:00
committed by GitHub
parent 8113840ac7
commit 761f7f6a51
1 changed files with 28 additions and 23 deletions
Download Patch File Download Diff File Expand all files Collapse all files

27
runtime/ops/formatter/mineru_formatter/process.py
View File

@@ -5,15 +5,15 @@
Description: MinerU PDF文本抽取 Description: MinerU PDF文本抽取
Create: 2025/10/29 17:24 Create: 2025/10/29 17:24
""" """
import asyncio
import os import os
import shutil import shutil
import time import time
from typing import Dict, Any from typing import Dict, Any
from datamate.common.utils.rest_client import http_request
from datamate.core.base_op import Mapper from datamate.core.base_op import Mapper
from loguru import logger from loguru import logger
from mineru.cli.common import do_parse, read_fn from mineru.cli.common import aio_do_parse, read_fn
from mineru.cli.fast_api import get_infer_result from mineru.cli.fast_api import get_infer_result
from pypdf import PdfReader from pypdf import PdfReader
@@ -30,17 +30,28 @@ class MineruFormatter(Mapper):
def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]: def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
start = time.time() start = time.time()
filename = sample[self.filename_key] filename = sample[self.filename_key]
filename_without_ext = os.path.splitext(filename)[0]
if not filename.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif", ".pdf")): if not filename.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif", ".pdf")):
return sample return sample
try: try:
sample[self.text_key] = asyncio.run(self.async_process_file(sample))
logger.info(
f"fileName: {filename}, method: MineruFormatter costs {(time.time() - start):6f} s")
except Exception as e:
logger.exception(f"fileName: {filename}, method: MineruFormatter causes error: {e}")
raise
return sample
async def async_process_file(self, sample):
filename = sample[self.filename_key]
filename_without_ext = os.path.splitext(filename)[0]
filepath = sample[self.filepath_key] filepath = sample[self.filepath_key]
parse_dir = os.path.join(self.output_dir, filename_without_ext, "vlm") parse_dir = os.path.join(self.output_dir, filename_without_ext, "vlm")
pdf_bytes = read_fn(filepath) pdf_bytes = read_fn(filepath)
total_page = len(PdfReader(filepath).pages) total_page = len(PdfReader(filepath).pages)
content = "" content = ""
for page in range(0, total_page, 10): for page in range(0, total_page, 10):
do_parse( logger.info(f"fileName: {filename}, total_page: {total_page}, page: {page}.")
await aio_do_parse(
output_dir=self.output_dir, output_dir=self.output_dir,
pdf_file_names=[filename_without_ext], pdf_file_names=[filename_without_ext],
pdf_bytes_list=[pdf_bytes], pdf_bytes_list=[pdf_bytes],
@@ -53,10 +64,4 @@ class MineruFormatter(Mapper):
if os.path.exists(parse_dir): if os.path.exists(parse_dir):
content += get_infer_result(".md", filename_without_ext, parse_dir) content += get_infer_result(".md", filename_without_ext, parse_dir)
shutil.rmtree(parse_dir) shutil.rmtree(parse_dir)
sample[self.text_key] = content return content
logger.info(
f"fileName: {filename}, method: MineruFormatter costs {(time.time() - start):6f} s")
except Exception as e:
logger.exception(f"fileName: {filename}, method: MineruFormatter causes error: {e}")
raise
return sample