download_models.py

import json
import os

import requests
from modelscope import snapshot_download


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):
    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.0.0':
            data = download_json(url)
    else:
        data = download_json(url)

    # 修改内容
    for key, value in modifications.items():
        data[key] = value

    # 保存修改后的内容
    with open(local_filename, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=4)


if __name__ == '__main__':
    mineru_patterns = [
        "models/Layout/LayoutLMv3/*",
        "models/Layout/YOLO/*",
        "models/MFD/YOLO/*",
        "models/MFR/unimernet_small/*",
        "models/TabRec/TableMaster/*",
        "models/TabRec/StructEqTable/*",
    ]
    
    # 在 Docker 环境中使用 /app 目录
    base_dir = "/app"
    cache_dir = os.path.join(base_dir, "model_cache")
    os.makedirs(cache_dir, exist_ok=True)
    
    # 下载模型，使用与原代码相同的参数
    try:
        model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0', 
                                    cache_dir=cache_dir,
                                    revision='master')
        model_dir = os.path.join(model_dir, 'models')
    except TypeError:
        # 如果 allow_patterns 不支持，则下载完整模型
        model_dir = snapshot_download('opendatalab/PDF-Extract-Kit-1.0',
                                    cache_dir=cache_dir)
        model_dir = os.path.join(model_dir, 'models')
    
    layoutreader_model_dir = snapshot_download('ppaanngggg/layoutreader',
                                             cache_dir=cache_dir)
    
    print(f'model_dir is: {model_dir}')
    print(f'layoutreader_model_dir is: {layoutreader_model_dir}')

    # 配置文件保存到 Docker 容器的应用目录
    json_url = 'https://gitee.com/myhloli/MinerU/raw/master/magic-pdf.template.json'
    config_file_name = 'magic-pdf.json'
    config_file = os.path.join(base_dir, config_file_name)

    json_mods = {
        'models-dir': model_dir,
        'layoutreader-model-dir': layoutreader_model_dir,
    }

    download_and_modify_json(json_url, config_file, json_mods)
    print(f'配置文件已成功配置，路径为: {config_file}')