From 4197ada6267c822f32c1d7ba2be6e7ce145a3404 Mon Sep 17 00:00:00 2001 From: Lai Zn Date: Fri, 17 Nov 2023 07:13:11 +0800 Subject: [PATCH] feat: enable resume download for hf_hub_download (#1249) --- scripts/setup | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/scripts/setup b/scripts/setup index 377bbe0b3..7c1f44464 100755 --- a/scripts/setup +++ b/scripts/setup @@ -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!")