Integre qualquer framework de ML com o Hub

O Hugging Face Hub facilita o hospedagem e o compartilhamento de modelos com a comunidade. Ele suporta dezenas de bibliotecas no ecossistema de Código Aberto. Estamos sempre trabalhando para expandir esse suporte e impulsionar o Aprendizado de Máquina colaborativo. A biblioteca huggingface_hub desempenha um papel fundamental nesse processo, permitindo que qualquer script em Python envie e carregue arquivos com facilidade.

Existem quatro maneiras principais de integrar uma biblioteca com o Hub:

1. Push to Hub: implemente um método para enviar um modelo ao Hub. Isso inclui os pesos do modelo, bem como o cartão do modelo e quaisquer outras informações ou dados relevantes necessários para executar o modelo (por exemplo, logs de treinamento). Esse método é frequentemente chamado de push_to_hub().
2. Download do Hub: implemente um método para carregar um modelo do Hub. O método deve baixar a configuração/weights do modelo e carregá-lo. Este método é frequentemente chamado de from_pretrained ou load_from_hub().
3. Widgets: exibe um widget na página inicial dos seus modelos no Hub. Permite que os usuários testem rapidamente um modelo a partir do navegador.

Neste guia, focaremos nos dois primeiros tópicos. Apresentaremos as duas principais abordagens que você pode usar para integrar uma biblioteca, com suas vantagens e desvantagens. Tudo é resumido no final do guia para ajudar você a escolher entre as duas. Tenha em mente que estas são apenas diretrizes que você é livre para adaptar aos seus requisitos.

Se você tem interesse em Inferência e Widgets, pode seguir este guia. Em ambos os casos, você pode nos contatar se estiver integrando uma biblioteca com o Hub e quiser ser listado na nossa documentação.

Uma abordagem flexível: auxiliares

A primeira abordagem para integrar uma biblioteca ao Hub é implementar você mesmo os métodos push_to_hub e from_pretrained. Isso lhe dá total flexibilidade sobre quais arquivos você precisa enviar/baixar e como lidar com entradas específicas do seu framework. Você pode consultar os dois guias de enviar arquivos e baixar arquivos para aprender mais sobre como fazer isso. É assim, por exemplo, que a integração do FastAI é implementada (veja push_to_hub_fastai() e from_pretrained_fastai()).

A implementação pode variar entre as bibliotecas, mas o fluxo de trabalho costuma ser semelhante.

from_pretrained

É assim que um método from_pretrained geralmente se parece:

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

O método push_to_hub geralmente exige um pouco mais de complexidade para lidar com a criação do repositório, gerar o cartão do modelo e salvar os pesos. Uma abordagem comum é salvar todos esses arquivos em uma pasta temporária, enviá-la e, em seguida, excluí-la.

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)

Isso é, claro, apenas um exemplo. Se você estiver interessado em manipulações mais complexas (excluir arquivos remotos, enviar pesos em tempo real, persistir pesos localmente, etc.), consulte o guia de upload de arquivos.

Limitações

Embora seja flexível, essa abordagem apresenta algumas desvantagens, especialmente em termos de manutenção. Os usuários do Hugging Face costumam se acostumar a recursos adicionais ao trabalhar com huggingface_hub. Por exemplo, ao carregar arquivos do Hub, é comum oferecer parâmetros como:

• token: para baixar de um repositório privado
• revision: para baixar de um branch específico
• cache_dir: para armazenar arquivos em cache em um diretório específico
• force_download/local_files_only: para reutilizar o cache ou não
• proxies: configura a sessão HTTP

Ao enviar modelos, parâmetros semelhantes são suportados:

• commit_message: mensagem de commit personalizada
• private: cria um repositório privado se não existir
• create_pr: cria um PR em vez de fazer push para main
• branch: faz push para um branch em vez do branch main
• allow_patterns/ignore_patterns: filtra quais arquivos serão enviados
• token
• …

Todos esses parâmetros podem ser adicionados às implementações que vimos acima e passados aos métodos do huggingface_hub. No entanto, se um parâmetro for alterado ou um novo recurso for adicionado, você precisará atualizar seu pacote. Suportar esses parâmetros também significa mais documentação para manter do seu lado. Para ver como mitigar essas limitações, vamos avançar para a nossa próxima seção herança de classe.

