| """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()) |
|
|
