Integra cualquier framework de ML con el Hub

El Hugging Face Hub facilita el alojamiento y la compartición de modelos con la comunidad. Admite decenas de librerías del ecosistema de código abierto. Estamos trabajando constantemente para ampliar este soporte y avanzar en el aprendizaje automático colaborativo. La librería huggingface_hub desempeña un papel clave en este proceso, ya que permite a cualquier script de Python subir y cargar archivos con facilidad.

Existen cuatro formas principales de integrar una biblioteca con el Hub:

1. Push to Hub: implementa un método para subir un modelo al Hub. Esto incluye los pesos del modelo, así como la tarjeta del modelo y cualquier otra información o datos relevantes necesarios para ejecutar el modelo (por ejemplo, registros de entrenamiento). Este método suele llamarse push_to_hub().
2. Descarga desde Hub: implemente un método para cargar un modelo desde el Hub. El método debe descargar la configuración/los pesos del modelo y cargar el modelo. Este método suele llamarse from_pretrained o load_from_hub().
3. Widgets: muestran un widget en la página de inicio de tus modelos en el Hub. Permite a los usuarios probar rápidamente un modelo desde el navegador.

En esta guía, nos centraremos en los dos primeros temas. Presentaremos los dos enfoques principales que puedes utilizar para integrar una biblioteca, junto con sus ventajas y desventajas. Todo se resume al final de la guía para ayudarte a elegir entre ambos. Ten en cuenta que estas son solo directrices que puedes adaptar libremente a tus requisitos.

Si te interesa el Inference y los Widgets, puedes seguir esta guía. En ambos casos, puedes ponerte en contacto con nosotros si estás integrando una biblioteca con el Hub y quieres que se incluya en nuestra documentación.

Un enfoque flexible: helpers

El primer enfoque para integrar una biblioteca en el Hub es implementar por ti mismo los métodos push_to_hub y from_pretrained. Esto te da total flexibilidad sobre qué archivos necesitas subir/descargar y cómo manejar las entradas específicas de tu marco de trabajo. Puedes consultar las guías de subir archivos y descargar archivos para aprender más sobre cómo hacerlo. Por ejemplo, así es como se implementa la integración de FastAI (ver push_to_hub_fastai() y from_pretrained_fastai()).

La implementación puede variar entre bibliotecas, pero el flujo de trabajo suele ser similar.

from_pretrained

Así suele verse un método 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

El método push_to_hub a menudo requiere un poco más de complejidad para gestionar la creación del repositorio, generar la tarjeta del modelo y guardar los pesos. Un enfoque común es guardar todos estos archivos en una carpeta temporal, subirlos y luego eliminarla.

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)

Esto es, por supuesto, solo un ejemplo. Si te interesan manipulaciones más complejas (eliminar archivos remotos, subir pesos sobre la marcha, persistir pesos localmente, etc.), consulta la guía de subir archivos.

Limitaciones

Aunque es flexible, este enfoque tiene algunas desventajas, especialmente en términos de mantenimiento. Los usuarios de Hugging Face suelen estar acostumbrados a funciones adicionales al trabajar con huggingface_hub. Por ejemplo, al cargar archivos desde el Hub, es común ofrecer parámetros como:

• token: para descargar de un repositorio privado
• revision: para descargar de una rama específica
• cache_dir: para cachear archivos en un directorio específico
• force_download/local_files_only: para reutilizar la caché o no
• proxies: configurar la sesión HTTP

Al subir modelos, se admiten parámetros similares:

• commit_message: mensaje de commit personalizado
• private: crea un repositorio privado si no existe
• create_pr: crea una PR en lugar de hacer push a main
• branch: enviar a una rama en lugar de a la rama main
• allow_patterns/ignore_patterns: filtra qué archivos se subirán
• token
• …

Todos estos parámetros pueden añadirse a las implementaciones que vimos anteriormente y pasarse a los métodos de huggingface_hub. Sin embargo, si un parámetro cambia o se añade una nueva funcionalidad, será necesario actualizar su paquete. Dar soporte a esos parámetros también implica más documentación que mantener de su parte. Para ver cómo mitigar estas limitaciones, pasemos a la siguiente sección herencia de clases.

Un enfoque más complejo: herencia de clases

