Intégrer n'importe quel framework ML avec le Hub

Le Hugging Face Hub facilite l’hébergement et le partage de modèles avec la communauté. Il prend en charge des dizaines de bibliothèques de l’écosystème Open Source. Nous travaillons en permanence à l’élargissement de cette prise en charge afin de faire progresser le Machine Learning collaboratif. La bibliothèque huggingface_hub joue un rôle clé dans ce processus, permettant à n’importe quel script Python de pousser et de charger des fichiers facilement.

Il existe quatre principales façons d'intégrer une bibliothèque avec le Hub :

1. Pousser vers le Hub : implémenter une méthode pour téléverser un modèle sur le Hub. Cela inclut les poids du modèle, ainsi que la carte du modèle et toute autre information ou donnée pertinente nécessaire à l'exécution du modèle (par exemple, les journaux d'entraînement). Cette méthode est souvent appelée push_to_hub().
2. Télécharger depuis le Hub : implémenter une méthode pour charger un modèle depuis le Hub. La méthode doit télécharger la configuration/les poids du modèle et charger le modèle. Cette méthode est souvent nommée from_pretrained ou load_from_hub().
3. Widgets : affiche un widget sur la page d'accueil de vos modèles sur le Hub. Il permet aux utilisateurs d'essayer rapidement un modèle depuis le navigateur.

Dans ce guide, nous nous concentrerons sur les deux premiers sujets. Nous présenterons les deux principales approches que vous pouvez utiliser pour intégrer une bibliothèque, ainsi que leurs avantages et leurs inconvénients. Tout est récapitulé à la fin du guide pour vous aider à choisir entre les deux. Gardez à l'esprit qu'il s'agit uniquement de lignes directrices que vous êtes libre d'adapter à vos besoins.

Si vous êtes intéressé par l’inférence et les widgets, vous pouvez suivre ce guide. Dans les deux cas, vous pouvez nous contacter si vous intégrez une bibliothèque avec le Hub et souhaitez être référencé dans notre documentation.

Une approche flexible : les utilitaires

La première approche pour intégrer une bibliothèque au Hub consiste à implémenter vous-même les méthodes push_to_hub et from_pretrained. Cela vous offre une flexibilité totale quant aux fichiers à téléverser/télécharger et à la manière de gérer les entrées spécifiques à votre framework. Vous pouvez consulter les guides téléverser des fichiers et télécharger des fichiers pour en savoir plus sur la marche à suivre. C’est, par exemple, ainsi que l’intégration FastAI est implémentée (voir push_to_hub_fastai() et from_pretrained_fastai()).

L'implémentation peut varier d'une bibliothèque à l'autre, mais le flux de travail est souvent similaire.

from_pretrained

Voici à quoi ressemble généralement une méthode 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

La méthode push_to_hub nécessite souvent un peu plus de complexité pour gérer la création du dépôt, générer la carte du modèle et enregistrer les poids. Une approche courante consiste à enregistrer tous ces fichiers dans un dossier temporaire, à le téléverser puis à le supprimer.

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)

Il s'agit bien sûr uniquement d'un exemple. Si vous êtes intéressé par des manipulations plus complexes (suppression de fichiers distants, téléversement de poids à la volée, persistance locale des poids, etc.), veuillez consulter le guide téléverser des fichiers.

Limites

Bien que flexible, cette approche présente certains inconvénients, notamment en termes de maintenance. Les utilisateurs de Hugging Face sont souvent habitués à des fonctionnalités supplémentaires lors de l’utilisation de huggingface_hub. Par exemple, lors du chargement de fichiers depuis le Hub, il est courant de proposer des paramètres tels que :

• token : pour télécharger depuis un dépôt privé
• revision : pour télécharger depuis une branche spécifique
• cache_dir : pour mettre les fichiers en cache dans un répertoire spécifique
• force_download/local_files_only : pour réutiliser le cache ou non
• proxies : configure la session HTTP

Lors de la publication de modèles, des paramètres similaires sont pris en charge :

• commit_message : message de commit personnalisé
• private : crée un dépôt privé s'il est manquant
• create_pr : créer une PR au lieu de pousser vers main
• branch : pousser vers une branche au lieu de la branche main
• allow_patterns/ignore_patterns : filtre les fichiers à téléverser
• token
• …

Tous ces paramètres peuvent être ajoutés aux implémentations que nous avons vues ci-dessus et passés aux méthodes de huggingface_hub. Cependant, si un paramètre change ou si une nouvelle fonctionnalité est ajoutée, vous devrez mettre à jour votre package. Prendre en charge ces paramètres signifie également plus de documentation à maintenir de votre côté. Pour voir comment atténuer ces limitations, passons à notre prochaine section héritage de classe.

Une approche plus complexe : l'héritage de classes

