Hugging Face Hub integration

[qubvel]

Hi @siyuanliii,

Thanks for this great work! I am excited for the code release!

I am reaching you about the integration with the 🤗 hub from day 0, so you can track the download number, have a nice model card and anyone can automatically load the model using from_pretrained (and push it using push_to_hub, similar to models in the Transformers library). It leverages the PyTorchModelHubMixin class which allows to inherit these methods.

Usage can be as follows:

from masa.models import MasaModel

model = MasaModel.from_pretrained("<your-username-on-hf-hub>/<model-name>")

This means people don't need to manually download a checkpoint first in their local environment, it just loads it automatically from the hub. The safetensors format is used to ensure safe serialization of the weights rather than pickle.

Integration usually is really simple, you just need to subclass your PyTorch model class from nn.Module together with PyTorchModelHubMixin class. For example:

import torch
import torch.nn as nn
+ from huggingface_hub import PyTorchModelHubMixin

class MasaModel(
        nn.Module,
+        PyTorchModelHubMixin,
+        library_name="masa",
+        repo_url="https://github.com/siyuanliii/masa",
        # ^ Optional metadata to generate model card
    ):
    def __init__(self, ...):
        ...

Then, you will be able to push the pretrained model to the hub to make it discoverable:

from masa.models import MasaModel

model = MasaModel()
model.save_pretrained("<your-username-on-hf-hub>/<model-name>", push_to_hub=True)

Would you consider this integration? Please, let me know if you have any questions or concerns!

Thank you, Pavel Iakubovskii ML Engineer @ HF 🤗

[siyuanliii]

Hi Pavel, thanks for reaching out and being interested in our work. Yes, I'd love to integrate the model weights into the hugging face hub!

[qubvel]

Awesome, let me know if you have any questions! Just ping me in this issue 🤗