download_models.py

"""
Download MinerU pipeline models from Hugging Face Hub.

Called once during Docker build: python download_models.py

Optimizations vs original:
  - Skip-if-exists: if models are already present (Docker layer cache reuse),
    the entire download is skipped without re-downloading anything.
  - MFR excluded: formula recognition models (unimernet, ~1-2 GB) are not
    downloaded because formula recognition is disabled in the config. Excluding
    them cuts download time and image size significantly.
  - layoutreader optional: if download fails (network issue, repo unavailable),
    the script logs a warning and continues. MinerU falls back to its built-in
    layout ordering without layoutreader.

── FORENSIC: OCR MODEL AVAILABILITY IN opendatalab/PDF-Extract-Kit-1.0 ───────

models_config.yml in magic-pdf 1.3.12 was patched in Dockerfile Layer 3.5 to
use the weight files that ARE present in the HF repo. After the patch:

  DEFAULT CPU path (ch_lite):
    det: ch_PP-OCRv5_det_infer.pth      ← present in HF repo ✓
    rec: ch_PP-OCRv5_rec_infer.pth      ← present in HF repo ✓

  en / latin det (patched):
    det: Multilingual_PP-OCRv3_det_infer.pth  ← present in HF repo ✓

  NOT IN REPO (these files do NOT exist and are not needed for default usage):
    ch_PP-OCRv3_det_infer.pth           ← absent; patched to v5
    en_PP-OCRv3_det_infer.pth           ← absent; patched to multilingual
    en_PP-OCRv4_rec_infer.pth           ← absent; en rec unavailable
    korean/japan/chinese_cht/etc v3 rec ← absent; those langs need explicit
                                           lang= arg which is not default

  BUNDLED IN WHEEL (not downloaded here):
    magic_pdf/resources/slanet_plus/slanet-plus.onnx  ← table model, in wheel
    magic_pdf/resources/fasttext-langdetect/lid.176.ftz
    magic_pdf/resources/yolov11-langdetect/yolo_v11_ft.pt

Models saved to /app/models/
Config written to /root/magic-pdf.json
"""

import json
import os
import sys

MODELS_DIR = "/app/models"
EXTRACT_KIT_DIR = os.path.join(MODELS_DIR, "PDF-Extract-Kit-1.0")
LAYOUTREADER_DIR = os.path.join(MODELS_DIR, "layoutreader")

# Canary files: both must exist for the skip-if-exists check to pass.
# Using the OCR det weight (not just the Layout dir) ensures a stale cache
# that predates the v3→v5 model rename cannot produce a false positive.
_CANARY_DIR  = os.path.join(EXTRACT_KIT_DIR, "models", "Layout")
_CANARY_FILE = os.path.join(
    EXTRACT_KIT_DIR, "models", "OCR", "paddleocr_torch",
    "ch_PP-OCRv5_det_infer.pth"           # patched det; absent in v3-era cache
)


def _models_present() -> bool:
    """Return True only when BOTH the Layout directory AND the OCR v5 det weight
    exist.  This prevents stale Docker layer caches (built before the v3→v5 patch)
    from reporting models as present when the required file is missing."""
    return os.path.isdir(_CANARY_DIR) and os.path.isfile(_CANARY_FILE)


def _write_config(layoutreader_dir: str) -> None:
    config = {
        "bucket_info": {},
        "models-dir": os.path.join(EXTRACT_KIT_DIR, "models"),
        "layoutreader-model-dir": layoutreader_dir,
        "device-mode": "cpu",
        "layout-config": {
            "model": "doclayout_yolo"
        },
        "formula-config": {
            "mfd_model": "yolo_v8_mfd",
            "mfr_model": "unimernet_small",
            "enable": False
        },
        "table-config": {
            "model": "rapid_table",
            "enable": True,
            "max_time": 400
        },
        "backend": "pipeline"
    }
    config_path = os.path.expanduser("~/magic-pdf.json")
    with open(config_path, "w") as f:
        json.dump(config, f, indent=2)
    print(f"Config written → {config_path}")
    return config_path


