Initial commit

mineru/cli/client.py

+# Copyright (c) Opendatalab. All rights reserved.
+import os
+import click
+from pathlib import Path
+from loguru import logger
+
+from mineru.utils.config_reader import get_device
+from mineru.utils.model_utils import get_vram
+from ..version import __version__
+from .common import do_parse, read_fn, pdf_suffixes, image_suffixes
+
+@click.command()
+@click.version_option(__version__,
+                      '--version',
+                      '-v',
+                      help='display the version and exit')
+@click.option(
+    '-p',
+    '--path',
+    'input_path',
+    type=click.Path(exists=True),
+    required=True,
+    help='local filepath or directory. support pdf, png, jpg, jpeg files',
+)
+@click.option(
+    '-o',
+    '--output',
+    'output_dir',
+    type=click.Path(),
+    required=True,
+    help='output local directory',
+)
+@click.option(
+    '-m',
+    '--method',
+    'method',
+    type=click.Choice(['auto', 'txt', 'ocr']),
+    help="""the method for parsing pdf:
+    auto: Automatically determine the method based on the file type.
+    txt: Use text extraction method.
+    ocr: Use OCR method for image-based PDFs.
+    Without method specified, 'auto' will be used by default.
+    Adapted only for the case where the backend is set to "pipeline".""",
+    default='auto',
+)
+@click.option(
+    '-b',
+    '--backend',
+    'backend',
+    type=click.Choice(['pipeline', 'vlm-transformers', 'vlm-sglang-engine', 'vlm-sglang-client']),
+    help="""the backend for parsing pdf:
+    pipeline: More general.
+    vlm-transformers: More general.
+    vlm-sglang-engine: Faster(engine).
+    vlm-sglang-client: Faster(client).
+    without method specified, pipeline will be used by default.""",
+    default='pipeline',
+)
+@click.option(
+    '-l',
+    '--lang',
+    'lang',
+    type=click.Choice(['ch', 'ch_server', 'ch_lite', 'en', 'korean', 'japan', 'chinese_cht', 'ta', 'te', 'ka']),
+    help="""
+    Input the languages in the pdf (if known) to improve OCR accuracy.  Optional.
+    Without languages specified, 'ch' will be used by default.
+    Adapted only for the case where the backend is set to "pipeline".
+    """,
+    default='ch',
+)
+@click.option(
+    '-u',
+    '--url',
+    'server_url',
+    type=str,
+    help="""
+    When the backend is `sglang-client`, you need to specify the server_url, for example:`http://127.0.0.1:30000`
+    """,
+    default=None,
+)
+@click.option(
+    '-s',
+    '--start',
+    'start_page_id',
+    type=int,
+    help='The starting page for PDF parsing, beginning from 0.',
+    default=0,
+)
+@click.option(
+    '-e',
+    '--end',
+    'end_page_id',
+    type=int,
+    help='The ending page for PDF parsing, beginning from 0.',
+    default=None,
+)
+@click.option(
+    '-f',
+    '--formula',
+    'formula_enable',
+    type=bool,
+    help='Enable formula parsing. Default is True. Adapted only for the case where the backend is set to "pipeline".',
+    default=True,
+)
+@click.option(
+    '-t',
+    '--table',
+    'table_enable',
+    type=bool,
+    help='Enable table parsing. Default is True. Adapted only for the case where the backend is set to "pipeline".',
+    default=True,
+)
+@click.option(
+    '-d',
+    '--device',
+    'device_mode',
+    type=str,
+    help='Device mode for model inference, e.g., "cpu", "cuda", "cuda:0", "npu", "npu:0", "mps". Adapted only for the case where the backend is set to "pipeline". ',
+    default=None,
+)
+@click.option(
+    '--vram',
+    'virtual_vram',
+    type=int,
+    help='Upper limit of GPU memory occupied by a single process. Adapted only for the case where the backend is set to "pipeline". ',
+    default=None,
+)
+@click.option(
+    '--source',
+    'model_source',
+    type=click.Choice(['huggingface', 'modelscope', 'local']),
+    help="""
+    The source of the model repository. Default is 'huggingface'.
+    """,
+    default='huggingface',
+)
+
+
+def main(input_path, output_dir, method, backend, lang, server_url, start_page_id, end_page_id, formula_enable, table_enable, device_mode, virtual_vram, model_source):
+
+    if not backend.endswith('-client'):
+        def get_device_mode() -> str:
+            if device_mode is not None:
+                return device_mode
+            else:
+                return get_device()
+        if os.getenv('MINERU_DEVICE_MODE', None) is None:
+            os.environ['MINERU_DEVICE_MODE'] = get_device_mode()
+
+        def get_virtual_vram_size() -> int:
+            if virtual_vram is not None:
+                return virtual_vram
+            if get_device_mode().startswith("cuda") or get_device_mode().startswith("npu"):
+                return round(get_vram(get_device_mode()))
+            return 1
+        if os.getenv('MINERU_VIRTUAL_VRAM_SIZE', None) is None:
+            os.environ['MINERU_VIRTUAL_VRAM_SIZE']= str(get_virtual_vram_size())
+
+        if os.getenv('MINERU_MODEL_SOURCE', None) is None:
+            os.environ['MINERU_MODEL_SOURCE'] = model_source
+
+    os.makedirs(output_dir, exist_ok=True)
+
+    def parse_doc(path_list: list[Path]):
+        try:
+            file_name_list = []
+            pdf_bytes_list = []
+            lang_list = []
+            for path in path_list:
+                file_name = str(Path(path).stem)
+                pdf_bytes = read_fn(path)
+                file_name_list.append(file_name)
+                pdf_bytes_list.append(pdf_bytes)
+                lang_list.append(lang)
+            do_parse(
+                output_dir=output_dir,
+                pdf_file_names=file_name_list,
+                pdf_bytes_list=pdf_bytes_list,
+                p_lang_list=lang_list,
+                backend=backend,
+                parse_method=method,
+                p_formula_enable=formula_enable,
+                p_table_enable=table_enable,
+                server_url=server_url,
+                start_page_id=start_page_id,
+                end_page_id=end_page_id
+            )
+        except Exception as e:
+            logger.exception(e)
+
+    if os.path.isdir(input_path):
+        doc_path_list = []
+        for doc_path in Path(input_path).glob('*'):
+            if doc_path.suffix in pdf_suffixes + image_suffixes:
+                doc_path_list.append(doc_path)
+        parse_doc(doc_path_list)
+    else:
+        parse_doc([Path(input_path)])
+
+if __name__ == '__main__':
+    main()

