fix: use snapshot_download to avoid OOMKilled during model pre-cache build step

download_models.py

@@ -1,38 +1,43 @@
 import os
 import sys
 
-# Monkeypatch for transformers/FlagEmbedding compatibility issue
-try:
-    import transformers.utils.import_utils
-    if not hasattr(transformers.utils.import_utils, 'is_torch_fx_available'):
-        transformers.utils.import_utils.is_torch_fx_available = lambda: False
-except Exception:
-    pass
+def download():
+    print("--- STARTING MODEL PRE-CACHE (file-only download) ---")
 
-from FlagEmbedding import BGEM3FlagModel
-from sentence_transformers import CrossEncoder
+    # 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
 
-def download():
-    print("--- STARTING MODEL PRE-CACHE ---")
-    
-    # 1. BGE-M3
+    # 1. BGE-M3 — download files only, no model instantiation
     model_name = "BAAI/bge-m3"
-    print(f"Downloading/Loading {model_name}...")
+    print(f"Downloading files for {model_name}...")
     try:
-        # This will trigger the download if not present
-        _ = BGEM3FlagModel(model_name, use_fp16=True)
-        print(f"Successfully cached {model_name}")
+        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"Error caching {model_name}: {e}")
+        print(f"Warning: could not pre-cache {model_name}: {e}")
 
-    # 2. Reranker
+    # 2. Reranker — download files only
     reranker_name = "cross-encoder/ms-marco-TinyBERT-L-2-v2"
-    print(f"Downloading/Loading {reranker_name}...")
+    print(f"Downloading files for {reranker_name}...")
     try:
-        _ = CrossEncoder(reranker_name)
-        print(f"Successfully cached {reranker_name}")
+        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"Error caching {reranker_name}: {e}")
+        print(f"Warning: could not pre-cache {reranker_name}: {e}")
 
     print("--- PRE-CACHE COMPLETE ---")