Preserve card metadata format/ordering on load->save

huggingface · Oct 2, 2024 · a0f1e75 · a0f1e75
1 parent ce38f46
commit a0f1e75
Show file tree

Hide file tree

Showing 3 changed files with 15 additions and 4 deletions.
diff --git a/src/huggingface_hub/repocard.py b/src/huggingface_hub/repocard.py
@@ -83,7 +83,7 @@ def __init__(self, content: str, ignore_metadata_errors: bool = False):
     def content(self):
         """The content of the RepoCard, including the YAML block and the Markdown body."""
         line_break = _detect_line_ending(self._content) or "\n"
-        return f"---{line_break}{self.data.to_yaml(line_break=line_break)}{line_break}---{line_break}{self.text}"
+        return f"---{line_break}{self.data.to_yaml(line_break=line_break, original_order=self._original_order)}{line_break}---{line_break}{self.text}"
 
     @content.setter
     def content(self, content: str):
@@ -110,6 +110,7 @@ def content(self, content: str):
             self.text = content
 
         self.data = self.card_data_class(**data_dict, ignore_metadata_errors=self.ignore_metadata_errors)
+        self._original_order = list(data_dict.keys())
 
     def __str__(self):
         return self.content
@@ -330,7 +331,7 @@ def from_template(
         if template_str is None:
             template_str = Path(cls.default_template_path).read_text()
         template = jinja2.Template(template_str)
-        content = template.render(card_data=card_data.to_yaml(), **kwargs)
+        content = template.render(card_data=card_data.to_yaml(original_order=cls._original_order), **kwargs)
         return cls(content)
 
 

diff --git a/src/huggingface_hub/repocard_data.py b/src/huggingface_hub/repocard_data.py
@@ -175,7 +175,7 @@ class CardData:
     def __init__(self, ignore_metadata_errors: bool = False, **kwargs):
         self.__dict__.update(kwargs)
 
-    def to_dict(self) -> Dict[str, Any]:
+    def to_dict(self):
         """Converts CardData to a dict.
 
         Returns:
@@ -195,7 +195,7 @@ def _to_dict(self, data_dict):
         """
         pass
 
-    def to_yaml(self, line_break=None) -> str:
+    def to_yaml(self, line_break=None, original_order: Optional[List[str]] = None) -> str:
         """Dumps CardData to a YAML block for inclusion in a README.md file.
 
         Args:
@@ -205,6 +205,10 @@ def to_yaml(self, line_break=None) -> str:
         Returns:
             `str`: CardData represented as a YAML block.
         """
+        if original_order:
+            self.__dict__ = {
+                k: self.__dict__[k] for k in original_order + list(set(self.__dict__.keys()) - set(original_order))
+            }
         return yaml_dump(self.to_dict(), sort_keys=False, line_break=line_break).strip()
 
     def __repr__(self):

diff --git a/tests/test_repocard.py b/tests/test_repocard.py
@@ -847,6 +847,12 @@ def test_model_card_from_template_eval_results(self):
         self.assertTrue(card.data.to_dict().get("eval_results") is None)
         self.assertEqual(str(card)[: len(DUMMY_MODELCARD_EVAL_RESULT)], DUMMY_MODELCARD_EVAL_RESULT)
 
+    def test_preserve_order_load_save(self):
+        model_card = ModelCard(DUMMY_MODELCARD)
+        model_card.data.license = "test"
+        model_card.content
+        self.assertEqual(model_card.content, "---\nlicense: test\ndatasets:\n- foo\n- bar\n---\n\nHello\n")
+
 
 class DatasetCardTest(TestCaseWithHfApi):
     def test_load_datasetcard_from_file(self):