Fix infinity in JSON serialized files (#42959)

* Handle inifinity and NaNs in JSON serialization

* Docs

* Tests

Author: LysandreJik

src/transformers/configuration_utils.py

@@ -16,6 +16,7 @@
 
 import copy
 import json
+import math
 import os
 import warnings
 from typing import TYPE_CHECKING, Any, Optional, TypeVar, Union
@@ -50,6 +51,9 @@
 # type hinting: specifying the type of config class that inherits from PreTrainedConfig
 SpecificPreTrainedConfigType = TypeVar("SpecificPreTrainedConfigType", bound="PreTrainedConfig")
 
+_FLOAT_TAG_KEY = "__float__"
+_FLOAT_TAG_VALUES = {"Infinity": float("inf"), "-Infinity": float("-inf"), "NaN": float("nan")}
+
 
 class PreTrainedConfig(PushToHubMixin, RotaryEmbeddingConfigMixin):
     # no-format
@@ -812,7 +816,56 @@ def from_json_file(
     def _dict_from_json_file(cls, json_file: str | os.PathLike):
         with open(json_file, encoding="utf-8") as reader:
             text = reader.read()
-        return json.loads(text)
+        config_dict = json.loads(text)
+
+        return cls._decode_special_floats(config_dict)
+
+    @classmethod
+    def _encode_special_floats(cls, obj: Any) -> Any:
+        """
+        Iterates over the passed object and encode specific floats that cannot be JSON-serialized. Python's JSON
+        engine saves floats like `Infinity` (+/-) or `NaN` which are not compatible with other JSON engines.
+
+        It serializes floats like `Infinity` as an object: `{'__float__': Infinity}`.
+        """
+        if isinstance(obj, float):
+            if math.isnan(obj):
+                return {_FLOAT_TAG_KEY: "NaN"}
+            if obj == float("inf"):
+                return {_FLOAT_TAG_KEY: "Infinity"}
+            if obj == float("-inf"):
+                return {_FLOAT_TAG_KEY: "-Infinity"}
+            return obj
+
+        if isinstance(obj, dict):
+            return {k: cls._encode_special_floats(v) for k, v in obj.items()}
+
+        if isinstance(obj, (list, tuple)):
+            return [cls._encode_special_floats(v) for v in obj]
+
+        return obj
+
+    @classmethod
+    def _decode_special_floats(cls, obj: Any) -> Any:
+        """
+        Iterates over the passed object and decode specific floats that cannot be JSON-serialized. Python's JSON
+        engine saves floats like `Infinity` (+/-) or `NaN` which are not compatible with other JSON engines.
+
+        This method deserializes objects like `{'__float__': Infinity}` to their float values like `Infinity`.
+        """
+        if isinstance(obj, dict):
+            if set(obj.keys()) == {_FLOAT_TAG_KEY} and isinstance(obj[_FLOAT_TAG_KEY], str):
+                tag = obj[_FLOAT_TAG_KEY]
+                if tag in _FLOAT_TAG_VALUES:
+                    return _FLOAT_TAG_VALUES[tag]
+                return obj
+
+            return {k: cls._decode_special_floats(v) for k, v in obj.items()}
+
+        if isinstance(obj, list):
+            return [cls._decode_special_floats(v) for v in obj]
+
+        return obj
 
     def __eq__(self, other):
         return isinstance(other, PreTrainedConfig) and (self.__dict__ == other.__dict__)
@@ -932,6 +985,10 @@ def to_json_string(self, use_diff: bool = True) -> str:
             config_dict = self.to_diff_dict()
         else:
             config_dict = self.to_dict()
+
+        # Handle +/-Infinity and NaNs
+        config_dict = self._encode_special_floats(config_dict)
+
         return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
 
     def to_json_file(self, json_file_path: str | os.PathLike, use_diff: bool = True):

tests/utils/test_configuration_utils.py

@@ -329,3 +329,44 @@ def test_bc_torch_dtype(self):
 
             config = PreTrainedConfig.from_pretrained(tmpdirname, torch_dtype="float32")
             self.assertEqual(config.dtype, "float32")
+
+    def test_unserializable_json_is_encoded(self):
+        class NewConfig(PreTrainedConfig):
+            def __init__(
+                self,
+                inf_positive: float = float("inf"),
+                inf_negative: float = float("-inf"),
+                nan: float = float("nan"),
+                **kwargs,
+            ):
+                self.inf_positive = inf_positive
+                self.inf_negative = inf_negative
+                self.nan = nan
+
+                super().__init__(**kwargs)
+
+        new_config = NewConfig()
+
+        # All floats should remain as floats when being accessed in the config
+        self.assertIsInstance(new_config.inf_positive, float)
+        self.assertIsInstance(new_config.inf_negative, float)
+        self.assertIsInstance(new_config.nan, float)
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            new_config.save_pretrained(tmpdirname)
+            config_file = Path(tmpdirname) / "config.json"
+            config_contents = json.loads(config_file.read_text())
+            new_config_instance = NewConfig.from_pretrained(tmpdirname)
+
+        # In the serialized JSON file, the non-JSON compatible floats should be updated
+        self.assertDictEqual(config_contents["inf_positive"], {"__float__": "Infinity"})
+        self.assertDictEqual(config_contents["inf_negative"], {"__float__": "-Infinity"})
+        self.assertDictEqual(config_contents["nan"], {"__float__": "NaN"})
+
+        with tempfile.TemporaryDirectory() as tmpdirname:
+            new_config.save_pretrained(tmpdirname)
+
+        # When reloading the config, it should have correct float values
+        self.assertIsInstance(new_config_instance.inf_positive, float)
+        self.assertIsInstance(new_config_instance.inf_negative, float)
+        self.assertIsInstance(new_config_instance.nan, float)