src/download_models.py

from sentence_transformers import SentenceTransformer, CrossEncoder
import os

CACHE_DIR = os.getenv("HUGGINGFACE_HUB_CACHE", "/app/huggingface_cache")

EMBEDDING_MODEL = "BAAI/bge-m3"
RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L-6-v2"

try:
    SentenceTransformer(EMBEDDING_MODEL, cache_folder=CACHE_DIR)
    print("download completed")
except Exception as e:
    print(f"error to download: {e}")
try:
    print(f"downloading rerank model: {RERANKER_MODEL} for {CACHE_DIR}")
    CrossEncoder(RERANKER_MODEL, cache_folder=CACHE_DIR)
    print("Download do reranker concluído.")
except Exception as e:
    print(f"error to download rerank model: {e}")

print("\ndownload conclude")