You've already forked DataMate
支持mineru npu处理 (#174)
* feature: unstructured支持简单pdf处理 * feature: update values.yaml to enhance ray-cluster configuration with security context, environment variables, and resource limits * feature: update deploy.yaml and process.py for mineru server configuration and PDF processing enhancements * feature: update deploy.yaml and process.py for mineru server configuration and PDF processing enhancements * feature: improve PDF processing logic and update dependencies in process.py and pyproject.toml * feature: improve PDF processing logic and update dependencies in process.py and pyproject.toml * feature: update Dockerfile for improved package source mirrors and add mineru-npu to build targets
This commit is contained in:
@@ -5,13 +5,17 @@
|
||||
Description: MinerU PDF文本抽取
|
||||
Create: 2025/10/29 17:24
|
||||
"""
|
||||
import json
|
||||
import os
|
||||
import shutil
|
||||
import time
|
||||
from loguru import logger
|
||||
from typing import Dict, Any
|
||||
|
||||
from datamate.core.base_op import Mapper
|
||||
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.fast_api import get_infer_result
|
||||
from pypdf import PdfReader
|
||||
|
||||
|
||||
class MineruFormatter(Mapper):
|
||||
@@ -19,21 +23,40 @@ class MineruFormatter(Mapper):
|
||||
|
||||
def __init__(self, *args, **kwargs):
|
||||
super(MineruFormatter, self).__init__(*args, **kwargs)
|
||||
self.base_url = "http://datamate-mineru:9001"
|
||||
self.pdf_extract_url = f"{self.base_url}/api/pdf-extract"
|
||||
self.server_url = "http://datamate-mineru:8000"
|
||||
self.backend = "vlm-http-client"
|
||||
self.output_dir = "/dataset/outputs"
|
||||
|
||||
def execute(self, sample: Dict[str, Any]) -> Dict[str, Any]:
|
||||
start = time.time()
|
||||
filename = sample[self.filename_key]
|
||||
if not filename.lower().endswith(".pdf"):
|
||||
filename_without_ext = os.path.splitext(filename)[0]
|
||||
if not filename.lower().endswith((".png", ".jpeg", ".jpg", ".webp", ".gif", ".pdf")):
|
||||
return sample
|
||||
try:
|
||||
data = {"source_path": sample[self.filepath_key], "export_path": sample[self.export_path_key]}
|
||||
response = http_request(method="POST", url=self.pdf_extract_url, data=data)
|
||||
sample[self.text_key] = json.loads(response.text).get("result")
|
||||
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(
|
||||
output_dir=self.output_dir,
|
||||
pdf_file_names=[filename_without_ext],
|
||||
pdf_bytes_list=[pdf_bytes],
|
||||
p_lang_list=["ch"],
|
||||
backend=self.backend,
|
||||
server_url=self.server_url,
|
||||
start_page_id=page,
|
||||
end_page_id=min(page + 9, total_page - 1),
|
||||
)
|
||||
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 UnicodeDecodeError as err:
|
||||
logger.exception(f"fileName: {filename}, method: MineruFormatter causes decode error: {err}")
|
||||
except Exception as e:
|
||||
logger.exception(f"fileName: {filename}, method: MineruFormatter causes error: {e}")
|
||||
raise
|
||||
return sample
|
||||
|
||||
@@ -11,7 +11,9 @@ dependencies = [
|
||||
"emoji>=2.15.0",
|
||||
"jieba>=0.42.1",
|
||||
"loguru>=0.7.3",
|
||||
"numpy==1.23.3",
|
||||
"mineru>=2.6.5",
|
||||
"numpy==1.24.3",
|
||||
"python-multipart>=0.0.20",
|
||||
"opencv-contrib-python-headless==4.7.0.72",
|
||||
"opencv-python-headless==4.7.0.72",
|
||||
"openslide-python>=1.4.3",
|
||||
@@ -29,4 +31,4 @@ dependencies = [
|
||||
"sqlalchemy>=2.0.44",
|
||||
"xmltodict>=1.0.2",
|
||||
"zhconv>=1.4.3",
|
||||
]
|
||||
]
|
||||
|
||||
Reference in New Issue
Block a user