Hub Python 库文档

将任何 ML 框架与 Hub 集成

Hugging Face's logo
加入 Hugging Face 社区

并获得增强的文档体验

开始使用

将任何 ML 框架与 Hub 集成

Hugging Face Hub 使托管和与社区共享模型变得容易。它支持开源生态系统中的数十个库。我们一直在努力扩展这种支持,以推动协作机器学习向前发展。huggingface_hub 库在此过程中起着关键作用,允许任何 Python 脚本轻松推送和加载文件。

将库与 Hub 集成主要有四种方式

  1. 推送到 Hub: 实现一种将模型上传到 Hub 的方法。这包括模型权重,以及模型卡片以及运行模型所需的任何其他相关信息或数据(例如,训练日志)。此方法通常称为 push_to_hub()
  2. 从 Hub 下载: 实现一种从 Hub 加载模型的方法。该方法应下载模型配置/权重并加载模型。此方法通常称为 from_pretrainedload_from_hub()
  3. Inference API: 使用我们的服务器免费对您的库支持的模型运行推理。
  4. Widgets: 在 Hub 上模型的登陆页面上显示一个 widget。它允许用户从浏览器快速尝试模型。

在本指南中,我们将重点关注前两个主题。我们将介绍您可以使用来集成库的两种主要方法,以及它们的优点和缺点。所有内容都总结在本指南的末尾,以帮助您在这两者之间做出选择。请记住,这些只是您可以自由地根据您的要求进行调整的指南。

如果您对 Inference 和 Widgets 感兴趣,您可以关注本指南。在这两种情况下,如果您正在将库与 Hub 集成,并希望在我们的文档中列出,您可以联系我们。

一种灵活的方法:助手函数

将库集成到 Hub 的第一种方法是实际自行实现 push_to_hubfrom_pretrained 方法。这使您可以完全灵活地选择需要上传/下载哪些文件以及如何处理特定于您的框架的输入。您可以参考两个上传文件下载文件指南,以了解有关如何执行此操作的更多信息。例如,FastAI 集成就是这样实现的(请参阅push_to_hub_fastai()from_pretrained_fastai())。

不同库之间的实现可能有所不同,但工作流程通常相似。

from_pretrained

这是一个 from_pretrained 方法通常的样子

def from_pretrained(model_id: str) -> MyModelClass:
   # Download model from Hub
   cached_model = hf_hub_download(
      repo_id=repo_id,
      filename="model.pkl",
      library_name="fastai",
      library_version=get_fastai_version(),
   )

   # Load model
    return load_model(cached_model)

push_to_hub

push_to_hub 方法通常需要更多的复杂性来处理仓库创建、生成模型卡片和保存权重。一种常见的方法是将所有这些文件保存在临时文件夹中,上传它,然后删除它。

def push_to_hub(model: MyModelClass, repo_name: str) -> None:
   api = HfApi()

   # Create repo if not existing yet and get the associated repo_id
   repo_id = api.create_repo(repo_name, exist_ok=True)

   # Save all files in a temporary directory and push them in a single commit
   with TemporaryDirectory() as tmpdir:
      tmpdir = Path(tmpdir)

      # Save weights
      save_model(model, tmpdir / "model.safetensors")

      # Generate model card
      card = generate_model_card(model)
      (tmpdir / "README.md").write_text(card)

      # Save logs
      # Save figures
      # Save evaluation metrics
      # ...

      # Push to hub
      return api.upload_folder(repo_id=repo_id, folder_path=tmpdir)

这当然只是一个例子。如果您对更复杂的操作(删除远程文件、动态上传权重、本地持久化权重等)感兴趣,请参阅上传文件指南。

局限性

虽然这种方法很灵活,但它有一些缺点,尤其是在维护方面。Hugging Face 用户在使用 huggingface_hub 时经常习惯于额外的功能。例如,当从 Hub 加载文件时,通常提供如下参数:

  • token:从私有仓库下载
  • revision:从特定分支下载
  • cache_dir:在特定目录中缓存文件
  • force_download/local_files_only:重用缓存与否
  • proxies:配置 HTTP 会话

推送模型时,也支持类似的参数

  • commit_message:自定义提交消息
  • private:如果缺少,则创建私有仓库
  • create_pr:创建 PR 而不是推送到 main
  • branch:推送到分支而不是 main 分支
  • allow_patterns/ignore_patterns:过滤要上传的文件
  • token

所有这些参数都可以添加到我们上面看到的实现中,并传递给 huggingface_hub 方法。但是,如果参数发生更改或添加了新功能,您将需要更新您的软件包。支持这些参数也意味着您需要维护更多的文档。要了解如何减轻这些限制,让我们跳转到下一节类继承

