"""Download the models and datasets at the revisions the benchmark used.

  python download.py                      # every model and every dataset
  python download.py bge-m3 Qwen3-Embedding-0.6B
  python download.py --data-only          # datasets only (enough for `eval.py prep`)

Files land under $EVAL_WORKDIR (default /workspace): models/<name>, data/<org>__<name>.
"""
import os
import sys

from huggingface_hub import snapshot_download

W = os.environ.get("EVAL_WORKDIR", "/workspace")

MODELS = {
    "Qwen/Qwen3-Embedding-0.6B": "97b0c614be4d77ee51c0cef4e5f07c00f9eb65b3",
    "Qwen/Qwen3-Embedding-4B": "5cf2132abc99cad020ac570b19d031efec650f2b",
    "Qwen/Qwen3-Reranker-0.6B": "e61197ed45024b0ed8a2d74b80b4d909f1255473",
    "BAAI/bge-m3": "5617a9f61b028005a4858fdac845db406aefb181",
    "BAAI/bge-reranker-v2-m3": "953dc6f6f85a1b2dbfca4c34a2796e7dde08d41e",
    "google/embeddinggemma-2": "914f7f89142e33e77833254d9c9b90c3cef7303b",
}
# Dataset revisions as of 26 October 2026 (all last modified before the benchmark ran).
DATA = {
    "C-MTEB/T2Retrieval": "8731a845f1bf500a4f111cf1070785c793d10e64",
    "C-MTEB/T2Retrieval-qrels": "1c83b8d1544e529875e3f6930f3a1fcf749a8e97",
    "C-MTEB/DuRetrieval": "a1a333e290fe30b10f3f56498e3a0d911a693ced",
    "C-MTEB/DuRetrieval-qrels": "497b7bd1bbb25cb3757ff34d95a8be50a3de2279",
    "mteb/scifact": "cf10ab6856b15b0e670ef8ae5dae4e266c12d035",
    "mteb/nfcorpus": "52ac3f19d3449632d9f00aab0ad34a110fc03816",
}
# bge-m3 ships its weights only as pytorch_model.bin.
PATTERNS = ["*.json", "*.safetensors", "pytorch_model.bin", "*.txt", "*.model", "*.jinja", "*.py", "*/config.json"]

data_only = "--data-only" in sys.argv[1:]
want = [a for a in sys.argv[1:] if a != "--data-only"]

for repo, rev in MODELS.items():
    if data_only or (want and repo.split("/")[1] not in want):
        continue
    p = snapshot_download(repo, revision=rev, local_dir=f"{W}/models/{repo.split('/')[1]}", allow_patterns=PATTERNS)
    print("model", repo, rev, p, flush=True)
for repo, rev in DATA.items():
    p = snapshot_download(repo, repo_type="dataset", revision=rev, local_dir=f"{W}/data/{repo.replace('/', '__')}")
    print("data", repo, rev, p, flush=True)