feat: enable resume download for hf_hub_download (#1249)

This commit is contained in:
Lai Zn 2023-11-17 07:13:11 +08:00 committed by GitHub
parent 09d9a91946
commit 4197ada626
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 9 additions and 0 deletions

View File

@ -1,11 +1,19 @@
#!/usr/bin/env python3
import os
import argparse
from huggingface_hub import hf_hub_download, snapshot_download
from private_gpt.paths import models_path, models_cache_path
from private_gpt.settings.settings import settings
resume_download = True
if __name__ == '__main__':
parser = argparse.ArgumentParser(prog='Setup: Download models from huggingface')
parser.add_argument('--resume', default=True, action=argparse.BooleanOptionalAction, help='Enable/Disable resume_download options to restart the download progress interrupted')
args = parser.parse_args()
resume_download = args.resume
os.makedirs(models_path, exist_ok=True)
embedding_path = models_path / "embedding"
@ -24,6 +32,7 @@ hf_hub_download(
filename=settings().local.llm_hf_model_file,
cache_dir=models_cache_path,
local_dir=models_path,
resume_download=resume_download,
)
print("LLM model downloaded!")