Repository: huggingface/huggingface_hub
Pull request #2159 (2024-03-26 10:44:09)

Fix ModelHubMixin when class is a dataclass
When the ModelHubMixin class happens to be a dataclass, `__init__` method will accept `**kwargs` when
inspecting it. However, due to how dataclasses work, we cannot forward arbitrary kwargs to the `__init__`.
This test ensures that the `from_pretrained` method does not raise an error when the class is a dataclass.

This PR fixes the SetFit integration (see https://github.com/huggingface/huggingface_hub/issues/2157). However, this is not a pattern we want to promote. Dataclasses are meant to store data. Model classes are more complex and doesn't play well with internal values.

HuggingFaceDocBuilderDev: The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/huggingface_hub/pr_2159). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.
Wauplin: Thanks for the review @LysandreJik. Will merge and make a patch release

Patch:
diff --git a/src/huggingface_hub/hub_mixin.py b/src/huggingface_hub/hub_mixin.py
index fa9ee5e7..eb42ef4e 100644
--- a/src/huggingface_hub/hub_mixin.py
+++ b/src/huggingface_hub/hub_mixin.py
@@ -403,7 +403,12 @@ class ModelHubMixin:
                 # Forward config to model initialization
                 model_kwargs["config"] = config
 
-            if any(param.kind == inspect.Parameter.VAR_KEYWORD for param in cls._hub_mixin_init_parameters.values()):
+            # Inject config if `**kwargs` are expected
+            if is_dataclass(cls):
+                for key in cls.__dataclass_fields__:
+                    if key not in model_kwargs and key in config:
+                        model_kwargs[key] = config[key]
+            elif any(param.kind == inspect.Parameter.VAR_KEYWORD for param in cls._hub_mixin_init_parameters.values()):
                 for key, value in config.items():
                     if key not in model_kwargs:
                         model_kwargs[key] = value


Test patch:
diff --git a/tests/test_hub_mixin.py b/tests/test_hub_mixin.py
index 09dd91c5..0047bdd4 100644
--- a/tests/test_hub_mixin.py
+++ b/tests/test_hub_mixin.py
@@ -99,6 +99,28 @@ class DummyModelSavingConfig(ModelHubMixin):
         (save_directory / "config.json").write_text(json.dumps({"custom_config": "custom_config"}))
 
 
+@dataclass
+class DummyModelThatIsAlsoADataclass(ModelHubMixin):
+    foo: int
+    bar: str
+
+    @classmethod
+    def _from_pretrained(
+        cls,
+        *,
+        model_id: str,
+        revision: Optional[str],
+        cache_dir: Optional[Union[str, Path]],
+        force_download: bool,
+        proxies: Optional[Dict],
+        resume_download: bool,
+        local_files_only: bool,
+        token: Optional[Union[str, bool]],
+        **model_kwargs,
+    ):
+        return cls(**model_kwargs)
+
+
 @pytest.mark.usefixtures("fx_cache_dir")
 class HubMixinTest(unittest.TestCase):
     cache_dir: Path
@@ -328,3 +350,18 @@ class HubMixinTest(unittest.TestCase):
         # config.json IS overwritten
         with open(self.cache_dir / "config.json") as f:
             assert json.load(f) == {"a": 1, "b": 2}
+
+    def test_from_pretrained_when_cls_is_a_dataclass(self):
+        """Regression test for #2157.
+
+        When the ModelHubMixin class happens to be a dataclass, `__init__` method will accept `**kwargs` when
+        inspecting it. However, due to how dataclasses work, we cannot forward arbitrary kwargs to the `__init__`.
+        This test ensures that the `from_pretrained` method does not raise an error when the class is a dataclass.
+
+        See https://github.com/huggingface/huggingface_hub/issues/2157.
+        """
+        (self.cache_dir / "config.json").write_text('{"foo": 42, "bar": "baz", "other": "value"}')
+        model = DummyModelThatIsAlsoADataclass.from_pretrained(self.cache_dir)
+        assert model.foo == 42
+        assert model.bar == "baz"
+        assert not hasattr(model, "other")