mineru/cli/common.py

+# Copyright (c) Opendatalab. All rights reserved.
+import io
+import json
+import os
+import copy
+from pathlib import Path
+
+import pypdfium2 as pdfium
+from loguru import logger
+
+from mineru.data.data_reader_writer import FileBasedDataWriter
+from mineru.utils.draw_bbox import draw_layout_bbox, draw_span_bbox
+from mineru.utils.enum_class import MakeMode
+from mineru.utils.pdf_image_tools import images_bytes_to_pdf_bytes
+from mineru.backend.vlm.vlm_middle_json_mkcontent import union_make as vlm_union_make
+from mineru.backend.vlm.vlm_analyze import doc_analyze as vlm_doc_analyze
+
+pdf_suffixes = [".pdf"]
+image_suffixes = [".png", ".jpeg", ".jpg"]
+
+
+def read_fn(path):
+    if not isinstance(path, Path):
+        path = Path(path)
+    with open(str(path), "rb") as input_file:
+        file_bytes = input_file.read()
+        if path.suffix in image_suffixes:
+            return images_bytes_to_pdf_bytes(file_bytes)
+        elif path.suffix in pdf_suffixes:
+            return file_bytes
+        else:
+            raise Exception(f"Unknown file suffix: {path.suffix}")
+
+
+def prepare_env(output_dir, pdf_file_name, parse_method):
+    local_md_dir = str(os.path.join(output_dir, pdf_file_name, parse_method))
+    local_image_dir = os.path.join(str(local_md_dir), "images")
+    os.makedirs(local_image_dir, exist_ok=True)
+    os.makedirs(local_md_dir, exist_ok=True)
+    return local_image_dir, local_md_dir
+
+
+def convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id=0, end_page_id=None):
+
+    # 从字节数据加载PDF
+    pdf = pdfium.PdfDocument(pdf_bytes)
+
+    # 确定结束页
+    end_page_id = end_page_id if end_page_id is not None and end_page_id >= 0 else len(pdf) - 1
+    if end_page_id > len(pdf) - 1:
+        logger.warning("end_page_id is out of range, use pdf_docs length")
+        end_page_id = len(pdf) - 1
+
+    # 创建一个新的PDF文档
+    output_pdf = pdfium.PdfDocument.new()
+
+    # 选择要导入的页面索引
+    page_indices = list(range(start_page_id, end_page_id + 1))
+
+    # 从原PDF导入页面到新PDF
+    output_pdf.import_pages(pdf, page_indices)
+
+    # 将新PDF保存到内存缓冲区
+    output_buffer = io.BytesIO()
+    output_pdf.save(output_buffer)
+
+    # 获取字节数据
+    output_bytes = output_buffer.getvalue()
+
+    pdf.close()  # 关闭原PDF文档以释放资源
+    output_pdf.close()  # 关闭新PDF文档以释放资源
+
+    return output_bytes
+
+
+def do_parse(
+    output_dir,
+    pdf_file_names: list[str],
+    pdf_bytes_list: list[bytes],
+    p_lang_list: list[str],
+    backend="pipeline",
+    parse_method="auto",
+    p_formula_enable=True,
+    p_table_enable=True,
+    server_url=None,
+    f_draw_layout_bbox=True,
+    f_draw_span_bbox=True,
+    f_dump_md=True,
+    f_dump_middle_json=True,
+    f_dump_model_output=True,
+    f_dump_orig_pdf=True,
+    f_dump_content_list=True,
+    f_make_md_mode=MakeMode.MM_MD,
+    start_page_id=0,
+    end_page_id=None,
+):
+
+    if backend == "pipeline":
+
+        from mineru.backend.pipeline.pipeline_middle_json_mkcontent import union_make as pipeline_union_make
+        from mineru.backend.pipeline.model_json_to_middle_json import result_to_middle_json as pipeline_result_to_middle_json
+        from mineru.backend.pipeline.pipeline_analyze import doc_analyze as pipeline_doc_analyze
+
+        for idx, pdf_bytes in enumerate(pdf_bytes_list):
+            new_pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
+            pdf_bytes_list[idx] = new_pdf_bytes
+
+        infer_results, all_image_lists, all_pdf_docs, lang_list, ocr_enabled_list = pipeline_doc_analyze(pdf_bytes_list, p_lang_list, parse_method=parse_method, formula_enable=p_formula_enable,table_enable=p_table_enable)
+
+        for idx, model_list in enumerate(infer_results):
+            model_json = copy.deepcopy(model_list)
+            pdf_file_name = pdf_file_names[idx]
+            local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
+            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
+
+            images_list = all_image_lists[idx]
+            pdf_doc = all_pdf_docs[idx]
+            _lang = lang_list[idx]
+            _ocr_enable = ocr_enabled_list[idx]
+
+            middle_json = pipeline_result_to_middle_json(model_list, images_list, pdf_doc, image_writer, _lang, _ocr_enable, p_formula_enable)
+
+            pdf_info = middle_json["pdf_info"]
+
+            pdf_bytes = pdf_bytes_list[idx]
+            if f_draw_layout_bbox:
+                draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_layout.pdf")
+
+            if f_draw_span_bbox:
+                draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_span.pdf")
+
+            if f_dump_orig_pdf:
+                md_writer.write(
+                    f"{pdf_file_name}_origin.pdf",
+                    pdf_bytes,
+                )
+
+            if f_dump_md:
+                image_dir = str(os.path.basename(local_image_dir))
+                md_content_str = pipeline_union_make(pdf_info, f_make_md_mode, image_dir)
+                md_writer.write_string(
+                    f"{pdf_file_name}.md",
+                    md_content_str,
+                )
+
+            if f_dump_content_list:
+                image_dir = str(os.path.basename(local_image_dir))
+                content_list = pipeline_union_make(pdf_info, MakeMode.CONTENT_LIST, image_dir)
+                md_writer.write_string(
+                    f"{pdf_file_name}_content_list.json",
+                    json.dumps(content_list, ensure_ascii=False, indent=4),
+                )
+
+            if f_dump_middle_json:
+                md_writer.write_string(
+                    f"{pdf_file_name}_middle.json",
+                    json.dumps(middle_json, ensure_ascii=False, indent=4),
+                )
+
+            if f_dump_model_output:
+                md_writer.write_string(
+                    f"{pdf_file_name}_model.json",
+                    json.dumps(model_json, ensure_ascii=False, indent=4),
+                )
+
+            logger.info(f"local output dir is {local_md_dir}")
+    else:
+
+        if backend.startswith("vlm-"):
+            backend = backend[4:]
+
+        f_draw_span_bbox = False
+        parse_method = "vlm"
+        for idx, pdf_bytes in enumerate(pdf_bytes_list):
+            pdf_file_name = pdf_file_names[idx]
+            pdf_bytes = convert_pdf_bytes_to_bytes_by_pypdfium2(pdf_bytes, start_page_id, end_page_id)
+            local_image_dir, local_md_dir = prepare_env(output_dir, pdf_file_name, parse_method)
+            image_writer, md_writer = FileBasedDataWriter(local_image_dir), FileBasedDataWriter(local_md_dir)
+            middle_json, infer_result = vlm_doc_analyze(pdf_bytes, image_writer=image_writer, backend=backend, server_url=server_url)
+
+            pdf_info = middle_json["pdf_info"]
+
+            if f_draw_layout_bbox:
+                draw_layout_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_layout.pdf")
+
+            if f_draw_span_bbox:
+                draw_span_bbox(pdf_info, pdf_bytes, local_md_dir, f"{pdf_file_name}_span.pdf")
+
+            if f_dump_orig_pdf:
+                md_writer.write(
+                    f"{pdf_file_name}_origin.pdf",
+                    pdf_bytes,
+                )
+
+            if f_dump_md:
+                image_dir = str(os.path.basename(local_image_dir))
+                md_content_str = vlm_union_make(pdf_info, f_make_md_mode, image_dir)
+                md_writer.write_string(
+                    f"{pdf_file_name}.md",
+                    md_content_str,
+                )
+
+            if f_dump_content_list:
+                image_dir = str(os.path.basename(local_image_dir))
+                content_list = vlm_union_make(pdf_info, MakeMode.CONTENT_LIST, image_dir)
+                md_writer.write_string(
+                    f"{pdf_file_name}_content_list.json",
+                    json.dumps(content_list, ensure_ascii=False, indent=4),
+                )
+
+            if f_dump_middle_json:
+                md_writer.write_string(
+                    f"{pdf_file_name}_middle.json",
+                    json.dumps(middle_json, ensure_ascii=False, indent=4),
+                )
+
+            if f_dump_model_output:
+                model_output = ("\n" + "-" * 50 + "\n").join(infer_result)
+                md_writer.write_string(
+                    f"{pdf_file_name}_model_output.txt",
+                    model_output,
+                )
+
+            logger.info(f"local output dir is {local_md_dir}")
+
+
+
+if __name__ == "__main__":
+    # pdf_path = "../../demo/pdfs/demo3.pdf"
+    pdf_path = "C:/Users/zhaoxiaomeng/Downloads/4546d0e2-ba60-40a5-a17e-b68555cec741.pdf"
+
+    try:
+       do_parse("./output", [Path(pdf_path).stem], [read_fn(Path(pdf_path))],["ch"],
+                end_page_id=10,
+                backend='vlm-huggingface'
+                # backend = 'pipeline'
+                )
+    except Exception as e:
+        logger.exception(e)