Uma abordagem mais complexa: herança de classes

Como vimos acima, existem dois métodos principais para incluir na sua biblioteca e integrá-la ao Hub: enviar arquivos (push_to_hub) e baixar arquivos (from_pretrained). Você pode implementar esses métodos por conta própria, mas isso vem com caveats. Para lidar com isso, o huggingface_hub fornece uma ferramenta que usa herança de classes. Vamos ver como funciona!

Em muitos casos, uma biblioteca já implementa seu modelo usando uma classe Python. A classe contém as propriedades do modelo e métodos para carregá-lo, executá-lo, treiná-lo e avaliá-lo. Nossa abordagem é estender essa classe para incluir recursos de upload e download usando mixins. Um Mixin é uma classe que visa estender uma classe existente com um conjunto de recursos específicos usando herança múltipla. O huggingface_hub fornece seu próprio mixin, o ModelHubMixin. O ponto-chave aqui é entender seu comportamento e como personalizá-lo.

A classe ModelHubMixin implementa 3 métodos públicos (push_to_hub, save_pretrained e from_pretrained). Esses são os métodos que seus usuários chamarão para carregar/salvar modelos com sua biblioteca. O ModelHubMixin também define 2 métodos privados (_save_pretrained e _from_pretrained). Esses são os que você deve implementar. Portanto, para integrar sua biblioteca, você deve:

1. Faça a sua classe Model herdar de ModelHubMixin.
2. Implemente os métodos privados:
   • _save_pretrained(): método que recebe como entrada um caminho para um diretório e salva o modelo nele. Você deve escrever toda a lógica para despejar seu modelo neste método: model card, pesos do modelo, arquivos de configuração, logs de treinamento e figuras. Qualquer informação relevante para este modelo deve ser tratada por este método. Model Cards são particularmente importantes para descrever seu modelo. Consulte nossa guia de implementação para mais detalhes.
   • _from_pretrained(): método de classe que recebe como entrada um model_id e retorna um modelo instanciado. O método deve baixar os arquivos relevantes e carregá-los.
3. Você terminou!

A vantagem de usar ModelHubMixin é que, uma vez que você cuida da serialização/carregamento dos arquivos, está pronto para ir. Você não precisa se preocupar com coisas como criação de repositório, commits, PRs ou revisões. O ModelHubMixin também garante que os métodos públicos estejam documentados e com anotações de tipo, e você poderá visualizar a contagem de downloads do seu modelo no Hub. Tudo isso é tratado pelo ModelHubMixin e disponível para seus usuários.

Um exemplo concreto: PyTorch

Um bom exemplo do que vimos acima é o PyTorchModelHubMixin, nossa integração para o framework PyTorch. Esta é uma integração pronta para uso.

Como usá-lo?

É assim que qualquer usuário pode carregar/salvar um modelo do PyTorch a partir/do 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"

Implementação

A implementação é, na verdade, muito direta, e a implementação completa pode ser encontrada aqui.

1. Primeiro, herde sua classe de ModelHubMixin:

from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)

2. Implemente o método _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. Implemente o método _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)

E é isso! Sua biblioteca agora permite que os usuários façam upload e download de arquivos para e do Hub.

Uso avançado

Na seção acima, discutimos rapidamente como o ModelHubMixin funciona. Nesta seção, veremos alguns de seus recursos mais avançados para melhorar a integração da sua biblioteca com o Hugging Face Hub.

Cartão do modelo

ModelHubMixin gera o cartão do modelo para você. Cartões de modelo são arquivos que acompanham os modelos e fornecem informações importantes sobre eles. Por baixo dos panos, cartões de modelo são arquivos Markdown simples com metadados adicionais. Cartões de modelo são essenciais para descoberta, reprodutibilidade e compartilhamento! Consulte o guia de cartões de modelo para mais detalhes.

Gerar cards de modelo de forma semiautomática é uma boa maneira de garantir que todos os modelos enviados com sua biblioteca compartilhem metadados comuns: library_name, tags, license, pipeline_tag, etc. Isso torna todos os modelos suportados por sua biblioteca facilmente pesquisáveis no Hub e fornece alguns links de recursos para usuários que acessam o Hub. Você pode definir os metadados diretamente ao herdar 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"]
):
   ...

