Интеграция любого ML-фреймворка с Hub

Hugging Face Hub упрощает размещение моделей и их обмен с сообществом. Он поддерживает десятки библиотек в экосистеме Open Source. Мы постоянно расширяем эту поддержку, чтобы продвигать совместное машинное обучение. Библиотека huggingface_hub играет ключевую роль в этом процессе, позволяя любому Python-скрипту легко загружать и читать файлы.

Существует четыре основных способа интеграции библиотеки с Hub:

1. Push to Hub: реализуйте метод для загрузки модели в Hub. Он должен включать веса модели, а также карточку модели и любую другую необходимую информацию или данные для запуска модели (например, журналы обучения). Этот метод обычно называется push_to_hub().
2. Загрузка из Hub: реализуйте метод для загрузки модели из Hub. Метод должен скачивать конфигурацию/веса модели и загружать её. Этот метод часто называется from_pretrained или load_from_hub().
3. Виджеты: отображают виджет на главной странице ваших моделей в Hub. Это позволяет пользователям быстро протестировать модель прямо из браузера.

В этом руководстве мы сосредоточимся на первых двух темах. Мы рассмотрим два основных подхода, которые можно использовать для интеграции библиотеки, а также их преимущества и недостатки. В конце руководства приведена сводная информация, которая поможет вам выбрать между этими двумя вариантами. Помните, что это лишь рекомендации, которые вы можете свободно адаптировать под свои требования.

Если вас интересует Inference и Widgets, вы можете следовать этому руководству. В обоих случаях вы можете связаться с нами, если интегрируете библиотеку с Hub и хотите, чтобы она была указана в нашей документации.

Гибкий подход: вспомогательные функции

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 and download files guides to learn more about how to do that. This is, for example how the FastAI integration is implemented (see push_to_hub_fastai() and 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=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 часто требует большей сложности для обработки создания репозитория, генерации карточки модели и сохранения весов. Обычный подход — сохранить все эти файлы во временной папке, загрузить её, а затем удалить.

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 files.

Ограничения

Несмотря на гибкость, этот подход имеет ряд недостатков, особенно в части поддержки. Пользователи 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. Класс содержит свойства модели и методы для её загрузки, запуска, обучения и оценки. Наш подход состоит в расширении этого класса с добавлением функций загрузки и скачивания с помощью миксинов. Миксин — это класс, предназначенный для расширения существующего класса набором конкретных функций с использованием множественного наследования. huggingface_hub предоставляет собственный миксин — ModelHubMixin. Ключевое здесь — понять его поведение и способы его настройки.

Класс ModelHubMixin реализует 3 публичных метода (push_to_hub, save_pretrained и from_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. Это готовая к использованию интеграция.

Как им пользоваться?

Вот как любой пользователь может загрузить/сохранить модель PyTorch из/в Hub:

>>> 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"

Реализация

Реализация на самом деле очень проста, а полный код можно найти здесь.

1. Сначала унаследуйте свой класс от ModelHubMixin:

from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)

2. Реализуйте метод _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))

3. Реализуйте метод _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,
      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. В этом разделе мы познакомимся с некоторыми его более продвинутыми возможностями, которые помогут улучшить интеграцию вашей библиотеки с Hugging Face Hub.

Карточка модели

ModelHubMixin автоматически генерирует карточку модели. Карточки моделей — это файлы, сопровождающие модели и содержащие важную информацию о них. По сути, карточки моделей — это простые Markdown-файлы с дополнительными метаданными. Карточки моделей необходимы для обнаружения, воспроизводимости и обмена моделями! Подробнее см. в руководстве по карточкам моделей.

Полуавтоматическая генерация карточек моделей — хороший способ обеспечить, чтобы все модели, загружаемые с помощью вашей библиотеки, имели общие метаданные: library_name, tags, license, pipeline_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. Это даёт два преимущества:

1. Пользователи смогут перезагрузить модель с точно такими же параметрами, как у вас.
2. Наличие файла config.json автоматически включает аналитику на Hub (то есть счётчик «скачиваний»).

Но как это работает на практике? Несколько правил делают процесс максимально плавным с точки зрения пользователя:

• если метод __init__ ожидает входной параметр config, он будет автоматически сохранён в репозитории как config.json.
• если параметр ввода config аннотирован типом dataclass (например, 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. Наследуйте от миксина (в данном случае PytorchModelHubMixin).
2. Передайте параметр coders при наследовании. Это словарь, где ключами являются пользовательские типы, которые вы хотите обработать. Значениями являются кортежи (encoder, decoder).
   • Кодировщик ожидает на вход объект указанного типа и возвращает значение, пригодное для сериализации в JSON. Это будет использоваться при сохранении модели с помощью save_pretrained.
   • Декодер ожидает на вход необработанные данные (как правило, словарь) и восстанавливает исходный объект. Это используется при загрузке модели с помощью from_pretrained.
3. Добавьте аннотацию типа к сигнатуре __init__. Это важно, чтобы миксин понимал, какой тип ожидается классом, и, соответственно, какой декодер использовать.

Ради простоты функции кодировщика/декодировщика в примере выше не являются устойчивыми. Для конкретной реализации, скорее всего, потребуется корректно обрабатывать граничные случаи.

Краткое сравнение

Давайте кратко подведём итоги двух рассмотренных подходов, указав их преимущества и недостатки. Приведённая ниже таблица носит лишь справочный характер. Ваша фреймворк может иметь свои особенности, которые необходимо учитывать. Это руководство предназначено только для предоставления рекомендаций и идей по вопросам интеграции. В любом случае, не стесняйтесь обращаться к нам, если у вас есть вопросы!