mineru/cli/models_download.py

+import json
+import os
+import sys
+import click
+import requests
+
+from mineru.utils.enum_class import ModelPath
+from mineru.utils.models_download_utils import auto_download_and_get_model_root_path
+
+
+def download_json(url):
+    """下载JSON文件"""
+    response = requests.get(url)
+    response.raise_for_status()
+    return response.json()
+
+
+def download_and_modify_json(url, local_filename, modifications):
+    """下载JSON并修改内容"""
+    if os.path.exists(local_filename):
+        data = json.load(open(local_filename))
+        config_version = data.get('config_version', '0.0.0')
+        if config_version < '1.3.0':
+            data = download_json(url)
+    else:
+        data = download_json(url)
+
+    # 修改内容
+    for key, value in modifications.items():
+        if key in data:
+            if isinstance(data[key], dict):
+                # 如果是字典，合并新值
+                data[key].update(value)
+            else:
+                # 否则直接替换
+                data[key] = value
+
+    # 保存修改后的内容
+    with open(local_filename, 'w', encoding='utf-8') as f:
+        json.dump(data, f, ensure_ascii=False, indent=4)
+
+
+def configure_model(model_dir, model_type):
+    """配置模型"""
+    json_url = 'https://gcore.jsdelivr.net/gh/opendatalab/MinerU@master/mineru.template.json'
+    config_file_name = os.getenv('MINERU_TOOLS_CONFIG_JSON', 'mineru.json')
+    home_dir = os.path.expanduser('~')
+    config_file = os.path.join(home_dir, config_file_name)
+
+    json_mods = {
+        'models-dir': {
+            f'{model_type}': model_dir
+        }
+    }
+
+    download_and_modify_json(json_url, config_file, json_mods)
+    print(f'The configuration file has been successfully configured, the path is: {config_file}')
+
+
+@click.command()
+@click.option(
+    '-s',
+    '--source',
+    'model_source',
+    type=click.Choice(['huggingface', 'modelscope']),
+    help="""
+        The source of the model repository. 
+        """,
+    default=None,
+)
+@click.option(
+    '-m',
+    '--model_type',
+    'model_type',
+    type=click.Choice(['pipeline', 'vlm', 'all']),
+    help="""
+        The type of the model to download.
+        """,
+    default=None,
+)
+def download_models(model_source, model_type):
+    """Download MinerU model files.
+
+    Supports downloading pipeline or VLM models from ModelScope or HuggingFace.
+    """
+    # 如果未显式指定则交互式输入下载来源
+    if model_source is None:
+        model_source = click.prompt(
+            "Please select the model download source: ",
+            type=click.Choice(['huggingface', 'modelscope']),
+            default='huggingface'
+        )
+
+    if os.getenv('MINERU_MODEL_SOURCE', None) is None:
+        os.environ['MINERU_MODEL_SOURCE'] = model_source
+
+    # 如果未显式指定则交互式输入模型类型
+    if model_type is None:
+        model_type = click.prompt(
+            "Please select the model type to download: ",
+            type=click.Choice(['pipeline', 'vlm', 'all']),
+            default='all'
+        )
+
+    click.echo(f"Downloading {model_type} model from {os.getenv('MINERU_MODEL_SOURCE', None)}...")
+
+    def download_pipeline_models():
+        """下载Pipeline模型"""
+        model_paths = [
+            ModelPath.doclayout_yolo,
+            ModelPath.yolo_v8_mfd,
+            ModelPath.unimernet_small,
+            ModelPath.pytorch_paddle,
+            ModelPath.layout_reader,
+            ModelPath.slanet_plus
+        ]
+        download_finish_path = ""
+        for model_path in model_paths:
+            click.echo(f"Downloading model: {model_path}")
+            download_finish_path = auto_download_and_get_model_root_path(model_path, repo_mode='pipeline')
+        click.echo(f"Pipeline models downloaded successfully to: {download_finish_path}")
+        configure_model(download_finish_path, "pipeline")
+
+    def download_vlm_models():
+        """下载VLM模型"""
+        download_finish_path = auto_download_and_get_model_root_path("/", repo_mode='vlm')
+        click.echo(f"VLM models downloaded successfully to: {download_finish_path}")
+        configure_model(download_finish_path, "vlm")
+
+    try:
+        if model_type == 'pipeline':
+            download_pipeline_models()
+        elif model_type == 'vlm':
+            download_vlm_models()
+        elif model_type == 'all':
+            download_pipeline_models()
+            download_vlm_models()
+        else:
+            click.echo(f"Unsupported model type: {model_type}", err=True)
+            sys.exit(1)
+
+    except Exception as e:
+        click.echo(f"Download failed: {str(e)}", err=True)
+        sys.exit(1)
+
+if __name__ == '__main__':
+    download_models()

mineru/cli/vlm_sglang_server.py

+from ..model.vlm_sglang_model.server import main
+
+if __name__ == "__main__":
+    main()

mineru/data/__init__.py

+# Copyright (c) Opendatalab. All rights reserved.

mineru/data/data_reader_writer/__init__.py

+from .base import DataReader, DataWriter
+from .dummy import DummyDataWriter
+from .filebase import FileBasedDataReader, FileBasedDataWriter
+from .multi_bucket_s3 import MultiBucketS3DataReader, MultiBucketS3DataWriter
+from .s3 import S3DataReader, S3DataWriter
+
+__all__ = [
+    "DataReader",
+    "DataWriter",
+    "FileBasedDataReader",
+    "FileBasedDataWriter",
+    "S3DataReader",
+    "S3DataWriter",
+    "MultiBucketS3DataReader",
+    "MultiBucketS3DataWriter",
+    "DummyDataWriter",
+]

mineru/data/data_reader_writer/base.py

+
+from abc import ABC, abstractmethod
+
+
+class DataReader(ABC):
+
+    def read(self, path: str) -> bytes:
+        """Read the file.
+
+        Args:
+            path (str): file path to read
+
+        Returns:
+            bytes: the content of the file
+        """
+        return self.read_at(path)
+
+    @abstractmethod
+    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
+        """Read the file at offset and limit.
+
+        Args:
+            path (str): the file path
+            offset (int, optional): the number of bytes skipped. Defaults to 0.
+            limit (int, optional): the length of bytes want to read. Defaults to -1.
+
+        Returns:
+            bytes: the content of the file
+        """
+        pass
+
+
+class DataWriter(ABC):
+    @abstractmethod
+    def write(self, path: str, data: bytes) -> None:
+        """Write the data to the file.
+
+        Args:
+            path (str): the target file where to write
+            data (bytes): the data want to write
+        """
+        pass
+
+    def write_string(self, path: str, data: str) -> None:
+        """Write the data to file, the data will be encoded to bytes.
+
+        Args:
+            path (str): the target file where to write
+            data (str): the data want to write
+        """
+
+        def safe_encode(data: str, method: str):
+            try:
+                bit_data = data.encode(encoding=method, errors='replace')
+                return bit_data, True
+            except:  # noqa
+                return None, False
+
+        for method in ['utf-8', 'ascii']:
+            bit_data, flag = safe_encode(data, method)
+            if flag:
+                self.write(path, bit_data)
+                break