Por padrão, um cartão de modelo genérico será gerado com as informações que você forneceu (exemplo: pyp1/VoiceCraft_giga830M). Mas você também pode definir seu próprio modelo de cartão de modelo!

Neste exemplo, todos os modelos enviados com a classe VoiceCraft incluirão automaticamente uma seção de citação e os detalhes da licença. Para mais detalhes sobre como definir um modelo de cartão de modelo, consulte o guia de Model Cards.

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

Por fim, se você deseja estender o processo de geração do cartão do modelo com valores dinâmicos, pode sobrescrever o método 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

Config

ModelHubMixin gerencia a configuração do modelo para você. Ele verifica automaticamente os valores de entrada ao instanciar o modelo e os serializa em um arquivo config.json. Isso oferece 2 benefícios:

1. Os usuários poderão recarregar o modelo com exatamente os mesmos parâmetros que você.
2. Ter um arquivo config.json ativa automaticamente a análise de dados no Hub (ou seja, a contagem de “downloads”).

Mas como isso funciona na prática? Várias regras tornam o processo o mais suave possível do ponto de vista do usuário:

• se o seu método __init__ espera uma entrada config, ele será salvo automaticamente no repositório como config.json.
• se o parâmetro de entrada config for anotado com um tipo de dataclass (por exemplo, config: Optional[MyConfigClass] = None), o valor de config será desserializado corretamente para você.
• todos os valores passados na inicialização também serão armazenados no arquivo de configuração. Isso significa que você não precisa necessariamente esperar uma entrada config para se beneficiar disso.

Exemplo:

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}

E se um valor não puder ser serializado como JSON? Por padrão, o valor será ignorado ao salvar o arquivo de configuração. No entanto, em alguns casos, sua biblioteca já espera um objeto personalizado como entrada que não pode ser serializado, e você não deseja atualizar sua lógica interna para alterar seu tipo. Sem problemas! Você pode passar codificadores/decodificadores personalizados para qualquer tipo ao herdar de ModelHubMixin. Isso exige um pouco mais de trabalho, mas garante que sua lógica interna permaneça intacta ao integrar sua biblioteca com o Hub.

Aqui está um exemplo concreto em que uma classe espera uma configuração argparse.Namespace como entrada:

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

Uma solução pode ser atualizar a assinatura de __init__ para def __init__(self, pattern: str, hidden_size: int) e atualizar todos os trechos de código que instanciam sua classe. Essa é uma maneira perfeitamente válida de corrigir o problema, mas pode quebrar aplicativos downstream que usam sua biblioteca.

Outra solução é fornecer um codificador/decodificador simples para converter argparse.Namespace em um dicionário.

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

No trecho acima, tanto a lógica interna quanto a assinatura de __init__ da classe não foram alteradas. Isso significa que todos os trechos de código existentes para a sua biblioteca continuarão a funcionar. Para alcançar isso, tivemos que:

1. Herde do mixin (PytorchModelHubMixin neste caso).
2. Passe um parâmetro coders na herança. Este é um dicionário em que as chaves são os tipos personalizados que você deseja processar. Os valores são uma tupla (encoder, decoder).
   • O codificador espera um objeto do tipo especificado como entrada e retorna um valor jsonable. Isso será usado ao salvar um modelo com save_pretrained.
   • O decodificador espera dados brutos (tipicamente um dicionário) como entrada e reconstrói o objeto inicial. Isso será usado ao carregar o modelo com from_pretrained.
3. Adicione uma anotação de tipo à assinatura de __init__. Isso é importante para que o mixin saiba qual tipo é esperado pela classe e, portanto, qual decodificador usar.

Por simplicidade, as funções de codificador/decodificador no exemplo acima não são robustas. Para uma implementação concreta, provavelmente seria necessário tratar adequadamente os casos extremos.

Comparação rápida

Vamos resumir rapidamente as duas abordagens que vimos, com suas vantagens e desvantagens. A tabela abaixo é apenas indicativa. Seu framework pode ter algumas especificidades que você precisa abordar. Este guia está aqui apenas para fornecer diretrizes e ideias sobre como lidar com a integração. De qualquer forma, não hesite em nos contatar se tiver alguma dúvida!