Como vimos arriba, hay dos métodos principales para incluir en tu biblioteca e integrarla con el Hub: subir archivos (push_to_hub) y descargar archivos (from_pretrained). Puedes implementar esos métodos por ti mismo, pero conlleva precauciones. Para abordar esto, huggingface_hub proporciona una herramienta que usa herencia de clases. ¡Veamos cómo funciona!

En muchos casos, una biblioteca ya implementa su modelo mediante una clase de Python. La clase contiene las propiedades del modelo y los métodos para cargarlo, ejecutarlo, entrenarlo y evaluarlo. Nuestro enfoque consiste en extender esta clase para incluir funciones de carga y descarga mediante mixins. Un Mixin es una clase diseñada para extender una clase existente con un conjunto de funciones específicas mediante herencia múltiple. huggingface_hub proporciona su propio mixin, el ModelHubMixin. La clave aquí es comprender su comportamiento y cómo personalizarlo.

La clase ModelHubMixin implementa 3 métodos públicos (push_to_hub, save_pretrained y from_pretrained). Esos son los métodos que tus usuarios llamarán para cargar/guardar modelos con tu biblioteca. ModelHubMixin también define 2 métodos privados (_save_pretrained y _from_pretrained). Esos son los que debes implementar. Por lo tanto, para integrar tu biblioteca, debes:

1. Haga que su clase Model herede de ModelHubMixin.
2. Implementa los métodos privados:
   • _save_pretrained(): método que toma como entrada una ruta a un directorio y guarda el modelo en él. Debes escribir toda la lógica para volcar tu modelo en este método: tarjeta del modelo, pesos del modelo, archivos de configuración, registros de entrenamiento y figuras. Toda la información relevante para este modelo debe ser gestionada por este método. Las tarjetas de modelo son especialmente importantes para describir tu modelo. Consulta nuestra guía de implementación para más detalles.
   • _from_pretrained(): método de clase que toma como entrada un model_id y devuelve un modelo instanciado. El método debe descargar los archivos relevantes y cargarlos.
3. ¡Has terminado!

La ventaja de usar ModelHubMixin es que, una vez que te encargas de la serialización/carga de los archivos, estás listo para empezar. No tienes que preocuparte por cosas como la creación de repositorios, commits, PRs o revisiones. El ModelHubMixin también se asegura de que los métodos públicos estén documentados y anotados con tipos, y podrás ver el número de descargas de tu modelo en el Hub. Todo esto lo gestiona el ModelHubMixin y está disponible para tus usuarios.

Un ejemplo concreto: PyTorch

Un buen ejemplo de lo que vimos anteriormente es PyTorchModelHubMixin, nuestra integración para el framework PyTorch. Se trata de una integración lista para usar.

¿Cómo se usa?

Así es como cualquier usuario puede cargar/guardar un modelo de PyTorch desde/en el 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"

Implementación

La implementación es en realidad muy sencilla, y la implementación completa se puede encontrar aquí.

1. Primero, hereda tu clase de ModelHubMixin:

from huggingface_hub import ModelHubMixin

class PyTorchModelHubMixin(ModelHubMixin):
   (...)

2. Implementa el 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. Implementa el 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)

¡Y eso es todo! Tu biblioteca ahora permite a los usuarios subir y descargar archivos desde y hacia el Hub.

Uso avanzado

En la sección anterior, discutimos brevemente cómo funciona ModelHubMixin. En esta sección, veremos algunas de sus funciones más avanzadas para mejorar la integración de tu biblioteca con el Hugging Face Hub.

Tarjeta del modelo

ModelHubMixin genera la tarjeta del modelo por ti. Las tarjetas de modelo son archivos que acompañan a los modelos y proporcionan información importante sobre ellos. Internamente, las tarjetas de modelo son simples archivos Markdown con metadatos adicionales. Las tarjetas de modelo son esenciales para la descubribilidad, la reproducibilidad y el intercambio. Consulta la guía de tarjetas de modelo para más detalles.

Generar tarjetas de modelo de forma semiautomática es una buena manera de garantizar que todos los modelos subidos con su biblioteca compartan metadatos comunes: library_name, tags, license, pipeline_tag, etc. Esto hace que todos los modelos respaldados por su biblioteca sean fácilmente buscables en el Hub y proporciona algunos enlaces de recursos para los usuarios que llegan al Hub. Puede definir los metadatos directamente al heredar 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 defecto, se generará una tarjeta de modelo genérica con la información que hayas proporcionado (ejemplo: pyp1/VoiceCraft_giga830M). ¡Pero también puedes definir tu propia plantilla de tarjeta de modelo!