mineru/data/data_reader_writer/dummy.py

+from .base import DataWriter
+
+
+class DummyDataWriter(DataWriter):
+    def write(self, path: str, data: bytes) -> None:
+        """Dummy write method that does nothing."""
+        pass
+
+    def write_string(self, path: str, data: str) -> None:
+        """Dummy write_string method that does nothing."""
+        pass

mineru/data/data_reader_writer/filebase.py

+import os
+
+from .base import DataReader, DataWriter
+
+
+class FileBasedDataReader(DataReader):
+    def __init__(self, parent_dir: str = ''):
+        """Initialized with parent_dir.
+
+        Args:
+            parent_dir (str, optional): the parent directory that may be used within methods. Defaults to ''.
+        """
+        self._parent_dir = parent_dir
+
+    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
+        """Read at offset and limit.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            offset (int, optional): the number of bytes skipped. Defaults to 0.
+            limit (int, optional): the length of bytes want to read. Defaults to -1.
+
+        Returns:
+            bytes: the content of file
+        """
+        fn_path = path
+        if not os.path.isabs(fn_path) and len(self._parent_dir) > 0:
+            fn_path = os.path.join(self._parent_dir, path)
+
+        with open(fn_path, 'rb') as f:
+            f.seek(offset)
+            if limit == -1:
+                return f.read()
+            else:
+                return f.read(limit)
+
+
+class FileBasedDataWriter(DataWriter):
+    def __init__(self, parent_dir: str = '') -> None:
+        """Initialized with parent_dir.
+
+        Args:
+            parent_dir (str, optional): the parent directory that may be used within methods. Defaults to ''.
+        """
+        self._parent_dir = parent_dir
+
+    def write(self, path: str, data: bytes) -> None:
+        """Write file with data.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            data (bytes): the data want to write
+        """
+        fn_path = path
+        if not os.path.isabs(fn_path) and len(self._parent_dir) > 0:
+            fn_path = os.path.join(self._parent_dir, path)
+
+        if not os.path.exists(os.path.dirname(fn_path)) and os.path.dirname(fn_path) != "":
+            os.makedirs(os.path.dirname(fn_path), exist_ok=True)
+
+        with open(fn_path, 'wb') as f:
+            f.write(data)

mineru/data/data_reader_writer/multi_bucket_s3.py

