服务器上更快地下载huggingface上的大模型

最新推荐文章于 2024-09-23 09:46:49 发布

小顾开心编程

最新推荐文章于 2024-09-23 09:46:49 发布

阅读量672

点赞数

文章标签：服务器

本文链接：https://blog.csdn.net/m0_59324917/article/details/130808369

版权

该代码示例展示了如何在服务器上使用HuggingFaceHub库，通过多线程和断点续传技术加快下载大模型的过程。首先，它获取指定模型库的信息，然后列出所有文件并构建wget命令进行断点下载。每个下载任务在一个独立的线程中执行，从而提高下载速度。

摘要由CSDN通过智能技术生成

服务器上更快地下载huggingface上的大模型。

只需要修改hf库名称和本地存储路径即可。

import datetime
import os
import threading

from huggingface_hub import hf_hub_url
from huggingface_hub.hf_api import HfApi
from huggingface_hub.utils import filter_repo_objects

# 执行命令
def execCmd(cmd):
    print("命令%s开始运行%s" % (cmd, datetime.datetime.now()))
    os.system(cmd)
    print("命令%s结束运行%s" % (cmd, datetime.datetime.now()))


if __name__ == '__main__':
    # 需下载的hf库名称
    repo_id = "Salesforce/blip2-opt-2.7b"
    # 本地存储路径
    save_path = './blip2-opt-2.7b'
    
    # 获取项目信息
    _api = HfApi()
    repo_info = _api.repo_info(
        repo_id=repo_id,
        repo_type="model",
        revision='main',
        token=None,
    )

    # 获取文件信息
    filtered_repo_files = list(
        filter_repo_objects(
            items=[f.rfilename for f in repo_info.siblings],
            allow_patterns=None,
            ignore_patterns=None,
        )
    )

    cmds = []
    threads = []

    # 需要执行的命令列表
    for file in filtered_repo_files:
        # 获取路径
        url = hf_hub_url(repo_id=repo_id, filename=file)
        # 断点下载指令
        cmds.append(f'wget -c {url} -P {save_path}')
    print(cmds)

    print("程序开始%s" % datetime.datetime.now())
    for cmd in cmds:
        th = threading.Thread(target=execCmd, args=(cmd,))
        th.start()
        threads.append(th)
    for th in threads:
        th.join()
    print("程序结束%s" % datetime.datetime.now())