diff options
| author | YurenHao0426 <blackhao0426@gmail.com> | 2025-12-17 04:29:37 -0600 |
|---|---|---|
| committer | YurenHao0426 <blackhao0426@gmail.com> | 2025-12-17 04:29:37 -0600 |
| commit | e43b3f8aa36c198b95c1e46bea2eaf3893b13dc3 (patch) | |
| tree | 6ce8a00d2f8b9ebd83c894a27ea01ac50cfb2ff5 /scripts/pull_models.py | |
Diffstat (limited to 'scripts/pull_models.py')
| -rw-r--r-- | scripts/pull_models.py | 76 |
1 files changed, 76 insertions, 0 deletions
diff --git a/scripts/pull_models.py b/scripts/pull_models.py new file mode 100644 index 0000000..1e4abc3 --- /dev/null +++ b/scripts/pull_models.py @@ -0,0 +1,76 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +import sys + +ROOT = Path(__file__).resolve().parents[1] +sys.path.insert(0, str(ROOT / "src")) + +from huggingface_hub import snapshot_download # type: ignore + +from personalization.config.settings import load_local_models_config + + +def pull_one(repo_id: str, dest: Path, force: bool = False) -> None: + dest.mkdir(parents=True, exist_ok=True) + if any(dest.iterdir()) and not force: + print(f"[skip] {dest} already populated. Use --force to overwrite.") + return + print(f"[pull] {repo_id} -> {dest}") + snapshot_download(repo_id=repo_id, local_dir=str(dest), local_dir_use_symlinks=False) + print(f"[done] {repo_id}") + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument( + "--target", + choices=[ + "llm", + "preference_extractor", + "embed_qwen3", + "embed_nemotron", + "embedders", + "reranker_qwen3", + "rerankers", + "all", + ], + default="all", + ) + parser.add_argument("--force", action="store_true") + args = parser.parse_args() + + cfg = load_local_models_config() + if args.target in ("llm", "all"): + pull_one(cfg.llm.hf_id, ROOT / cfg.llm.local_path, force=args.force) + if args.target in ("preference_extractor", "all"): + pull_one( + cfg.preference_extractor.hf_id, + ROOT / cfg.preference_extractor.local_path, + force=args.force, + ) + if args.target in ("embed_qwen3", "embedders", "all") and cfg.embedding and cfg.embedding.qwen3: + pull_one( + cfg.embedding.qwen3.hf_id, + ROOT / cfg.embedding.qwen3.local_path, + force=args.force, + ) + if args.target in ("embed_nemotron", "embedders", "all") and cfg.embedding and cfg.embedding.nemotron: + pull_one( + cfg.embedding.nemotron.hf_id, + ROOT / cfg.embedding.nemotron.local_path, + force=args.force, + ) + if args.target in ("reranker_qwen3", "rerankers", "all") and cfg.reranker and cfg.reranker.qwen3_8b: + pull_one( + cfg.reranker.qwen3_8b.hf_id, + ROOT / cfg.reranker.qwen3_8b.local_path, + force=args.force, + ) + + +if __name__ == "__main__": + main() + + |