+
+from ..utils.exceptions import InvalidConfig, InvalidParams
+from .base import DataReader, DataWriter
+from ..io.s3 import S3Reader, S3Writer
+from ..utils.schemas import S3Config
+from ..utils.path_utils import parse_s3_range_params, parse_s3path, remove_non_official_s3_args
+
+
+class MultiS3Mixin:
+    def __init__(self, default_prefix: str, s3_configs: list[S3Config]):
+        """Initialized with multiple s3 configs.
+
+        Args:
+            default_prefix (str): the default prefix of the relative path. for example, {some_bucket}/{some_prefix} or {some_bucket}
+            s3_configs (list[S3Config]): list of s3 configs, the bucket_name must be unique in the list.
+
+        Raises:
+            InvalidConfig: default bucket config not in s3_configs.
+            InvalidConfig: bucket name not unique in s3_configs.
+            InvalidConfig: default bucket must be provided.
+        """
+        if len(default_prefix) == 0:
+            raise InvalidConfig('default_prefix must be provided')
+
+        arr = default_prefix.strip('/').split('/')
+        self.default_bucket = arr[0]
+        self.default_prefix = '/'.join(arr[1:])
+
+        found_default_bucket_config = False
+        for conf in s3_configs:
+            if conf.bucket_name == self.default_bucket:
+                found_default_bucket_config = True
+                break
+
+        if not found_default_bucket_config:
+            raise InvalidConfig(
+                f'default_bucket: {self.default_bucket} config must be provided in s3_configs: {s3_configs}'
+            )
+
+        uniq_bucket = set([conf.bucket_name for conf in s3_configs])
+        if len(uniq_bucket) != len(s3_configs):
+            raise InvalidConfig(
+                f'the bucket_name in s3_configs: {s3_configs} must be unique'
+            )
+
+        self.s3_configs = s3_configs
+        self._s3_clients_h: dict = {}
+
+
+class MultiBucketS3DataReader(DataReader, MultiS3Mixin):
+    def read(self, path: str) -> bytes:
+        """Read the path from s3, select diffect bucket client for each request
+        based on the bucket, also support range read.
+
+        Args:
+            path (str): the s3 path of file, the path must be in the format of s3://bucket_name/path?offset,limit.
+            for example: s3://bucket_name/path?0,100.
+
+        Returns:
+            bytes: the content of s3 file.
+        """
+        may_range_params = parse_s3_range_params(path)
+        if may_range_params is None or 2 != len(may_range_params):
+            byte_start, byte_len = 0, -1
+        else:
+            byte_start, byte_len = int(may_range_params[0]), int(may_range_params[1])
+        path = remove_non_official_s3_args(path)
+        return self.read_at(path, byte_start, byte_len)
+
+    def __get_s3_client(self, bucket_name: str):
+        if bucket_name not in set([conf.bucket_name for conf in self.s3_configs]):
+            raise InvalidParams(
+                f'bucket name: {bucket_name} not found in s3_configs: {self.s3_configs}'
+            )
+        if bucket_name not in self._s3_clients_h:
+            conf = next(
+                filter(lambda conf: conf.bucket_name == bucket_name, self.s3_configs)
+            )
+            self._s3_clients_h[bucket_name] = S3Reader(
+                bucket_name,
+                conf.access_key,
+                conf.secret_key,
+                conf.endpoint_url,
+                conf.addressing_style,
+            )
+        return self._s3_clients_h[bucket_name]
+
+    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
+        """Read the file with offset and limit, select diffect bucket client
+        for each request based on the bucket.
+
+        Args:
+            path (str): the file path.
+            offset (int, optional): the number of bytes skipped. Defaults to 0.
+            limit (int, optional): the number of bytes want to read. Defaults to -1 which means infinite.
+
+        Returns:
+            bytes: the file content.
+        """
+        if path.startswith('s3://'):
+            bucket_name, path = parse_s3path(path)
+            s3_reader = self.__get_s3_client(bucket_name)
+        else:
+            s3_reader = self.__get_s3_client(self.default_bucket)
+            if self.default_prefix:
+                path = self.default_prefix + '/' + path
+        return s3_reader.read_at(path, offset, limit)
+
+
+class MultiBucketS3DataWriter(DataWriter, MultiS3Mixin):
+    def __get_s3_client(self, bucket_name: str):
+        if bucket_name not in set([conf.bucket_name for conf in self.s3_configs]):
+            raise InvalidParams(
+                f'bucket name: {bucket_name} not found in s3_configs: {self.s3_configs}'
+            )
+        if bucket_name not in self._s3_clients_h:
+            conf = next(
+                filter(lambda conf: conf.bucket_name == bucket_name, self.s3_configs)
+            )
+            self._s3_clients_h[bucket_name] = S3Writer(
+                bucket_name,
+                conf.access_key,
+                conf.secret_key,
+                conf.endpoint_url,
+                conf.addressing_style,
+            )
+        return self._s3_clients_h[bucket_name]
+
+    def write(self, path: str, data: bytes) -> None:
+        """Write file with data, also select diffect bucket client for each
+        request based on the bucket.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            data (bytes): the data want to write.
+        """
+        if path.startswith('s3://'):
+            bucket_name, path = parse_s3path(path)
+            s3_writer = self.__get_s3_client(bucket_name)
+        else:
+            s3_writer = self.__get_s3_client(self.default_bucket)
+            if self.default_prefix:
+                path = self.default_prefix + '/' + path
+        return s3_writer.write(path, data)

mineru/data/data_reader_writer/s3.py