一种更复杂的方法:类继承

正如我们上面看到的,将库包含到 Hub 中主要有两种方法:上传文件 (push_to_hub) 和下载文件 (from_pretrained)。您可以自行实现这些方法,但会带来一些问题。为了解决这个问题,huggingface_hub 提供了一个使用类继承的工具。让我们看看它是如何工作的!

在很多情况下,库已经使用 Python 类实现了其模型。该类包含模型的属性以及加载、运行、训练和评估模型的方法。我们的方法是扩展此类,以使用 mixin 包含上传和下载功能。Mixin 是一个类,旨在通过多重继承使用一组特定功能扩展现有类。huggingface_hub 提供了自己的 mixin,即ModelHubMixin。这里的关键是了解其行为以及如何自定义它。

ModelHubMixin 类实现了 3 个公共方法 (push_to_hubsave_pretrainedfrom_pretrained)。这些是您的用户将调用的用于使用您的库加载/保存模型的方法。ModelHubMixin 还定义了 2 个私有方法 (_save_pretrained_from_pretrained)。这些是您必须实现的方法。因此,要集成您的库,您应该

  1. 使您的 Model 类继承自 ModelHubMixin
  2. 实现私有方法
    • _save_pretrained():方法将目录路径作为输入,并将模型保存到该目录。您必须编写所有逻辑以在此方法中转储您的模型:模型卡片、模型权重、配置文件、训练日志和图形。此模型的所有相关信息都必须由此方法处理。模型卡片对于描述您的模型尤其重要。有关更多详细信息,请查看我们的实现指南
    • _from_pretrained()类方法,将 model_id 作为输入,并返回实例化的模型。该方法必须下载相关文件并加载它们。
  3. 您已完成!

使用 ModelHubMixin 的优点在于,一旦您处理了文件的序列化/加载,就可以开始了。您无需担心仓库创建、提交、PR 或修订等问题。ModelHubMixin 还确保公共方法已记录并进行类型注释,并且您将能够在 Hub 上查看模型的下载计数。所有这些都由 ModelHubMixin 处理,并可供您的用户使用。

一个具体的例子:PyTorch

我们上面看到的一个很好的例子是 PyTorchModelHubMixin,我们为 PyTorch 框架的集成。这是一个即用型集成。

如何使用它?

以下是任何用户如何从/向 Hub 加载/保存 PyTorch 模型

>>> import torch
>>> import torch.nn as nn
>>> from huggingface_hub import PyTorchModelHubMixin


# Define your Pytorch model exactly the same way you are used to
>>> class MyModel(
...         nn.Module,
...         PyTorchModelHubMixin, # multiple inheritance
...         library_name="keras-nlp",
...         tags=["keras"],
...         repo_url="https://github.com/keras-team/keras-nlp",
...         docs_url="https://keras.org.cn/keras_nlp/",
...         # ^ optional metadata to generate model card
...     ):
...     def __init__(self, hidden_size: int = 512, vocab_size: int = 30000, output_size: int = 4):
...         super().__init__()
...         self.param = nn.Parameter(torch.rand(hidden_size, vocab_size))
...         self.linear = nn.Linear(output_size, vocab_size)

...     def forward(self, x):
...         return self.linear(x + self.param)

# 1. Create model
>>> model = MyModel(hidden_size=128)

# Config is automatically created based on input + default values
>>> model.param.shape[0]
128

# 2. (optional) Save model to local directory
>>> model.save_pretrained("path/to/my-awesome-model")

# 3. Push model weights to the Hub
>>> model.push_to_hub("my-awesome-model")

# 4. Initialize model from the Hub => config has been preserved
>>> model = MyModel.from_pretrained("username/my-awesome-model")
>>> model.param.shape[0]
128

# Model card has been correctly populated
>>> from huggingface_hub import ModelCard
>>> card = ModelCard.load("username/my-awesome-model")
>>> card.data.tags
["keras", "pytorch_model_hub_mixin", "model_hub_mixin"]
>>> card.data.library_name
"keras-nlp"

实现

该实现实际上非常简单,完整的实现可以在此处找到。

  1. 首先,从 ModelHubMixin 继承您的类
from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)
  1. 实现 _save_pretrained 方法
from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)

    def _save_pretrained(self, save_directory: Path) -> None:
        """Save weights from a Pytorch model to a local directory."""
        save_model_as_safetensor(self.module, str(save_directory / SAFETENSORS_SINGLE_FILE))
  1. 实现 _from_pretrained 方法
class PyTorchModelHubMixin(ModelHubMixin):
   (...)

   @classmethod # Must be a classmethod!
   def _from_pretrained(
      cls,
      *,
      model_id: str,
      revision: str,
      cache_dir: str,
      force_download: bool,
      proxies: Optional[Dict],
      resume_download: bool,
      local_files_only: bool,
      token: Union[str, bool, None],
      map_location: str = "cpu", # additional argument
      strict: bool = False, # additional argument
      **model_kwargs,
   ):
      """Load Pytorch pretrained weights and return the loaded model."""
        model = cls(**model_kwargs)
        if os.path.isdir(model_id):
            print("Loading weights from local directory")
            model_file = os.path.join(model_id, SAFETENSORS_SINGLE_FILE)
            return cls._load_as_safetensor(model, model_file, map_location, strict)

         model_file = hf_hub_download(
            repo_id=model_id,
            filename=SAFETENSORS_SINGLE_FILE,
            revision=revision,
            cache_dir=cache_dir,
            force_download=force_download,
            proxies=proxies,
            resume_download=resume_download,
            token=token,
            local_files_only=local_files_only,
            )
         return cls._load_as_safetensor(model, model_file, map_location, strict)

就是这样!您的库现在使用户能够将文件上传和下载到 Hub 以及从 Hub 下载文件。

高级用法

在上面的部分中,我们快速讨论了 ModelHubMixin 的工作原理。在本节中,我们将看到它的一些更高级的功能,以改进您的库与 Hugging Face Hub 的集成。

模型卡片

ModelHubMixin 为您生成模型卡片。模型卡片是随模型一起提供的文件,并提供有关模型的重要信息。在底层,模型卡片是带有附加元数据的简单 Markdown 文件。模型卡片对于可发现性、可重复性和共享至关重要!有关更多详细信息,请查看模型卡片指南

半自动生成模型卡片是确保使用您的库推送的所有模型都共享通用元数据的好方法:library_nametagslicensepipeline_tag 等。这使得 Hub 上支持的所有模型都可以轻松搜索,并为登陆 Hub 的用户提供一些资源链接。您可以直接在从 ModelHubMixin 继承时定义元数据

class UniDepthV1(
   nn.Module,
   PyTorchModelHubMixin,
   library_name="unidepth",
   repo_url="https://github.com/lpiccinelli-eth/UniDepth",
   docs_url=...,
   pipeline_tag="depth-estimation",
   license="cc-by-nc-4.0",
   tags=["monocular-metric-depth-estimation", "arxiv:1234.56789"]
):
   ...

默认情况下,将生成带有您提供的信息的通用模型卡片(示例:pyp1/VoiceCraft_giga830M)。但是您也可以定义自己的模型卡片模板!

在此示例中,所有使用 VoiceCraft 类推送的模型将自动包含引用部分和许可证详细信息。有关如何定义模型卡片模板的更多详细信息,请查看 模型卡片指南

MODEL_CARD_TEMPLATE = """
---
# For reference on model card metadata, see the spec: https://github.com/huggingface/hub-docs/blob/main/modelcard.md?plain=1
# Doc / guide: https://huggingface.co/docs/hub/model-cards
{{ card_data }}
---

This is a VoiceCraft model. For more details, please check out the official Github repo: https://github.com/jasonppy/VoiceCraft. This model is shared under a Attribution-NonCommercial-ShareAlike 4.0 International license.

## Citation

@article{peng2024voicecraft,
  author    = {Peng, Puyuan and Huang, Po-Yao and Li, Daniel and Mohamed, Abdelrahman and Harwath, David},
  title     = {VoiceCraft: Zero-Shot Speech Editing and Text-to-Speech in the Wild},
  journal   = {arXiv},
  year      = {2024},
}
"""

class VoiceCraft(
   nn.Module,
   PyTorchModelHubMixin,
   library_name="voicecraft",
   model_card_template=MODEL_CARD_TEMPLATE,
   ...
):
   ...

最后,如果您想使用动态值扩展模型卡片生成过程,您可以覆盖 generate_model_card() 方法

from huggingface_hub import ModelCard, PyTorchModelHubMixin

class UniDepthV1(nn.Module, PyTorchModelHubMixin, ...):
   (...)

   def generate_model_card(self, *args, **kwargs) -> ModelCard:
      card = super().generate_model_card(*args, **kwargs)
      card.data.metrics = ...  # add metrics to the metadata
      card.text += ... # append section to the modelcard
      return card

配置

ModelHubMixin 为您处理模型配置。它会在您实例化模型时自动检查输入值,并将它们序列化到 config.json 文件中。这提供了 2 个好处

  1. 用户将能够使用与您完全相同的参数重新加载模型。
  2. 拥有 config.json 文件会自动在 Hub 上启用分析(即“下载”计数)。

但它在实践中是如何运作的呢?以下几个规则使该过程从用户角度来看尽可能顺畅

  • 如果您的 __init__ 方法期望 config 输入,它将自动保存在仓库中作为 config.json
  • 如果 config 输入参数使用数据类类型进行注释(例如 config: Optional[MyConfigClass] = None),那么 config 值将被正确地为您反序列化。
  • 在初始化时传递的所有值也将存储在配置文件中。这意味着您不一定需要期望 config 输入才能从中受益。

示例

class MyModel(ModelHubMixin):
   def __init__(value: str, size: int = 3):
      self.value = value
      self.size = size

   (...) # implement _save_pretrained / _from_pretrained

model = MyModel(value="my_value")
model.save_pretrained(...)

# config.json contains passed and default values
{"value": "my_value", "size": 3}

但是,如果某个值无法序列化为 JSON 怎么办?默认情况下,在保存配置文件时,该值将被忽略。但是,在某些情况下,您的库已经期望一个无法序列化的自定义对象作为输入,并且您不想更新您的内部逻辑来更新其类型。不用担心!当从 ModelHubMixin 继承时,您可以为任何类型传递自定义编码器/解码器。这需要更多的工作,但可以确保在将您的库与 Hub 集成时,您的内部逻辑不受影响。

这是一个具体的例子,其中一个类期望 argparse.Namespace 配置作为输入

class VoiceCraft(nn.Module):
    def __init__(self, args):
      self.pattern = self.args.pattern
      self.hidden_size = self.args.hidden_size
      ...

一种解决方案是将 __init__ 签名更新为 def __init__(self, pattern: str, hidden_size: int) 并更新所有实例化您的类的代码片段。这是一种完全有效的修复方法,但可能会破坏使用您的库的下游应用程序。

另一种解决方案是提供一个简单的编码器/解码器,将 argparse.Namespace 转换为字典。

from argparse import Namespace

class VoiceCraft(
   nn.Module,
   PyTorchModelHubMixin,  # inherit from mixin
   coders={
      Namespace : (
         lambda x: vars(x),  # Encoder: how to convert a `Namespace` to a valid jsonable value?
         lambda data: Namespace(**data),  # Decoder: how to reconstruct a `Namespace` from a dictionary?
      )
   }
):
    def __init__(self, args: Namespace): # annotate `args`
      self.pattern = self.args.pattern
      self.hidden_size = self.args.hidden_size
      ...

在上面的代码片段中,类的内部逻辑和 __init__ 签名都没有更改。这意味着您库的所有现有代码片段将继续工作。为了实现这一点,我们必须

  1. 从 mixin 继承(在本例中为 PytorchModelHubMixin)。
  2. 在继承中传递 coders 参数。这是一个字典,其中键是您要处理的自定义类型。值是一个元组 (encoder, decoder)
    • 编码器期望指定类型的对象作为输入,并返回一个可 JSON 化的值。这将在使用 save_pretrained 保存模型时使用。
    • 解码器期望原始数据(通常是字典)作为输入,并重建初始对象。这将在使用 from_pretrained 加载模型时使用。
  3. __init__ 签名添加类型注释。这对于让 mixin 知道类期望的类型,以及因此要使用哪个解码器非常重要。

为了简单起见,上面示例中的编码器/解码器函数并不健壮。对于具体的实现,您很可能必须正确处理 corner cases。

快速比较

让我们快速总结一下我们看到的两种方法的优点和缺点。下表仅供参考。您的框架可能有一些您需要解决的特殊性。本指南仅旨在提供关于如何处理集成的指南和想法。在任何情况下,如果您有任何问题,请随时与我们联系!

集成 使用助手 使用 ModelHubMixin
用户体验 model = load_from_hub(...)
push_to_hub(model, ...)
model = MyModel.from_pretrained(...)
model.push_to_hub(...)
灵活性 非常灵活。
您完全控制实现。
灵活性较低。
您的框架必须有一个模型类。
维护 需要更多维护来添加对配置和新功能的支持。可能还需要修复用户报告的问题。 维护较少,因为与 Hub 的大多数交互都在 huggingface_hub 中实现。
文档 / 类型注解 需要手动编写。 部分由 huggingface_hub 处理。
下载计数器 需要手动处理。 如果类具有 config 属性,则默认启用。
模型卡片 需要手动处理 默认情况下使用 library_name、tags 等生成。
< > 在 GitHub 上更新