Comme nous l’avons vu plus haut, il existe deux méthodes principales à inclure dans votre bibliothèque pour l’intégrer au Hub : le téléversement de fichiers (push_to_hub) et le téléchargement de fichiers (from_pretrained). Vous pouvez implémenter ces méthodes vous-même, mais cela comporte des précautions. Pour y remédier, huggingface_hub fournit un outil qui utilise l’héritage de classes. Voyons comment il fonctionne !

Dans de nombreux cas, une bibliothèque implémente déjà son modèle à l'aide d'une classe Python. Cette classe contient les propriétés du modèle ainsi que des méthodes pour le charger, l'exécuter, l'entraîner et l'évaluer. Notre approche consiste à étendre cette classe pour inclure des fonctionnalités de téléversement et de téléchargement à l'aide de mixins. Un Mixin est une classe conçue pour étendre une classe existante avec un ensemble de fonctionnalités spécifiques en utilisant l'héritage multiple. huggingface_hub fournit son propre mixin, le ModelHubMixin. L'essentiel est de comprendre son comportement et la manière de le personnaliser.

La classe ModelHubMixin implémente 3 méthodes publiques (push_to_hub, save_pretrained et from_pretrained). Ce sont les méthodes que vos utilisateurs appelleront pour charger/enregistrer des modèles avec votre bibliothèque. ModelHubMixin définit également 2 méthodes privées (_save_pretrained et _from_pretrained). Ce sont celles que vous devez implémenter. Pour intégrer votre bibliothèque, vous devez donc :

1. Faites hériter votre classe Model de ModelHubMixin.
2. Implémentez les méthodes privées :
   • _save_pretrained() : méthode qui prend en entrée un chemin vers un répertoire et y enregistre le modèle. Vous devez écrire toute la logique permettant de décharger votre modèle dans cette méthode : carte du modèle, poids du modèle, fichiers de configuration, journaux d’entraînement et figures. Toute information pertinente pour ce modèle doit être gérée par cette méthode. Les cartes de modèle sont particulièrement importantes pour décrire votre modèle. Consultez notre guide d’implémentation pour plus de détails.
   • _from_pretrained() : méthode de classe qui prend en entrée un model_id et renvoie un modèle instancié. La méthode doit télécharger les fichiers pertinents et les charger.
3. Vous avez terminé !

L’avantage d’utiliser ModelHubMixin est que, une fois la sérialisation/le chargement des fichiers pris en charge, vous êtes prêt à partir. Vous n’avez pas à vous soucier de la création de dépôts, des commits, des PR ou des révisions. Le ModelHubMixin garantit également que les méthodes publiques sont documentées et annotées avec des types, et vous pourrez consulter le nombre de téléchargements de votre modèle sur le Hub. Tout cela est géré par le ModelHubMixin et mis à disposition de vos utilisateurs.

Un exemple concret : PyTorch

Un bon exemple de ce que nous avons vu plus haut est PyTorchModelHubMixin, notre intégration pour le framework PyTorch. Il s’agit d’une intégration prête à l’emploi.

Comment l'utiliser ?

Voici comment n’importe quel utilisateur peut charger/enregistrer un modèle PyTorch depuis/vers le 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"

Implémentation

La mise en œuvre est en réalité très simple, et l'implémentation complète peut être trouvée ici.

1. Premièrement, faites hériter votre classe de ModelHubMixin :

from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)

2. Implémenter la méthode _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. Implémenter la méthode _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)

Et c’est tout ! Votre bibliothèque permet désormais aux utilisateurs de téléverser et de télécharger des fichiers vers et depuis le Hub.

Utilisation avancée

Dans la section précédente, nous avons brièvement expliqué le fonctionnement de ModelHubMixin. Dans cette section, nous allons découvrir certaines de ses fonctionnalités avancées afin d'améliorer l'intégration de votre bibliothèque avec le Hugging Face Hub.

Fiche du modèle

ModelHubMixin génère la carte du modèle pour vous. Les cartes de modèle sont des fichiers qui accompagnent les modèles et fournissent des informations importantes à leur sujet. En interne, les cartes de modèle sont de simples fichiers Markdown avec des métadonnées supplémentaires. Les cartes de modèle sont essentielles pour la découvrabilité, la reproductibilité et le partage ! Consultez le guide des cartes de modèle pour plus de détails.

Générer les fiches de modèle de manière semi-automatique est une bonne façon de s’assurer que tous les modèles publiés avec votre bibliothèque partagent des métadonnées communes : library_name, tags, license, pipeline_tag, etc. Cela rend tous les modèles soutenus par votre bibliothèque facilement recherchables sur le Hub et fournit certains liens de ressources pour les utilisateurs qui atterrissent sur le Hub. Vous pouvez définir les métadonnées directement lors de l’héritage de 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"]
):
   ...

