File size: 2,120 Bytes
cede62f | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 | """Download only the two immutable model snapshots needed by the Space."""
from __future__ import annotations
import argparse
from pathlib import Path
from huggingface_hub import snapshot_download
EMBEDDING_MODEL = "intfloat/multilingual-e5-small"
EMBEDDING_REVISION = "fd1525a9fd15316a2d503bf26ab031a61d056e98"
RERANKER_MODEL = "cross-encoder/ms-marco-MiniLM-L2-v2"
RERANKER_REVISION = "1b5cd67b15209f24824c50370e0397743aa9b787"
def download_models(cache_dir: Path) -> tuple[Path, Path]:
cache_dir = cache_dir.resolve()
cache_dir.mkdir(parents=True, exist_ok=True)
embedding = Path(
snapshot_download(
repo_id=EMBEDDING_MODEL,
revision=EMBEDDING_REVISION,
cache_dir=str(cache_dir),
allow_patterns=(
"1_Pooling/config.json",
"config.json",
"model.safetensors",
"modules.json",
"sentence_bert_config.json",
"sentencepiece.bpe.model",
"special_tokens_map.json",
"tokenizer.json",
"tokenizer_config.json",
),
)
)
reranker = Path(
snapshot_download(
repo_id=RERANKER_MODEL,
revision=RERANKER_REVISION,
cache_dir=str(cache_dir),
allow_patterns=(
"config.json",
"model.safetensors",
"openvino/openvino_model_qint8_quantized.bin",
"openvino/openvino_model_qint8_quantized.xml",
"special_tokens_map.json",
"tokenizer.json",
"tokenizer_config.json",
"vocab.txt",
),
)
)
return embedding, reranker
def main() -> int:
parser = argparse.ArgumentParser()
parser.add_argument("--cache-dir", type=Path, required=True)
args = parser.parse_args()
embedding, reranker = download_models(args.cache_dir)
print(f"embedding_snapshot={embedding}")
print(f"reranker_snapshot={reranker}")
return 0
if __name__ == "__main__":
raise SystemExit(main())
|