+from .multi_bucket_s3 import MultiBucketS3DataReader, MultiBucketS3DataWriter
+from ..utils.schemas import S3Config
+
+
+class S3DataReader(MultiBucketS3DataReader):
+    def __init__(
+        self,
+        default_prefix_without_bucket: str,
+        bucket: str,
+        ak: str,
+        sk: str,
+        endpoint_url: str,
+        addressing_style: str = 'auto',
+    ):
+        """s3 reader client.
+
+        Args:
+            default_prefix_without_bucket: prefix that not contains bucket
+            bucket (str): bucket name
+            ak (str): access key
+            sk (str): secret key
+            endpoint_url (str): endpoint url of s3
+            addressing_style (str, optional): Defaults to 'auto'. Other valid options here are 'path' and 'virtual'
+            refer to https://boto3.amazonaws.com/v1/documentation/api/1.9.42/guide/s3.html
+        """
+        super().__init__(
+            f'{bucket}/{default_prefix_without_bucket}',
+            [
+                S3Config(
+                    bucket_name=bucket,
+                    access_key=ak,
+                    secret_key=sk,
+                    endpoint_url=endpoint_url,
+                    addressing_style=addressing_style,
+                )
+            ],
+        )
+
+
+class S3DataWriter(MultiBucketS3DataWriter):
+    def __init__(
+        self,
+        default_prefix_without_bucket: str,
+        bucket: str,
+        ak: str,
+        sk: str,
+        endpoint_url: str,
+        addressing_style: str = 'auto',
+    ):
+        """s3 writer client.
+
+        Args:
+            default_prefix_without_bucket: prefix that not contains bucket
+            bucket (str): bucket name
+            ak (str): access key
+            sk (str): secret key
+            endpoint_url (str): endpoint url of s3
+            addressing_style (str, optional): Defaults to 'auto'. Other valid options here are 'path' and 'virtual'
+            refer to https://boto3.amazonaws.com/v1/documentation/api/1.9.42/guide/s3.html
+        """
+        super().__init__(
+            f'{bucket}/{default_prefix_without_bucket}',
+            [
+                S3Config(
+                    bucket_name=bucket,
+                    access_key=ak,
+                    secret_key=sk,
+                    endpoint_url=endpoint_url,
+                    addressing_style=addressing_style,
+                )
+            ],
+        )

mineru/data/io/__init__.py

+
+from .base import IOReader, IOWriter
+from .http import HttpReader, HttpWriter
+from .s3 import S3Reader, S3Writer
+
+__all__ = ['IOReader', 'IOWriter', 'HttpReader', 'HttpWriter', 'S3Reader', 'S3Writer']
\ No newline at end of file

mineru/data/io/base.py

+from abc import ABC, abstractmethod
+
+
+class IOReader(ABC):
+    @abstractmethod
+    def read(self, path: str) -> bytes:
+        """Read the file.
+
+        Args:
+            path (str): file path to read
+
+        Returns:
+            bytes: the content of the file
+        """
+        pass
+
+    @abstractmethod
+    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
+        """Read at offset and limit.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            offset (int, optional): the number of bytes skipped. Defaults to 0.
+            limit (int, optional): the length of bytes want to read. Defaults to -1.
+
+        Returns:
+            bytes: the content of file
+        """
+        pass
+
+
+class IOWriter(ABC):
+
+    @abstractmethod
+    def write(self, path: str, data: bytes) -> None:
+        """Write file with data.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            data (bytes): the data want to write
+        """
+        pass

mineru/data/io/http.py

+
+import io
+
+import requests
+
+from .base import IOReader, IOWriter
+
+
+class HttpReader(IOReader):
+
+    def read(self, url: str) -> bytes:
+        """Read the file.
+
+        Args:
+            path (str): file path to read
+
+        Returns:
+            bytes: the content of the file
+        """
+        return requests.get(url).content
+
+    def read_at(self, path: str, offset: int = 0, limit: int = -1) -> bytes:
+        """Not Implemented."""
+        raise NotImplementedError
+
+
+class HttpWriter(IOWriter):
+    def write(self, url: str, data: bytes) -> None:
+        """Write file with data.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            data (bytes): the data want to write
+        """
+        files = {'file': io.BytesIO(data)}
+        response = requests.post(url, files=files)
+        assert 300 > response.status_code and response.status_code > 199

mineru/data/io/s3.py

