返回

如何使用Hugging Face Hub批量下载模型

日志

Hugging Face Hub是一个用于存储和共享自然语言处理模型的平台。它提供了大量的预训练模型,可以帮助您快速开始您的项目。本文将介绍如何使用Hugging Face Hub批量下载模型。

要批量下载模型,您可以使用Hugging Face官方提供的huggingface_hub库。这个库提供了一些方法,可以帮助您创建、删除、更新和检索repo的信息,也可以从repo中下载文件或将它们集成到您的库中。

例如,您可以使用以下代码来下载一个模型:

from huggingface_hub import snapshot_download

MODEL_URL = "https://huggingface.co/bert-base-uncased"
SNAPSHOT_PATH = snapshot_download(MODEL_URL)

from transformers import AutoModel

model = AutoModel.from_pretrained(SNAPSHOT_PATH)

上面的代码将从Hugging Face Hub下载bert-base-uncased模型,并使用transformers库中的AutoModel类加载它。

当您使用snapshot_download函数下载模型时,它会将模型文件下载到一个临时目录中。您可以通过检查snapshot_download函数的返回值来找到这个目录的位置。例如:

from huggingface_hub import snapshot_download

MODEL_URL = "THUDM/chatglm-6b"
SNAPSHOT_PATH = snapshot_download(MODEL_URL)

print(f"Downloaded files are located in: {SNAPSHOT_PATH}")

上面的代码将打印出模型文件所在的目录。您可以在那个目录中找到下载的文件。

此外,您还可以使用snapshot_download函数的cache_dir参数来指定下载位置。例如,如果您想要将模型文件下载到当前目录下的my_model文件夹中,您可以这样调用snapshot_download函数:

from huggingface_hub import snapshot_download

MODEL_URL = "THUDM/chatglm-6b"
SNAPSHOT_PATH = snapshot_download(MODEL_URL, cache_dir="./my_model")

print(f"Downloaded files are located in: {SNAPSHOT_PATH}")

上面的代码将把模型文件下载到当前目录下的my_model文件夹中。

Hugging Face Hub提供了一种简单快捷的方式来批量下载模型。您可以使用huggingface_hub库中的snapshot_download函数来轻松完成这项任务。