Repository: huggingface/huggingface_hub

Issue:
[Feature request] Add a hook to allow user to manipulate model init kwargs
The current `from_pretrained` implementation in `ModelHubMixin` reads in `config.json` (if available), and adds it as a kwarg named `'config'` in the model init kwargs. This is in line with `transformers` style, more or less.

I'd rather explicitly write config kwargs in my model's `__init__`. To do that, we have to unpack the config dict instead of having it be a kwarg. The solution right now is to basically copy paste 99% of `ModelHubMixin.from_pretrained` to update one line of code.

---

What do you think about adding a simple hook that wraps this line and allows users to manipulate `model_kwargs` if they'd like by overriding?

https://github.com/huggingface/huggingface_hub/blob/379414dade73e0816a996d83ead7e4ffd6986a53/src/huggingface_hub/hub_mixin.py#L196


Discussion:
**TL;DR:** Sounds good to me, do you want to give it a try?

---

Longer-term I think we should consider making `ModelHubMixin` more agnostic and **require** the user to define a `_save_pretrained`  and a `_from_pretrained`  method for it to work.

And then we would implement a `TorchModelHubMixin` , `ScikitLearnModelHubMixin` , etc. that are framework-specific and slightly more opinionated (and then in your case you can either tweak `TorchModelHubMixin`, or just inherit from the general ModelHubMixin and implement your logic the way you want)
Sure, I'll give it a go!

Resolving patch:
diff --git a/src/huggingface_hub/__init__.py b/src/huggingface_hub/__init__.py
--- a/src/huggingface_hub/__init__.py
+++ b/src/huggingface_hub/__init__.py
@@ -31,7 +31,7 @@
 )
 from .file_download import cached_download, hf_hub_download, hf_hub_url
 from .hf_api import HfApi, HfFolder, repo_type_and_id_from_hf_id
-from .hub_mixin import ModelHubMixin
+from .hub_mixin import ModelHubMixin, PyTorchModelHubMixin
 from .inference_api import InferenceApi
 from .repository import Repository
 from .snapshot_download import snapshot_download
diff --git a/src/huggingface_hub/hub_mixin.py b/src/huggingface_hub/hub_mixin.py
--- a/src/huggingface_hub/hub_mixin.py
+++ b/src/huggingface_hub/hub_mixin.py
@@ -20,29 +20,11 @@
 
 
 class ModelHubMixin:
