Add paper URL to hub mixin (#2917)

* First draft

* Add test

* Apply style fixes

* fix tests

* style

---------

Co-authored-by: github-actions[bot] <github-actions[bot]@users.noreply.github.com>
Co-authored-by: Celina Hanouti <[email protected]>

Author: 3 people

src/huggingface_hub/hub_mixin.py

@@ -58,7 +58,8 @@ class DataclassInstance(Protocol):
 ---
 
 This model has been pushed to the Hub using the [PytorchModelHubMixin](https://huggingface.co/docs/huggingface_hub/package_reference/mixins#huggingface_hub.PyTorchModelHubMixin) integration:
-- Library: {{ repo_url | default("[More Information Needed]", true) }}
+- Code: {{ repo_url | default("[More Information Needed]", true) }}
+- Paper: {{ paper_url | default("[More Information Needed]", true) }}
 - Docs: {{ docs_url | default("[More Information Needed]", true) }}
 """
 
@@ -67,8 +68,9 @@ class DataclassInstance(Protocol):
 class MixinInfo:
     model_card_template: str
     model_card_data: ModelCardData
-    repo_url: Optional[str] = None
     docs_url: Optional[str] = None
+    paper_url: Optional[str] = None
+    repo_url: Optional[str] = None
 
 
 class ModelHubMixin:
@@ -88,6 +90,8 @@ class ModelHubMixin:
     Args:
         repo_url (`str`, *optional*):
             URL of the library repository. Used to generate model card.
+        paper_url (`str`, *optional*):
+            URL of the library paper. Used to generate model card.
         docs_url (`str`, *optional*):
             URL of the library documentation. Used to generate model card.
         model_card_template (`str`, *optional*):
@@ -110,7 +114,7 @@ class ModelHubMixin:
         pipeline_tag (`str`, *optional*):
             Tag of the pipeline. Used to generate model card. E.g. "text-classification".
         tags (`List[str]`, *optional*):
-            Tags to be added to the model card. Used to generate model card. E.g. ["x-custom-tag", "arxiv:2304.12244"]
+            Tags to be added to the model card. Used to generate model card. E.g. ["computer-vision"]
         coders (`Dict[Type, Tuple[Callable, Callable]]`, *optional*):
             Dictionary of custom types and their encoders/decoders. Used to encode/decode arguments that are not
             jsonable by default. E.g dataclasses, argparse.Namespace, OmegaConf, etc.
@@ -124,8 +128,9 @@ class ModelHubMixin:
     >>> class MyCustomModel(
     ...         ModelHubMixin,
     ...         library_name="my-library",
-    ...         tags=["x-custom-tag", "arxiv:2304.12244"],
+    ...         tags=["computer-vision"],
     ...         repo_url="https://github.com/huggingface/my-cool-library",
+    ...         paper_url="https://arxiv.org/abs/2304.12244",
     ...         docs_url="https://huggingface.co/docs/my-cool-library",
     ...         # ^ optional metadata to generate model card
     ...     ):
@@ -194,6 +199,7 @@ def __init_subclass__(
         *,
         # Generic info for model card
         repo_url: Optional[str] = None,
+        paper_url: Optional[str] = None,
         docs_url: Optional[str] = None,
         # Model card template
         model_card_template: str = DEFAULT_MODEL_CARD,
@@ -234,6 +240,7 @@ def __init_subclass__(
 
             # Inherit other info
             info.docs_url = cls._hub_mixin_info.docs_url
+            info.paper_url = cls._hub_mixin_info.paper_url
             info.repo_url = cls._hub_mixin_info.repo_url
         cls._hub_mixin_info = info
 
@@ -242,6 +249,8 @@ def __init_subclass__(
             info.model_card_template = model_card_template
         if repo_url is not None:
             info.repo_url = repo_url
+        if paper_url is not None:
+            info.paper_url = paper_url
         if docs_url is not None:
             info.docs_url = docs_url
         if language is not None:
@@ -692,6 +701,7 @@ def generate_model_card(self, *args, **kwargs) -> ModelCard:
             card_data=self._hub_mixin_info.model_card_data,
             template_str=self._hub_mixin_info.model_card_template,
             repo_url=self._hub_mixin_info.repo_url,
+            paper_url=self._hub_mixin_info.paper_url,
             docs_url=self._hub_mixin_info.docs_url,
             **kwargs,
         )
@@ -718,6 +728,7 @@ class PyTorchModelHubMixin(ModelHubMixin):
     ...         PyTorchModelHubMixin,
     ...         library_name="keras-nlp",
     ...         repo_url="https://github.com/keras-team/keras-nlp",
+    ...         paper_url="https://arxiv.org/abs/2304.12244",
     ...         docs_url="https://keras.io/keras_nlp/",
     ...         # ^ optional metadata to generate model card
     ...     ):

tests/test_hub_mixin.py

@@ -91,7 +91,12 @@ def _from_pretrained(
         return cls(**kwargs)
 
 
-class BaseModelForInheritance(ModelHubMixin, repo_url="https://hf.co/my-repo", library_name="my-cool-library"):
+class BaseModelForInheritance(
+    ModelHubMixin,
+    repo_url="https://hf.co/my-repo",
+    paper_url="https://arxiv.org/abs/2304.12244",
+    library_name="my-cool-library",
+):
     pass
 
 
@@ -452,6 +457,7 @@ def test_inherited_class(self):
         """Test MixinInfo attributes are inherited from the parent class."""
         model = DummyModelInherited()
         assert model._hub_mixin_info.repo_url == "https://hf.co/my-repo"
+        assert model._hub_mixin_info.paper_url == "https://arxiv.org/abs/2304.12244"
         assert model._hub_mixin_info.model_card_data.library_name == "my-cool-library"
 
     def test_autocomplete_works_as_expected(self):

tests/test_hub_mixin_pytorch.py

@@ -36,9 +36,12 @@
 ---
 
 This is a dummy model card with kwargs.
-Arxiv ID: 1234.56789
 
 {{ custom_data }}
+
+- Code: {{ repo_url }}
+- Paper: {{ paper_url }}
+- Docs: {{ docs_url }}
 """
 
 if is_torch_available():
@@ -93,6 +96,9 @@ class DummyModelWithModelCardAndCustomKwargs(
         nn.Module,
         PyTorchModelHubMixin,
         model_card_template=DUMMY_MODEL_CARD_TEMPLATE_WITH_CUSTOM_KWARGS,
+        docs_url="https://hf.co/docs/my-repo",
+        paper_url="https://arxiv.org/abs/2304.12244",
+        repo_url="https://hf.co/my-repo",
     ):
         def __init__(self, linear_layer: int = 4):
             super().__init__()
@@ -331,7 +337,6 @@ def test_generate_model_card(self):
         assert card.data.license == "apache-2.0"
         assert card.data.pipeline_tag == "text-classification"
         assert card.data.tags == ["model_hub_mixin", "pytorch_model_hub_mixin", "tag1", "tag2"]
-
         # Model card template has been used
         assert "This is a dummy model card" in str(card)
 
@@ -438,7 +443,9 @@ def test_model_card_with_custom_kwargs(self):
         model = DummyModelWithModelCardAndCustomKwargs()
         card = model.generate_model_card(**model_card_kwargs)
         assert model_card_kwargs["custom_data"] in str(card)
-
+        assert "Code: https://hf.co/my-repo" in str(card)
+        assert "Paper: https://arxiv.org/abs/2304.12244" in str(card)
+        assert "Docs: https://hf.co/docs/my-repo" in str(card)
         # Test saving card => model card is saved and restored with custom data
         model.save_pretrained(self.cache_dir, model_card_kwargs=model_card_kwargs)
         card_reloaded = ModelCard.load(self.cache_dir / "README.md")