def download() -> None:
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("ERROR: huggingface_hub not installed", file=sys.stderr)
        sys.exit(1)

    # ── Skip-if-exists ────────────────────────────────────────────────────────
    if _models_present():
        print("Models already present — skipping download (Docker layer cache).")
        # Config may still need writing if this is a fresh container from cached layer
        lr_dir = LAYOUTREADER_DIR if os.path.isdir(LAYOUTREADER_DIR) else ""
        _write_config(lr_dir)
        return

    os.makedirs(MODELS_DIR, exist_ok=True)

    # ── PDF-Extract-Kit-1.0 ───────────────────────────────────────────────────
    # Excluded via ignore_patterns:
    #   models/MFR — formula recognition (unimernet). Disabled in config.
    #                Saves ~1-2 GB of download and disk.
    #
    # OCR models (models/OCR/paddleocr_torch/) ARE included (not ignored).
    # After the Layer 3.5 patch, the files models_config.yml references match
    # what is actually present in the repo:
    #   ch_PP-OCRv5_det_infer.pth     ← present ✓
    #   ch_PP-OCRv5_rec_infer.pth     ← present ✓
    #   Multilingual_PP-OCRv3_det_infer.pth ← present ✓
    print("=" * 60)
    print("Downloading PDF-Extract-Kit-1.0 ...")
    print("  (MFR/formula-recognition excluded — disabled in config)")
    print("=" * 60)
    snapshot_download(
        repo_id="opendatalab/PDF-Extract-Kit-1.0",
        local_dir=EXTRACT_KIT_DIR,
        ignore_patterns=[
            "*.git*",
            ".gitattributes",
            "models/MFR*",
            "models/MFR/*",
        ],
    )
    print(f"  → {EXTRACT_KIT_DIR}")

    # Verify canary file landed correctly
    if not os.path.isfile(_CANARY_FILE):
        print(
            f"\nERROR: Expected OCR model not found after download:\n"
            f"  {_CANARY_FILE}\n"
            f"The HF repo may have changed its file structure.\n"
            f"Run:  python3 -c \"from huggingface_hub import list_repo_files; "
            f"[print(f) for f in list_repo_files('opendatalab/PDF-Extract-Kit-1.0')]\"\n"
            f"to inspect the current repo contents.",
            file=sys.stderr,
        )
        sys.exit(1)
    print(f"  Canary verified: {os.path.basename(_CANARY_FILE)} ✓")

    # ── layoutreader (optional) ───────────────────────────────────────────────
    # Improves reading-order accuracy. If unavailable, MinerU uses fallback.
    layoutreader_dir = ""
    print("=" * 60)
    print("Downloading layoutreader (optional, improves reading order) ...")
    print("=" * 60)
    try:
        snapshot_download(
            repo_id="hantian/layoutreader",
            local_dir=LAYOUTREADER_DIR,
            ignore_patterns=["*.git*", ".gitattributes"],
        )
        layoutreader_dir = LAYOUTREADER_DIR
        print(f"  → {LAYOUTREADER_DIR}")
    except Exception as exc:
        print(f"  WARNING: layoutreader download failed ({exc})")
        print("  Continuing without layoutreader — MinerU will use fallback ordering.")

    # ── Write config ──────────────────────────────────────────────────────────
    _write_config(layoutreader_dir)

    print("\n✓ Model setup complete.")
    print(f"  models-dir          : {os.path.join(EXTRACT_KIT_DIR, 'models')}")
    print(f"  layoutreader-dir    : {layoutreader_dir or '(not available)'}")
    print(f"  device-mode         : cpu")
    print(f"  formula recognition : disabled (MFR models excluded)")
    print(f"  table recognition   : enabled (slanet-plus.onnx bundled in wheel)")


if __name__ == "__main__":
    download()