-    def __init__(self, *args, **kwargs):
-        """
-        Mix this class with your torch-model class for ease process of saving & loading from huggingface-hub
-
-        Example::
-
-            >>> from huggingface_hub import ModelHubMixin
-
-            >>> class MyModel(nn.Module, ModelHubMixin):
-            ...    def __init__(self, **kwargs):
-            ...        super().__init__()
-            ...        self.config = kwargs.pop("config", None)
-            ...        self.layer = ...
-            ...    def forward(self, ...)
-            ...        return ...
-
-            >>> model = MyModel()
-            >>> model.save_pretrained("mymodel", push_to_hub=False) # Saving model weights in the directory
-            >>> model.push_to_hub("mymodel", "model-1") # Pushing model-weights to hf-hub
-
-            >>> # Downloading weights from hf-hub & model will be initialized from those weights
-            >>> model = MyModel.from_pretrained("username/mymodel@main")
-        """
+    """
+    A Generic Base Model Hub Mixin. Define your own mixin for anything by inheriting from this class
+    and overwriting _from_pretrained and _save_pretrained to define custom logic for saving/loading
+    your classes. See ``huggingface_hub.PyTorchModelHubMixin`` for an example.
+    """
 
     def save_pretrained(
         self,
@@ -75,26 +57,22 @@ def save_pretrained(
             with open(path, "w") as f:
                 json.dump(config, f)
 
-        # saving model weights
-        path = os.path.join(save_directory, PYTORCH_WEIGHTS_NAME)
-        self._save_pretrained(path)
+        # saving model weights/files
+        self._save_pretrained(save_directory)
 
         if push_to_hub:
             return self.push_to_hub(save_directory, **kwargs)
 
-    def _save_pretrained(self, path):
+    def _save_pretrained(self, save_directory):
         """
-        Overwrite this method in case you don't want to save complete model, rather some specific layers
+        Overwrite this method in subclass to define how to save your model.
         """
-        model_to_save = self.module if hasattr(self, "module") else self
-        torch.save(model_to_save.state_dict(), path)
+        raise NotImplementedError
 
     @classmethod
     def from_pretrained(
         cls,
         pretrained_model_name_or_path: Optional[str],
-        strict: bool = True,
-        map_location: Optional[str] = "cpu",
         force_download: bool = False,
         resume_download: bool = False,
         proxies: Dict = None,
@@ -145,7 +123,6 @@ def from_pretrained(
         """
 
         model_id = pretrained_model_name_or_path
-        map_location = torch.device(map_location)
 
         revision = None
         if len(model_id.split("@")) == 2:
@@ -171,35 +148,38 @@ def from_pretrained(
                 logger.warning("config.json NOT FOUND in HuggingFace Hub")
                 config_file = None
 
-        if os.path.isdir(model_id):
-            print("LOADING weights from local directory")
-            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
-        else:
-            model_url = hf_hub_url(
-                model_id, filename=PYTORCH_WEIGHTS_NAME, revision=revision
-            )
-            model_file = cached_download(
-                model_url,
-                cache_dir=cache_dir,
-                force_download=force_download,
-                proxies=proxies,
-                resume_download=resume_download,
-                local_files_only=local_files_only,
-                use_auth_token=use_auth_token,
-            )
-
         if config_file is not None:
             with open(config_file, "r", encoding="utf-8") as f:
                 config = json.load(f)
             model_kwargs.update({"config": config})
 
-        model = cls(**model_kwargs)
-
-        state_dict = torch.load(model_file, map_location=map_location)
-        model.load_state_dict(state_dict, strict=strict)
-        model.eval()
+        return cls._from_pretrained(
+            model_id,
+            revision,
+            cache_dir,
+            force_download,
+            proxies,
+            resume_download,
+            local_files_only,
+            use_auth_token,
+            **model_kwargs,
+        )
 
-        return model
+    @classmethod
+    def _from_pretrained(
+        cls,
+        model_id,
+        revision,
+        cache_dir,
+        force_download,
+        proxies,
+        resume_download,
+        local_files_only,
+        use_auth_token,
+        **model_kwargs,
+    ):
+        """Overwrite this method in subclass to define how to load your model from pretrained"""
+        raise NotImplementedError
 
     def push_to_hub(
         self,
@@ -302,3 +282,82 @@ def push_to_hub(
         repo.git_add()
         repo.git_commit(commit_message)
         return repo.git_push()
+
+
+class PyTorchModelHubMixin(ModelHubMixin):
+    def __init__(self, *args, **kwargs):
+        """
+        Mix this class with your torch-model class for ease process of saving & loading from huggingface-hub
+
+        Example::
+
+            >>> from huggingface_hub import PyTorchModelHubMixin
+
+            >>> class MyModel(nn.Module, PyTorchModelHubMixin):
+            ...    def __init__(self, **kwargs):
+            ...        super().__init__()
+            ...        self.config = kwargs.pop("config", None)
+            ...        self.layer = ...
+            ...    def forward(self, ...)
+            ...        return ...
+
+            >>> model = MyModel()
+            >>> model.save_pretrained("mymodel", push_to_hub=False) # Saving model weights in the directory
+            >>> model.push_to_hub("mymodel", "model-1") # Pushing model-weights to hf-hub
+
+            >>> # Downloading weights from hf-hub & model will be initialized from those weights
+            >>> model = MyModel.from_pretrained("username/mymodel@main")
+        """
+
+    def _save_pretrained(self, save_directory):
+        """
+        Overwrite this method in case you don't want to save complete model, rather some specific layers
+        """
+        path = os.path.join(save_directory, PYTORCH_WEIGHTS_NAME)
+        model_to_save = self.module if hasattr(self, "module") else self
+        torch.save(model_to_save.state_dict(), path)
+
+    @classmethod
+    def _from_pretrained(
+        cls,
+        model_id,
+        revision,
+        cache_dir,
+        force_download,
+        proxies,
+        resume_download,
+        local_files_only,
+        use_auth_token,
+        map_location="cpu",
+        strict=False,
+        **model_kwargs,
+    ):
+        """
+        Overwrite this method in case you wish to initialize your model in a different way.
+        """
+        map_location = torch.device(map_location)
+
+        if os.path.isdir(model_id):
+            print("LOADING weights from local directory")
+            model_file = os.path.join(model_id, PYTORCH_WEIGHTS_NAME)
+        else:
+            model_url = hf_hub_url(
+                model_id, filename=PYTORCH_WEIGHTS_NAME, revision=revision
+            )
+            model_file = cached_download(
+                model_url,
+                cache_dir=cache_dir,
+                force_download=force_download,
+                proxies=proxies,
+                resume_download=resume_download,
+                local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
+            )
+
+        model = cls(**model_kwargs)
+
+        state_dict = torch.load(model_file, map_location=map_location)
+        model.load_state_dict(state_dict, strict=strict)
+        model.eval()
+
+        return model


Test patch:
diff --git a/tests/test_hubmixin.py b/tests/test_hubmixin.py
--- a/tests/test_hubmixin.py
+++ b/tests/test_hubmixin.py
@@ -5,7 +5,7 @@
 
 from huggingface_hub import HfApi
 from huggingface_hub.file_download import is_torch_available
-from huggingface_hub.hub_mixin import ModelHubMixin
+from huggingface_hub.hub_mixin import PyTorchModelHubMixin
 
 from .testing_constants import ENDPOINT_STAGING, PASS, USER
 from .testing_utils import set_write_permission_and_retry
@@ -37,7 +37,7 @@ def require_torch(test_case):
 
 if is_torch_available():
 
-    class DummyModel(nn.Module, ModelHubMixin):
+    class DummyModel(nn.Module, PyTorchModelHubMixin):
         def __init__(self, **kwargs):
             super().__init__()
             self.config = kwargs.pop("config", None)