MinerU/magic_pdf/libs/config_reader.py

"""
根据bucket的名字返回对应的s3 AK， SK，endpoint三元组

"""

import json
import os

from loguru import logger

from magic_pdf.libs.commons import parse_bucket_key

# 定义配置文件名常量
CONFIG_FILE_NAME = "magic-pdf.json"


def read_config():
    home_dir = os.path.expanduser("~")

    config_file = os.path.join(home_dir, CONFIG_FILE_NAME)

    if not os.path.exists(config_file):
        raise FileNotFoundError(f"{config_file} not found")

    with open(config_file, "r", encoding="utf-8") as f:
        config = json.load(f)
    return config


def get_s3_config(bucket_name: str):
    """
    ~/magic-pdf.json 读出来
    """
    config = read_config()

    bucket_info = config.get("bucket_info")
    if bucket_name not in bucket_info:
        access_key, secret_key, storage_endpoint = bucket_info["[default]"]
    else:
        access_key, secret_key, storage_endpoint = bucket_info[bucket_name]

    if access_key is None or secret_key is None or storage_endpoint is None:
        raise Exception(f"ak, sk or endpoint not found in {CONFIG_FILE_NAME}")

    # logger.info(f"get_s3_config: ak={access_key}, sk={secret_key}, endpoint={storage_endpoint}")

    return access_key, secret_key, storage_endpoint


def get_s3_config_dict(path: str):
    access_key, secret_key, storage_endpoint = get_s3_config(get_bucket_name(path))
    return {"ak": access_key, "sk": secret_key, "endpoint": storage_endpoint}


def get_bucket_name(path):
    bucket, key = parse_bucket_key(path)
    return bucket


def get_local_models_dir():
    config = read_config()
    models_dir = config.get("models-dir")
    if models_dir is None:
        logger.warning(f"'models-dir' not found in {CONFIG_FILE_NAME}, use '/tmp/models' as default")
        return "/tmp/models"
    else:
        return models_dir


def get_device():
    config = read_config()
    device = config.get("device-mode")
    if device is None:
        logger.warning(f"'device-mode' not found in {CONFIG_FILE_NAME}, use 'cpu' as default")
        return "cpu"
    else:
        return device


def get_table_recog_config():
    config = read_config()
    table_config = config.get("table-config")
    if table_config is None:
        logger.warning(f"'table-config' not found in {CONFIG_FILE_NAME}, use 'False' as default")
        return json.loads('{"is_table_recog_enable": false, "max_time": 400}')
    else:
        return table_config


if __name__ == "__main__":
    ak, sk, endpoint = get_s3_config("llm-raw")