Fix: correctly encode/decode config in ModelHubMixin if custom coders (#2337)

* Fix: correctly encode/decode config in ModelHubMixin if custom coders

* make style

* make quality

* Update tests/test_hub_mixin_pytorch.py

Co-authored-by: Hafedh <[email protected]>

---------

Co-authored-by: Hafedh <[email protected]>

Author: Wauplin

docs/source/ko/guides/integrations.md

@@ -365,9 +365,9 @@ from argparse import Namespace
 
 class VoiceCraft(
    nn.Module,
-   PytorchModelHubMixin,  # 믹스인을 상속합니다.
-   coders: {
-      Namespace = (
+   PyTorchModelHubMixin,  # 믹스인을 상속합니다.
+   coders={
+      Namespace: (
          lambda x: vars(x),  # Encoder: `Namespace`를 유효한 JSON 형태로 변환하는 방법은 무엇인가요?
          lambda data: Namespace(**data),  # Decoder: 딕셔너리에서 Namespace를 재구성하는 방법은 무엇인가요?
       )

src/huggingface_hub/hub_mixin.py

@@ -15,7 +15,6 @@
     Type,
     TypeVar,
     Union,
-    get_args,
 )
 
 from .constants import CONFIG_NAME, PYTORCH_WEIGHTS_NAME, SAFETENSORS_SINGLE_FILE
@@ -326,12 +325,11 @@ def __new__(cls, *args, **kwargs) -> "ModelHubMixin":
                 if instance._is_jsonable(value)  # Only if jsonable or we have a custom encoder
             },
         }
-        init_config.pop("config", {})
+        passed_config = init_config.pop("config", {})
 
         # Populate `init_config` with provided config
-        provided_config = passed_values.get("config")
-        if isinstance(provided_config, dict):
-            init_config.update(provided_config)
+        if isinstance(passed_config, dict):
+            init_config.update(passed_config)
 
         # Set `config` attribute and return
         if init_config != {}:
@@ -362,9 +360,14 @@ def _decode_arg(cls, expected_type: Type[ARGS_T], value: Any) -> Optional[ARGS_T
             if value is None:
                 return None
             expected_type = unwrap_simple_optional_type(expected_type)
+        # Dataclass => handle it
+        if is_dataclass(expected_type):
+            return _load_dataclass(expected_type, value)  # type: ignore[return-value]
+        # Otherwise => check custom decoders
         for type_, (_, decoder) in cls._hub_mixin_coders.items():
             if inspect.isclass(expected_type) and issubclass(expected_type, type_):
                 return decoder(value)
+        # Otherwise => don't decode
         return value
 
     def save_pretrained(
@@ -531,18 +534,9 @@ def from_pretrained(
 
             # Check if `config` argument was passed at init
             if "config" in cls._hub_mixin_init_parameters and "config" not in model_kwargs:
-                # Check if `config` argument is a dataclass
+                # Decode `config` argument if it was passed
                 config_annotation = cls._hub_mixin_init_parameters["config"].annotation
-                if config_annotation is inspect.Parameter.empty:
-                    pass  # no annotation
-                elif is_dataclass(config_annotation):
-                    config = _load_dataclass(config_annotation, config)
-                else:
-                    # if Optional/Union annotation => check if a dataclass is in the Union
-                    for _sub_annotation in get_args(config_annotation):
-                        if is_dataclass(_sub_annotation):
-                            config = _load_dataclass(_sub_annotation, config)
-                            break
+                config = cls._decode_arg(config_annotation, config)
 
                 # Forward config to model initialization
                 model_kwargs["config"] = config

tests/test_hub_mixin_pytorch.py

@@ -2,6 +2,7 @@
 import os
 import struct
 import unittest
+from argparse import Namespace
 from pathlib import Path
 from typing import Any, Dict, Optional, TypeVar
 from unittest.mock import Mock, patch
@@ -95,6 +96,21 @@ class DummyModelWithModelCardAndCustomKwargs(
         def __init__(self, linear_layer: int = 4):
             super().__init__()
 
+    class DummyModelWithEncodedConfig(
+        nn.Module,
+        PyTorchModelHubMixin,
+        coders={
+            Namespace: (
+                lambda x: vars(x),
+                lambda data: Namespace(**data),
+            )
+        },
+    ):
+        # Regression test for https://github.com/huggingface/huggingface_hub/issues/2334
+        def __init__(self, config: Namespace):
+            super().__init__()
+            self.config = config
+
 else:
     DummyModel = None
     DummyModelWithModelCard = None
@@ -419,3 +435,19 @@ def test_model_card_with_custom_kwargs(self):
         model.save_pretrained(self.cache_dir, model_card_kwargs=model_card_kwargs)
         card_reloaded = ModelCard.load(self.cache_dir / "README.md")
         assert str(card) == str(card_reloaded)
+
+    def test_config_with_custom_coders(self):
+        """
+        Regression test for #2334. When `config` is encoded with custom coders, it should be decoded correctly.
+
+        See https://github.com/huggingface/huggingface_hub/issues/2334.
+        """
+        model = DummyModelWithEncodedConfig(Namespace(a=1, b=2))
+        model.save_pretrained(self.cache_dir)
+        assert model._hub_mixin_config["a"] == 1
+        assert model._hub_mixin_config["b"] == 2
+
+        reloaded = DummyModelWithEncodedConfig.from_pretrained(self.cache_dir)
+        assert isinstance(reloaded.config, Namespace)
+        assert reloaded.config.a == 1
+        assert reloaded.config.b == 2