# Build from the base olmocr image
FROM alleninstituteforai/olmocr:latest

# Allow specifying which model to include at build time
# Default to the latest olmOCR model
ARG MODEL_NAME=allenai/olmOCR-2-7B-1025-FP8

# Set model cache directory to a fixed location in the image
ENV HF_HOME=/opt/models
ENV TRANSFORMERS_CACHE=/opt/models
ENV OLMOCR_MODEL=${MODEL_NAME}

# Pre-download the olmOCR model into the image
# This adds ~16GB to the image size but eliminates runtime downloads
RUN python -c "from huggingface_hub import snapshot_download; \
    import os; \
    model = os.environ.get('OLMOCR_MODEL'); \
    print(f'Downloading model: {model}'); \
    snapshot_download(model, cache_dir='/opt/models/hub'); \
    print(f'Model {model} successfully downloaded and cached in image')"

# Verify the model is present
RUN python -c "import os; \
    model = os.environ.get('OLMOCR_MODEL').replace('/', '--'); \
    model_path = f'/opt/models/hub/models--{model}'; \
    assert os.path.exists(model_path), f'Model not found at {model_path}'; \
    size = sum(os.path.getsize(os.path.join(dp, f)) for dp, dn, fn in os.walk(model_path) for f in fn) / (1024**3); \
    print(f'Model size: {size:.2f} GB')"

# The entrypoint is already set to /bin/bash in the base image