backend/scripts/prefetch_models.py

"""Download the ML models at image-build time so the container is fully
self-contained: no first-request download, fast cold starts, works offline.

Run during `docker build`. Honours HF_HOME / EASYOCR_MODULE_PATH / TORCH_HOME
so the cache lands in a fixed, world-readable location the runtime user shares.
"""
from __future__ import annotations

import os

os.environ.setdefault("HF_HUB_DISABLE_SYMLINKS_WARNING", "1")


def fetch_docling() -> None:
    # Layout + TableFormer models (ds4sd/docling-models).
    from docling.pipeline.standard_pdf_pipeline import StandardPdfPipeline

    path = StandardPdfPipeline.download_models_hf()
    print(f"[prefetch] docling models -> {path}", flush=True)


def fetch_easyocr() -> None:
    # EasyOCR detection + recognition models (used for scanned PDFs / images).
    try:
        import easyocr  # noqa: F401

        easyocr.Reader(["en"], gpu=False)
        print("[prefetch] easyocr (en) models ready", flush=True)
    except Exception as e:  # OCR is optional; never fail the build over it.
        print(f"[prefetch] easyocr skipped: {e}", flush=True)


if __name__ == "__main__":
    fetch_docling()
    fetch_easyocr()
    print("[prefetch] done", flush=True)