download_models.py

import os
import sys

def download():
    print("--- STARTING MODEL PRE-CACHE (file-only download) ---")

    # Use snapshot_download to pull files WITHOUT loading model into RAM.
    # Loading BGEM3FlagModel() during build causes OOMKilled on HF Spaces
    # because the build container has very limited memory.
    try:
        from huggingface_hub import snapshot_download
    except ImportError:
        print("huggingface_hub not available, skipping pre-cache")
        return

    # 1. BGE-M3 — download files only, no model instantiation
    model_name = "BAAI/bge-m3"
    print(f"Downloading files for {model_name}...")
    try:
        snapshot_download(
            repo_id=model_name,
            repo_type="model",
            ignore_patterns=["*.msgpack", "*.h5", "flax_model*", "tf_model*", "rust_model*"],
        )
        print(f"✓ Files cached: {model_name}")
    except Exception as e:
        print(f"Warning: could not pre-cache {model_name}: {e}")

    # 2. Reranker — download files only
    reranker_name = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
    print(f"Downloading files for {reranker_name}...")
    try:
        snapshot_download(
            repo_id=reranker_name,
            repo_type="model",
            ignore_patterns=["*.msgpack", "*.h5", "flax_model*", "tf_model*"],
        )
        print(f"✓ Files cached: {reranker_name}")
    except Exception as e:
        print(f"Warning: could not pre-cache {reranker_name}: {e}")

    print("--- PRE-CACHE COMPLETE ---")

if __name__ == "__main__":
    download()