返回
如何使用Hugging Face Hub批量下载模型
日志
2024-01-09 22:13:51
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函数来轻松完成这项任务。