huggingface / huggingface_hub

The official Python client for the Huggingface Hub.
https://huggingface.co/docs/huggingface_hub
Apache License 2.0
1.83k stars 471 forks source link

Support custom kwargs for model card in save_pretrained #2310

Closed qubvel closed 1 month ago

qubvel commented 1 month ago

Hi!

There is an option to customize the model card for Mixins, it would be great to support custom kwargs for a card also.

Here is an example:

import torch
from huggingface_hub import PyTorchModelHubMixin

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 }}
---
{{ custom_data }}
"""

class Model(torch.nn.Module, PyTorchModelHubMixin, model_card_template=MODEL_CARD_TEMPLATE):
    pass

model = Model()
model_card_kwargs = {"custom_data": "This is an awesome model..."}
model.save_pretrained("test-model", model_card_kwargs=model_card_kwargs, push_to_hub=True)

Generated readme:

---
tags:
- pytorch_model_hub_mixin
- model_hub_mixin
---
This is awesome model...

Let me know what you think about this feature! If it's OK, I can add the necessary tests.

HuggingFaceDocBuilderDev commented 1 month ago

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

Wauplin commented 1 month ago

The mixin currently have a set of very relevant model card metadata. For more customization, it's possible to overwrite the generate_model_card method. This method allows more flexibility and access to model's values/config. It return a ModelCard object. For example:

class Model(nn.Module, PytorchModelHubMixin):
    ...

    def generate_model_card(self) -> ModelCard:
        card = super(self).generate_model_card()
        card.data.custom_key = "custom value"
        return card

I'd prefer not to add too many features to the main way of defining tags/library_name/... to avoid complexity and promote this flexible method instead. WDYT?

qubvel commented 1 month ago

The feature is inspired by this PR, as you may see I try to make it possible to pass a custom metrics dictionary and dataset name to be able to generate cards like this. https://huggingface.co/qubvel-hf/oxford-pet-segmentation

model.save_pretrained('./my_model', metrics={'accuracy': 0.95}, dataset='my_dataset')

You might want to pass arguments not at the time of model class creation, but at the time of model saving, for example when the model is trained. I overridden the generate_model_card() method but I didn't find a way to pass any parameters from the save_pretrained() method. The only way I found is to set them as private attributes in save_pretrained and then remove them. https://github.com/qubvel/segmentation_models.pytorch/blob/3d6da1d74636873372c265f300862a6a6d01777d/segmentation_models_pytorch/base/hub_mixin.py#L107 Now it looks like a dirty hack :(

qubvel commented 1 month ago

@Wauplin thank you for the feedback! I will update the integration, would be cool to include this PR changes, to get rid of private attributes :)

Could you add a small test in test_hub_mixin_pytorch.py with a new class that overwrites generate_model_card + pass data in save_pretrained and check you have the expected result. Thanks in advance!

Done! Please, have a look!

qubvel commented 1 month ago

Thanks for the quick review and merge 🤗