Repo metadata load and save (#339)

* repocard utilities + tests

* mix those methods on Repository

* handle case of no README.md inside repo

* suggestion from @osanseviero

Co-authored-by: Omar Sanseviero <[email protected]>

* tiny doc add

* `MODELCARD_NAME` => more general `REPOCARD_NAME`

* Additional test

Co-authored-by: Omar Sanseviero <[email protected]>
Co-authored-by: Lysandre <[email protected]>

setup.py

@@ -15,6 +15,7 @@ def get_version() -> str:
     "filelock",
     "requests",
     "tqdm",
+    "pyyaml",
     "typing-extensions",
     "importlib_metadata;python_version<'3.8'",
     "packaging>=20.9",

src/huggingface_hub/constants.py

@@ -13,6 +13,7 @@
 TF_WEIGHTS_NAME = "model.ckpt"
 FLAX_WEIGHTS_NAME = "flax_model.msgpack"
 CONFIG_NAME = "config.json"
+REPOCARD_NAME = "README.md"
 
 HUGGINGFACE_CO_URL_HOME = "https://huggingface.co/"
 

src/huggingface_hub/repocard.py

@@ -0,0 +1,38 @@
+import re
+from pathlib import Path
+from typing import Dict, Optional, Union
+
+import yaml
+
+
+# exact same regex as in the Hub server. Please keep in sync.
+REGEX_YAML_BLOCK = re.compile(r"---[\n\r]+([\S\s]*?)[\n\r]+---[\n\r]")
+
+
+def metadata_load(local_path: Union[str, Path]) -> Optional[Dict]:
+    content = Path(local_path).read_text()
+    match = REGEX_YAML_BLOCK.search(content)
+    if match:
+        yaml_block = match.group(1)
+        data = yaml.safe_load(yaml_block)
+        if isinstance(data, dict):
+            return data
+        else:
+            raise ValueError("repo card metadata block should be a dict")
+    else:
+        return None
+
+
+def metadata_save(local_path: Union[str, Path], data: Dict) -> None:
+    data_yaml = yaml.dump(data, sort_keys=False)
+    # sort_keys: keep dict order
+    content = Path(local_path).read_text() if Path(local_path).is_file() else ""
+    match = REGEX_YAML_BLOCK.search(content)
+    if match:
+        output = (
+            content[: match.start()] + f"---\n{data_yaml}---\n" + content[match.end() :]
+        )
+    else:
+        output = f"---\n{data_yaml}---\n{content}"
+
+    Path(local_path).write_text(output)

src/huggingface_hub/repository.py

@@ -7,11 +7,12 @@
 import time
 from contextlib import contextmanager
 from pathlib import Path
-from typing import Callable, Iterator, List, Optional, Tuple, Union
+from typing import Callable, Dict, Iterator, List, Optional, Tuple, Union
 
 from tqdm.auto import tqdm
 
-from huggingface_hub.constants import REPO_TYPES_URL_PREFIXES
+from huggingface_hub.constants import REPO_TYPES_URL_PREFIXES, REPOCARD_NAME
+from huggingface_hub.repocard import metadata_load, metadata_save
 
 from .hf_api import ENDPOINT, HfApi, HfFolder, repo_type_and_id_from_hf_id
 from .lfs import LFS_MULTIPART_UPLOAD_COMMAND
@@ -1141,6 +1142,14 @@ def commit(
 
             os.chdir(current_working_directory)
 
+    def repocard_metadata_load(self) -> Optional[Dict]:
+        filepath = os.path.join(self.local_dir, REPOCARD_NAME)
+        if os.path.isfile(filepath):
+            return metadata_load(filepath)
+
+    def repocard_metadata_save(self, data: Dict) -> None:
+        return metadata_save(os.path.join(self.local_dir, REPOCARD_NAME), data)
+
     @property
     def commands_failed(self):
         return [c for c in self.command_queue if c.status > 0]

tests/test_repocard.py

@@ -0,0 +1,102 @@
+# Copyright 2021 The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+import os
+import shutil
+import unittest
+from pathlib import Path
+
+from huggingface_hub.constants import REPOCARD_NAME
+from huggingface_hub.repocard import metadata_load, metadata_save
+
+from .testing_utils import set_write_permission_and_retry
+
+
+DUMMY_MODELCARD = """
+
+Hi
+
+---
+license: mit
+datasets:
+- foo
+- bar
+---
+
+Hello
+"""
+
+DUMMY_MODELCARD_TARGET = """
+
+Hi
+
+---
+meaning_of_life: 42
+---
+
+Hello
+"""
+
+DUMMY_MODELCARD_TARGET_NO_YAML = """---
+meaning_of_life: 42
+---
+Hello
+"""
+
+DUMMY_MODELCARD_TARGET_NO_TAGS = """
+Hello
+"""
+
+REPOCARD_DIR = os.path.join(
+    os.path.dirname(os.path.abspath(__file__)), "fixtures/repocard"
+)
+
+
+class RepocardTest(unittest.TestCase):
+    def setUp(self):
+        os.makedirs(REPOCARD_DIR, exist_ok=True)
+
+    def tearDown(self) -> None:
+        try:
+            shutil.rmtree(REPOCARD_DIR, onerror=set_write_permission_and_retry)
+        except FileNotFoundError:
+            pass
+
+    def test_metadata_load(self):
+        filepath = Path(REPOCARD_DIR) / REPOCARD_NAME
+        filepath.write_text(DUMMY_MODELCARD)
+        data = metadata_load(filepath)
+        self.assertDictEqual(data, {"license": "mit", "datasets": ["foo", "bar"]})
+
+    def test_metadata_save(self):
+        filename = "dummy_target.md"
+        filepath = Path(REPOCARD_DIR) / filename
+        filepath.write_text(DUMMY_MODELCARD)
+        metadata_save(filepath, {"meaning_of_life": 42})
+        content = filepath.read_text()
+        self.assertEqual(content, DUMMY_MODELCARD_TARGET)
+
+    def test_metadata_save_from_file_no_yaml(self):
+        filename = "dummy_target_2.md"
+        filepath = Path(REPOCARD_DIR) / filename
+        filepath.write_text("Hello\n")
+        metadata_save(filepath, {"meaning_of_life": 42})
+        content = filepath.read_text()
+        self.assertEqual(content, DUMMY_MODELCARD_TARGET_NO_YAML)
+
+    def test_no_metadata_returns_none(self):
+        filename = "dummy_target_3.md"
+        filepath = Path(REPOCARD_DIR) / filename
+        filepath.write_text(DUMMY_MODELCARD_TARGET_NO_TAGS)
+        data = metadata_load(filepath)
+        self.assertEqual(data, None)