使用snapshot_download配置代理多线程下载模型
snapshot_download
huggingface 官方提供了snapshot_download 方法下载完整模型,参数众多、比较完善。支持断点续传、多线程、指定路径、配置代理、排除特定文件等功能。然而有两个缺点:
- 1))该方法依赖于 transformers 库,而这个库是个开发用的库,对于自动化运维有点重;
- 2) 该方法调用比较复杂,参数较多,例如默认会检查用户缓存目录下是否已有对应模型,如已有则会创建符号链接,不理解的容易导致问题。外加需要配置代理。常见参数配置如下:
因为即使开了代理网络也不太稳定,因此需要搞个死循环,然后批量下载的代码如下
from huggingface_hub import snapshot_download
import sys
from tqdm import tqdm
model_list = ['vikp/layout_segmenter', 'vikp/column_detector', 'facebook/nougat-base', 'vikp/pdf_postprocessor_t5']
while True:
try:
for model_str in tqdm(model_list):
snapshot_download(
repo_id=model_str,
proxies={"http": "http://localhost:7890"},
resume_download=True, # 设置断点续传,版本大于1.0不需要设置
max_workers=8
)
except Exception as e:
# 获取当前异常的信息
exc_type, exc_value, exc_traceback = sys.exc_info()
# 只打印异常类型
print(f"捕获到异常: {exc_type.__name__}")
# time.sleep(5)
else:
print('下载完成')
break
效果大概如下图所示
对于需要登录的模型,还需要两行额外代码:
import huggingface_hub
huggingface_hub.login("HF_TOKEN") # token 从 https://huggingface.co/settings/tokens 获取
速度和效果如下所示
参考内容: