# 将任意 ML 框架与 Hub 集成

Hugging Face Hub 让托管模型并与社区共享变得轻而易举。它支持开源生态系统中的[数十个库](https://huggingface.co/docs/hub/models-libraries)。我们一直在努力扩展这些支持，以推动协作式机器学习的发展。`huggingface_hub` 库在此过程中发挥着关键作用，它允许任何 Python 脚本轻松推送和加载文件。

将库与 Hub 集成主要有四种方式：
1. **推送到 Hub：**实现一种将模型上传到 Hub 的方法。这包括模型权重，以及[模型卡片](https://huggingface.co/docs/huggingface_hub/how-to-model-cards)和运行模型所需的其他相关信息或数据（例如，训练日志）。该方法通常称为 `push_to_hub()`。
2. **从 Hub 下载：**实现一个从 Hub 加载模型的方法。该方法应下载模型配置/权重并加载模型。此方法通常命名为 `from_pretrained` 或 `load_from_hub()`。
3. **小部件：**在 Hub 上您的模型落地页上显示一个小部件。它允许用户从浏览器中快速试用模型。

在本指南中，我们将重点讨论前两个主题。我们将介绍集成库的两种主要方法，并分析各自的优缺点。指南末尾会对所有内容进行总结，以帮助你在这两种方法之间做出选择。请记住，这些只是指导原则，你可以根据自己的需求自由调整。

如果您对 Inference 和 Widgets 感兴趣，可以参阅[此指南](https://huggingface.co/docs/hub/models-adding-libraries#set-up-the-inference-api)。在这两种情况下，如果您正在将某个库集成到 Hub 中并希望被收录到[我们的文档](https://huggingface.co/docs/hub/models-libraries)中，可以联系我们。

## 一种灵活的方法：辅助函数

The first approach to integrate a library to the Hub is to actually implement the `push_to_hub` and `from_pretrained` methods by yourself. This gives you full flexibility on which files you need to upload/download and how to handle inputs specific to your framework. You can refer to the two [upload files](./upload) and [download files](./download) guides to learn more about how to do that. This is, for example how the FastAI integration is implemented (see [push_to_hub_fastai()](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.push_to_hub_fastai) and [from_pretrained_fastai()](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.from_pretrained_fastai)).

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

### from_pretrained

`from_pretrained` 方法通常如下所示：

```python
def from_pretrained(model_id: str) -> MyModelClass:
    # Download model from Hub
    cached_model = hf_hub_download(
        repo_id=model_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` 方法通常需要处理更多复杂性，例如创建仓库、生成模型卡片以及保存权重。一种常见的做法是将所有这些文件保存到临时文件夹中，上传后再将其删除。

```python
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)
```

这当然只是一个示例。如果您对更复杂的操作（删除远程文件、即时上传权重、在本地持久化权重等）感兴趣，请参阅[上传文件](./upload)指南。

### 局限性

这种方法虽然灵活，但也存在一些缺点，尤其是在维护方面。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](https://stackoverflow.com/a/547714) 是一种旨在通过多重继承为现有类添加一组特定功能的类。`huggingface_hub` 提供了它自己的混入，即 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin)。关键在于理解其行为以及如何对其进行自定义。

[ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 类实现了 3 个 *公开* 方法（`push_to_hub`、`save_pretrained` 和 `from_pretrained`）。这些是用户调用以使用你的库加载/保存模型的方法。[ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 还定义了 2 个 *私有* 方法（`_save_pretrained` 和 `_from_pretrained`）。这些是你必须实现的方法。因此，要集成你的库，你应该：

1. 让您的 Model 类继承自 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin)。
2. 实现私有方法：
    - [_save_pretrained()](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin._save_pretrained)：该方法接收一个目录路径作为输入，并将模型保存至该目录。您必须在此方法中编写所有用于导出模型的逻辑：模型卡片、模型权重、配置文件、训练日志以及图表。与该模型相关的所有信息都必须由该方法处理。[模型卡片](https://huggingface.co/docs/hub/model-cards)对于描述您的模型尤为重要。有关更多详细信息，请参阅[我们的实现指南](./model-cards)。
    - [_from_pretrained()](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin._from_pretrained)：**类方法**，接收 `model_id` 作为输入并返回一个已实例化的模型。该方法必须下载相关文件并将其加载。
3. 完成了！

使用 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 的优势在于，一旦你处理好了文件的序列化/加载，就可以直接上手了。你无需操心诸如创建仓库、提交、PR 或修订版本等事宜。[ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 还会确保公开方法有文档说明和类型注解，并且你可以在 Hub 上查看模型的下载量。所有这些都由 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 处理，并可供你的用户使用。

### 一个具体示例：PyTorch

上述内容的良好示例是 [PyTorchModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin)，即我们针对 PyTorch 框架的集成。这是一个开箱即用的集成。

#### 如何使用？

以下是任何用户如何从 Hub 加载/保存 PyTorch 模型的方法：

```python
>>> 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.io/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"
```

#### 实现

实现实际上非常简单，完整的实现可以在[此处](https://github.com/huggingface/huggingface_hub/blob/main/src/huggingface_hub/hub_mixin.py)找到。

1. 首先，让你的类继承自 `ModelHubMixin`：

```python
from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)
```

2. 实现 `_save_pretrained` 方法：

```py
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))

```

3. 实现 `_from_pretrained` 方法：

```python
class PyTorchModelHubMixin(ModelHubMixin):
   (...)

   @classmethod # Must be a classmethod!
   def _from_pretrained(
      cls,
      *,
      model_id: str,
      revision: str,
      cache_dir: str,
      force_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,
            token=token,
            local_files_only=local_files_only,
            )
         return cls._load_as_safetensor(model, model_file, map_location, strict)
```

就是这样！你的库现在支持用户向 Hub 上传和下载文件。

### 高级用法

在上一节中，我们简要介绍了 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 的工作原理。在本节中，我们将了解它的一些更高级的功能，以改进你的库与 Hugging Face Hub 的集成。

#### 模型卡片

[ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 会为您生成模型卡片。模型卡片是随模型一同提供的文件，包含关于模型的重要信息。在底层，模型卡片是带有额外元数据的简单 Markdown 文件。模型卡片对于可发现性、可复现性和共享至关重要！有关更多详细信息，请参阅 [模型卡片指南](https://huggingface.co/docs/hub/model-cards)。

半自动生成模型卡片是确保通过您的库推送的所有模型共享通用元数据（如 `library_name`、`tags`、`license`、`pipeline_tag` 等）的好方法。这使得基于您库的所有模型都能在 Hub 上轻松搜索，并为访问 Hub 的用户提供一些资源链接。您可以在继承 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 时直接定义元数据：

```py
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](https://huggingface.co/pyp1/VoiceCraft_giga830M)）。但您也可以自定义自己的模型卡片模板！

在此示例中，所有通过 `VoiceCraft` 类推送的模型都会自动包含引用部分和许可证详情。有关如何定义模型卡模板的更多详细信息，请参阅 [Model Cards 指南](./model-cards)。

```py
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()` 方法：

```py
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](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 会为您处理模型配置。在实例化模型时，它会自动检查输入值，并将其序列化到 `config.json` 文件中。这带来了两个好处：
1. 用户将能够使用与您完全相同的参数重新加载模型。
2. 拥有 `config.json` 文件会自动在 Hub 上启用分析功能（即“下载量”计数）。

但在实际使用中它是如何工作的呢？有几条规则可以确保从用户角度来看，整个过程尽可能顺畅：
- 如果你的 `__init__` 方法需要一个 `config` 输入，它会自动以 `config.json` 的形式保存在仓库中。
- 如果 `config` 输入参数被标注为 dataclass 类型（例如 `config: Optional[MyConfigClass] = None`），那么 `config` 的值将被正确反序列化。
- 在初始化时传入的所有值也会存储在配置文件中。这意味着您不一定需要预期一个`config`输入才能从中受益。

示例：

```py
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](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.ModelHubMixin) 时，你可以为任意类型传入自定义的编码器/解码器。这虽然需要多做一些工作，但能确保在将你的库与 Hub 集成时，内部逻辑保持不变。

下面是一个具体示例，其中某个类期望接收一个 `argparse.Namespace` 配置作为输入：

```py
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` 转换为字典。

```py
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. 继承自混入类（此处为 `PytorchModelHubMixin`）。
2. 在继承中传入 `coders` 参数。这是一个字典，其中键是您想要处理的自定义类型，值是一个元组 `(encoder, decoder)`。
   - 编码器期望接收指定类型的对象作为输入，并返回一个可序列化为 JSON 的值。在通过 `save_pretrained` 保存模型时，将使用此值。
   - 解码器期望接收原始数据（通常为字典）作为输入，并重建初始对象。在通过 `from_pretrained` 加载模型时，将使用此功能。
3. 为 `__init__` 签名添加类型注解。这一点很重要，它能让 mixin 了解该类期望的类型，从而确定应使用哪个解码器。

为了简洁起见，上述示例中的编码器/解码器函数并不健壮。在实际实现中，你很可能需要妥善处理各种边界情况。

## 快速对比

让我们快速总结一下上述两种方法及其各自的优缺点。下表仅供参考。您的框架可能有一些特定之处需要处理。本指南旨在提供关于如何处理集成的指导方针和思路。无论如何，如果您有任何疑问，请随时联系我们！

|           集成           |                                                      使用辅助函数                                                       |                                     使用 [ModelHubMixin](/docs/huggingface_hub/v2.4.0/zh/package_reference/mixins#huggingface_hub.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 等。                        |