Par défaut, une carte de modèle générique sera générée avec les informations que vous avez fournies (exemple : pyp1/VoiceCraft_giga830M). Mais vous pouvez également définir votre propre modèle de carte de modèle !

Dans cet exemple, tous les modèles publiés avec la classe VoiceCraft incluront automatiquement une section de citation et les détails de la licence. Pour plus de détails sur la définition d'un modèle de carte de modèle, consultez le guide des cartes de modèle.

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,
   ...
):
   ...

Enfin, si vous souhaitez étendre le processus de génération de la carte de modèle avec des valeurs dynamiques, vous pouvez surcharger la méthode 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

Configuration

ModelHubMixin gère la configuration du modèle pour vous. Il vérifie automatiquement les valeurs d’entrée lors de l’instantiation du modèle et les sérialise dans un fichier config.json. Cela offre deux avantages :

1. Les utilisateurs pourront recharger le modèle avec exactement les mêmes paramètres que vous.
2. La présence d'un fichier config.json active automatiquement les analyses sur le Hub (c'est-à-dire le compteur de « téléchargements »).

Mais comment cela fonctionne-t-il en pratique ? Plusieurs règles rendent le processus aussi fluide que possible du point de vue de l'utilisateur :

• si votre méthode __init__ attend un paramètre config, il sera automatiquement enregistré dans le dépôt sous le nom config.json.
• si le paramètre d’entrée config est annoté avec un type dataclass (par ex. config: Optional[MyConfigClass] = None), la valeur de config sera correctement désérialisée pour vous.
• toutes les valeurs passées à l’initialisation seront également stockées dans le fichier de configuration. Cela signifie que vous n’avez pas forcément besoin d’attendre une entrée config pour en bénéficier.

Exemple :

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}

Mais que se passe-t-il si une valeur ne peut pas être sérialisée en JSON ? Par défaut, la valeur est ignorée lors de l’enregistrement du fichier de configuration. Cependant, dans certains cas, votre bibliothèque s’attend déjà à un objet personnalisé en entrée qui ne peut pas être sérialisé, et vous ne souhaitez pas modifier votre logique interne pour changer son type. Pas de panique ! Vous pouvez passer des encodeurs/décodeurs personnalisés pour n’importe quel type lors de l’héritage de ModelHubMixin. Cela demande un peu plus de travail, mais garantit que votre logique interne reste intacte lors de l’intégration de votre bibliothèque avec le Hub.

Voici un exemple concret où une classe attend une configuration argparse.Namespace en entrée :

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

Une solution consiste à mettre à jour la signature de __init__ en def __init__(self, pattern: str, hidden_size: int) et à mettre à jour tous les extraits de code qui instancient votre classe. C'est une approche parfaitement valide pour corriger le problème, mais elle risque de casser les applications en aval qui utilisent votre bibliothèque.

Une autre solution consiste à fournir un encodeur/décodeur simple pour convertir argparse.Namespace en dictionnaire.

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

Dans l'extrait ci-dessus, à la fois la logique interne et la signature de __init__ de la classe n'ont pas changé. Cela signifie que tous les extraits de code existants pour votre bibliothèque continueront à fonctionner. Pour y parvenir, nous avons dû :

1. Hériter du mixin (PytorchModelHubMixin dans ce cas).
2. Passez un coders paramètre dans l'héritage. Il s'agit d'un dictionnaire dont les clés sont les types personnalisés que vous souhaitez traiter. Les valeurs sont un tuple (encoder, decoder).
   • L'encodeur s'attend à un objet du type spécifié en entrée et renvoie une valeur jsonable. Cela sera utilisé lors de l'enregistrement d'un modèle avec save_pretrained.
   • Le décodeur s'attend à des données brutes (généralement un dictionnaire) en entrée et reconstruit l'objet initial. Cela sera utilisé lors du chargement du modèle avec from_pretrained.
3. Ajoutez une annotation de type à la signature de __init__. C'est important pour que le mixin sache quel type est attendu par la classe et, par conséquent, quel décodeur utiliser.

Par souci de simplicité, les fonctions encodeur/décodeur de l’exemple ci-dessus ne sont pas robustes. Pour une implémentation concrète, il faudrait très probablement traiter correctement les cas limites.

Comparaison rapide

Faisons un rapide résumé des deux approches que nous avons vues, avec leurs avantages et leurs inconvénients. Le tableau ci-dessous est purement indicatif. Votre framework peut présenter des spécificités auxquelles vous devez répondre. Ce guide est là uniquement pour fournir des directives et des idées sur la manière de gérer l’intégration. Dans tous les cas, n’hésitez pas à nous contacter si vous avez des questions !