You've already forked DataMate
fix: optimize PDF parsing by implementing concurrent processing with … (#177)
* 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:
@@ -5,15 +5,15 @@
|
||||
Description: MinerU PDF文本抽取
|
||||
Create: 2025/10/29 17:24
|
||||
"""
|
||||
import asyncio
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from typing import Dict, Any
|
||||
|
||||
from datamate.common.utils.rest_client import http_request
|
||||
from datamate.core.base_op import Mapper
|
||||
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 pypdf import PdfReader
|
||||
|
||||
@@ -30,17 +30,28 @@ class MineruFormatter(Mapper):
|
||||
def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
|
||||
start = time.time()
|
||||
filename = sample[self.filename_key]
|
||||
filename_without_ext = os.path.splitext(filename)[0]
|
||||
if not filename.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif", ".pdf")):
|
||||
return sample
|
||||
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]
|
||||
parse_dir = os.path.join(self.output_dir, filename_without_ext, "vlm")
|
||||
pdf_bytes = read_fn(filepath)
|
||||
total_page = len(PdfReader(filepath).pages)
|
||||
content = ""
|
||||
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,
|
||||
pdf_file_names=[filename_without_ext],
|
||||
pdf_bytes_list=[pdf_bytes],
|
||||
@@ -53,10 +64,4 @@ class MineruFormatter(Mapper):
|
||||
if os.path.exists(parse_dir):
|
||||
content += get_infer_result(".md", filename_without_ext, parse_dir)
|
||||
shutil.rmtree(parse_dir)
|
||||
sample[self.text_key] = 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
|
||||
return content
|
||||
|
||||
Reference in New Issue
Block a user