social-data-attribution-w2/scripts/modal/image.py

"""Modal image definition for SOC-91 WebOrganizer enrichment."""

from pathlib import Path

import modal

from config import ALL_MODELS, MODEL_VOLUME_NAME

_local_path = Path(__file__).resolve()


def download_models() -> None:
    from huggingface_hub import snapshot_download

    for model_name in ALL_MODELS:
        snapshot_download(model_name)


if len(_local_path.parents) > 2:
    _REPO_ROOT = _local_path.parents[2]
    _SRC_DOLMA = str(_REPO_ROOT / "src" / "dolma")
    _CONFIG_PY = str(_REPO_ROOT / "scripts" / "modal" / "config.py")

    _base_image = (
        modal.Image.debian_slim(python_version="3.12")
        .pip_install(
            "torch==2.9.0",
            "transformers>=4.57.1,<5.0.0",
            "zstandard>=0.24.0",
            "pyarrow>=18.0.0",
            "huggingface_hub>=0.25",
            "xformers>=0.0.33.post1",
        )
        .run_commands("python -c 'import torch; print(torch.__version__)'")
        .env({"PYTHONPATH": "/root/src:/root"})
        .add_local_file(_CONFIG_PY, remote_path="/root/config.py", copy=True)
        .run_function(download_models)
        .add_local_dir(_SRC_DOLMA, remote_path="/root/src/dolma")
    )
else:
    _base_image = modal.Image.debian_slim(python_version="3.12")

image_with_models = _base_image

model_volume = modal.Volume.from_name(MODEL_VOLUME_NAME, create_if_missing=True)