+import boto3
+from botocore.config import Config
+
+from ..io.base import IOReader, IOWriter
+
+
+class S3Reader(IOReader):
+    def __init__(
+        self,
+        bucket: str,
+        ak: str,
+        sk: str,
+        endpoint_url: str,
+        addressing_style: str = 'auto',
+    ):
+        """s3 reader client.
+
+        Args:
+            bucket (str): bucket name
+            ak (str): access key
+            sk (str): secret key
+            endpoint_url (str): endpoint url of s3
+            addressing_style (str, optional): Defaults to 'auto'. Other valid options here are 'path' and 'virtual'
+            refer to https://boto3.amazonaws.com/v1/documentation/api/1.9.42/guide/s3.html
+        """
+        self._bucket = bucket
+        self._ak = ak
+        self._sk = sk
+        self._s3_client = boto3.client(
+            service_name='s3',
+            aws_access_key_id=ak,
+            aws_secret_access_key=sk,
+            endpoint_url=endpoint_url,
+            config=Config(
+                s3={'addressing_style': addressing_style},
+                retries={'max_attempts': 5, 'mode': 'standard'},
+            ),
+        )
+
+    def read(self, key: str) -> bytes:
+        """Read the file.
+
+        Args:
+            path (str): file path to read
+
+        Returns:
+            bytes: the content of the file
+        """
+        return self.read_at(key)
+
+    def read_at(self, key: str, offset: int = 0, limit: int = -1) -> bytes:
+        """Read at offset and limit.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            offset (int, optional): the number of bytes skipped. Defaults to 0.
+            limit (int, optional): the length of bytes want to read. Defaults to -1.
+
+        Returns:
+            bytes: the content of file
+        """
+        if limit > -1:
+            range_header = f'bytes={offset}-{offset+limit-1}'
+            res = self._s3_client.get_object(
+                Bucket=self._bucket, Key=key, Range=range_header
+            )
+        else:
+            res = self._s3_client.get_object(
+                Bucket=self._bucket, Key=key, Range=f'bytes={offset}-'
+            )
+        return res['Body'].read()
+
+
+class S3Writer(IOWriter):
+    def __init__(
+        self,
+        bucket: str,
+        ak: str,
+        sk: str,
+        endpoint_url: str,
+        addressing_style: str = 'auto',
+    ):
+        """s3 reader client.
+
+        Args:
+            bucket (str): bucket name
+            ak (str): access key
+            sk (str): secret key
+            endpoint_url (str): endpoint url of s3
+            addressing_style (str, optional): Defaults to 'auto'. Other valid options here are 'path' and 'virtual'
+            refer to https://boto3.amazonaws.com/v1/documentation/api/1.9.42/guide/s3.html
+        """
+        self._bucket = bucket
+        self._ak = ak
+        self._sk = sk
+        self._s3_client = boto3.client(
+            service_name='s3',
+            aws_access_key_id=ak,
+            aws_secret_access_key=sk,
+            endpoint_url=endpoint_url,
+            config=Config(
+                s3={'addressing_style': addressing_style},
+                retries={'max_attempts': 5, 'mode': 'standard'},
+            ),
+        )
+
+    def write(self, key: str, data: bytes):
+        """Write file with data.
+
+        Args:
+            path (str): the path of file, if the path is relative path, it will be joined with parent_dir.
+            data (bytes): the data want to write
+        """
+        self._s3_client.put_object(Bucket=self._bucket, Key=key, Body=data)

mineru/data/utils/__init__.py

+# Copyright (c) Opendatalab. All rights reserved.

mineru/data/utils/exceptions.py

+# Copyright (c) Opendatalab. All rights reserved.
+
+class FileNotExisted(Exception):
+
+    def __init__(self, path):
+        self.path = path
+
+    def __str__(self):
+        return f'File {self.path} does not exist.'
+
+
+class InvalidConfig(Exception):
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __str__(self):
+        return f'Invalid config: {self.msg}'
+
+
+class InvalidParams(Exception):
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __str__(self):
+        return f'Invalid params: {self.msg}'
+
+
+class EmptyData(Exception):
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __str__(self):
+        return f'Empty data: {self.msg}'
+
+class CUDA_NOT_AVAILABLE(Exception):
+    def __init__(self, msg):
+        self.msg = msg
+
+    def __str__(self):
+        return f'CUDA not available: {self.msg}'
\ No newline at end of file

mineru/data/utils/path_utils.py

+# Copyright (c) Opendatalab. All rights reserved.
+
+
+def remove_non_official_s3_args(s3path):
+    """
+    example: s3://abc/xxxx.json?bytes=0,81350 ==> s3://abc/xxxx.json
+    """
+    arr = s3path.split("?")
+    return arr[0]
+
+def parse_s3path(s3path: str):
+    # from s3pathlib import S3Path
+    # p = S3Path(remove_non_official_s3_args(s3path))
+    # return p.bucket, p.key
+    s3path = remove_non_official_s3_args(s3path).strip()
+    if s3path.startswith(('s3://', 's3a://')):
+        prefix, path = s3path.split('://', 1)
+        bucket_name, key = path.split('/', 1)
+        return bucket_name, key
+    elif s3path.startswith('/'):
+        raise ValueError("The provided path starts with '/'. This does not conform to a valid S3 path format.")
+    else:
+        raise ValueError("Invalid S3 path format. Expected 's3://bucket-name/key' or 's3a://bucket-name/key'.")
+
+
+def parse_s3_range_params(s3path: str):
+    """
+    example: s3://abc/xxxx.json?bytes=0,81350 ==> [0, 81350]
+    """
+    arr = s3path.split("?bytes=")
+    if len(arr) == 1:
+        return None
+    return arr[1].split(",")

mineru/data/utils/schemas.py

+# Copyright (c) Opendatalab. All rights reserved.
+
+from pydantic import BaseModel, Field
+
+
+class S3Config(BaseModel):
+    """S3 config
+    """
+    bucket_name: str = Field(description='s3 bucket name', min_length=1)
+    access_key: str = Field(description='s3 access key', min_length=1)
+    secret_key: str = Field(description='s3 secret key', min_length=1)
+    endpoint_url: str = Field(description='s3 endpoint url', min_length=1)
+    addressing_style: str = Field(description='s3 addressing style', default='auto', min_length=1)
+
+
+class PageInfo(BaseModel):
+    """The width and height of page
+    """
+    w: float = Field(description='the width of page')
+    h: float = Field(description='the height of page')

mineru/model/__init__.py

+# Copyright (c) Opendatalab. All rights reserved.