En este ejemplo, todos los modelos publicados con la clase VoiceCraft incluirán automáticamente una sección de citación y los detalles de la licencia. Para más detalles sobre cómo definir una plantilla de tarjeta de modelo, consulta la guía de tarjetas de modelo.

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 último, si deseas extender el proceso de generación de la tarjeta del modelo con valores dinámicos, puedes sobrescribir el 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 gestiona la configuración del modelo por ti. Comprueba automáticamente los valores de entrada al instanciar el modelo y los serializa en un archivo config.json. Esto ofrece dos ventajas:

1. Los usuarios podrán recargar el modelo con los mismos parámetros exactos que tú.
2. Tener un archivo config.json habilita automáticamente las analíticas en el Hub (es decir, el recuento de “descargas”).

Pero ¿cómo funciona en la práctica? Varias reglas hacen que el proceso sea lo más fluido posible desde la perspectiva del usuario:

• si su método __init__ espera una entrada config, se guardará automáticamente en el repositorio como config.json.
• si el parámetro de entrada config está anotado con un tipo de dataclass (p. ej. config: Optional[MyConfigClass] = None), el valor de config se deserializará correctamente para usted.
• Todos los valores pasados en la inicialización también se almacenarán en el archivo de configuración. Esto significa que no es necesario esperar una entrada config para beneficiarse de ello.

Ejemplo:

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}

¿Y si un valor no se puede serializar como JSON? De forma predeterminada, el valor se ignorará al guardar el archivo de configuración. Sin embargo, en algunos casos, tu biblioteca ya espera un objeto personalizado como entrada que no se puede serializar, y no deseas actualizar tu lógica interna para cambiar su tipo. ¡No te preocupes! Puedes pasar codificadores/decodificadores personalizados para cualquier tipo al heredar de ModelHubMixin. Esto requiere un poco más de trabajo, pero garantiza que tu lógica interna permanezca intacta al integrar tu biblioteca con el Hub.

Aquí hay un ejemplo concreto en el que una clase espera una configuración argparse.Namespace como entrada:

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

Una solución puede ser actualizar la firma de __init__ a def __init__(self, pattern: str, hidden_size: int) y actualizar todos los fragmentos de código que instancian su clase. Esta es una forma perfectamente válida de solucionarlo, pero podría romper aplicaciones de nivel inferior que utilizan su biblioteca.

Otra solución es proporcionar un codificador/descodificador simple para convertir argparse.Namespace en un diccionario.

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

En el fragmento anterior, tanto la lógica interna como la firma de __init__ de la clase no cambiaron. Esto significa que todos los fragmentos de código existentes de su biblioteca seguirán funcionando. Para lograrlo, tuvimos que:

1. Hereda del mixin (PytorchModelHubMixin en este caso).
2. Pasa un coders parámetro en la herencia. Es un diccionario donde las claves son los tipos personalizados que deseas procesar. Los valores son una tupla (encoder, decoder).
   • El codificador espera un objeto del tipo especificado como entrada y devuelve un valor serializable en JSON. Esto se utilizará al guardar un modelo con save_pretrained.
   • El decodificador espera datos crudos (típicamente un diccionario) como entrada y reconstruye el objeto inicial. Esto se utilizará al cargar el modelo con from_pretrained.
3. Añade una anotación de tipo a la firma de __init__. Esto es importante para que el mixin sepa qué tipo espera la clase y, por lo tanto, qué decodificador utilizar.

Por simplicidad, las funciones de codificador/decodificador del ejemplo anterior no son robustas. Para una implementación concreta, probablemente sea necesario manejar los casos límite de manera adecuada.

Comparación rápida

Resumamos rápidamente los dos enfoques que vimos, junto con sus ventajas y desventajas. La tabla siguiente es solo indicativa. Es posible que su marco de trabajo tenga algunas especificidades que deba abordar. Esta guía solo está aquí para ofrecer directrices e ideas sobre cómo manejar la integración. En cualquier caso, no dude en contactarnos si tiene alguna pregunta.