py: let users add full base model and dataset to model_card

if they choose to. Also added more tests for this feature.

Author: mofosyne

gguf-py/gguf/metadata.py

@@ -367,19 +367,24 @@ def use_array_model_card_metadata(metadata_key: str, model_card_key: str):
                 for model_id in metadata_base_models:
                     # NOTE: model size of base model is assumed to be similar to the size of the current model
                     base_model = {}
-                    if isinstance(model_id, str) and (model_id.startswith("http://") or model_id.startswith("https://")):
-                        base_model["repo_url"] = model_id
+                    if isinstance(model_id, str):
+                        if model_id.startswith("http://") or model_id.startswith("https://") or model_id.startswith("ssh://"):
+                            base_model["repo_url"] = model_id
+                        else:
+                            # Likely a Hugging Face ID
+                            model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params)
+                            if model_full_name_component is not None:
+                                base_model["name"] = Metadata.id_to_title(model_full_name_component)
+                            if org_component is not None:
+                                base_model["organization"] = Metadata.id_to_title(org_component)
+                            if version is not None:
+                                base_model["version"] = version
+                            if org_component is not None and model_full_name_component is not None:
+                                base_model["repo_url"] = f"https://huggingface.co/{org_component}/{model_full_name_component}"
+                    elif isinstance(model_id, dict):
+                        base_model = model_id
                     else:
-                        # Likely a Hugging Face ID
-                        model_full_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(model_id, total_params)
-                        if model_full_name_component is not None:
-                            base_model["name"] = Metadata.id_to_title(model_full_name_component)
-                        if org_component is not None:
-                            base_model["organization"] = Metadata.id_to_title(org_component)
-                        if version is not None:
-                            base_model["version"] = version
-                        if org_component is not None and model_full_name_component is not None:
-                            base_model["repo_url"] = f"https://huggingface.co/{org_component}/{model_full_name_component}"
+                        logger.error(f"base model entry '{str(model_id)}' not in a known format")
                     metadata.base_models.append(base_model)
 
             if "datasets" in model_card or "dataset" in model_card:
@@ -399,19 +404,24 @@ def use_array_model_card_metadata(metadata_key: str, model_card_key: str):
                 for dataset_id in metadata_datasets:
                     # NOTE: model size of base model is assumed to be similar to the size of the current model
                     dataset = {}
-                    if isinstance(dataset_id, str) and (dataset_id.startswith("http://") or dataset_id.startswith("https://")):
-                        dataset["repo_url"] = dataset_id
+                    if isinstance(dataset_id, str):
+                        if dataset_id.startswith("http://") or dataset_id.startswith("https://") or dataset_id.startswith("ssh://"):
+                            dataset["repo_url"] = dataset_id
+                        else:
+                            # Likely a Hugging Face ID
+                            dataset_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(dataset_id, total_params)
+                            if dataset_name_component is not None:
+                                dataset["name"] = Metadata.id_to_title(dataset_name_component)
+                            if org_component is not None:
+                                dataset["organization"] = Metadata.id_to_title(org_component)
+                            if version is not None:
+                                dataset["version"] = version
+                            if org_component is not None and dataset_name_component is not None:
+                                dataset["repo_url"] = f"https://huggingface.co/{org_component}/{dataset_name_component}"
+                    elif isinstance(dataset_id, dict):
+                        dataset = dataset_id
                     else:
-                        # Likely a Hugging Face ID
-                        dataset_name_component, org_component, basename, finetune, version, size_label = Metadata.get_model_id_components(dataset_id, total_params)
-                        if dataset_name_component is not None:
-                            dataset["name"] = Metadata.id_to_title(dataset_name_component)
-                        if org_component is not None:
-                            dataset["organization"] = Metadata.id_to_title(org_component)
-                        if version is not None:
-                            dataset["version"] = version
-                        if org_component is not None and dataset_name_component is not None:
-                            dataset["repo_url"] = f"https://huggingface.co/{org_component}/{dataset_name_component}"
+                        logger.error(f"dataset entry '{str(dataset_id)}' not in a known format")
                     metadata.datasets.append(dataset)
 
             use_model_card_metadata("license", "license")

gguf-py/tests/test_metadata.py

@@ -183,7 +183,42 @@ def test_apply_metadata_heuristic_from_model_card(self):
         expect.tags=['Llama-3', 'instruct', 'finetune', 'chatml', 'DPO', 'RLHF', 'gpt4', 'synthetic data', 'distillation', 'function calling', 'json mode', 'axolotl']
         expect.languages=['en']
         expect.datasets=[{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}]
+        self.assertEqual(got, expect)
 
+        # Base Model spec is inferred from model id
+        model_card = {'base_models': ['teknium/OpenHermes-2.5']}
+        expect = gguf.Metadata(base_models=[{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}])
+        got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
+        self.assertEqual(got, expect)
+
+        # Base Model spec is only url
+        model_card = {'base_models': ['https://huggingface.co/teknium/OpenHermes-2.5']}
+        expect = gguf.Metadata(base_models=[{'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}])
+        got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
+        self.assertEqual(got, expect)
+
+        # Base Model spec is given directly
+        model_card = {'base_models': [{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}]}
+        expect = gguf.Metadata(base_models=[{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}])
+        got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
+        self.assertEqual(got, expect)
+
+        # Dataset spec is inferred from model id
+        model_card = {'datasets': ['teknium/OpenHermes-2.5']}
+        expect = gguf.Metadata(datasets=[{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}])
+        got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
+        self.assertEqual(got, expect)
+
+        # Dataset spec is only url
+        model_card = {'datasets': ['https://huggingface.co/teknium/OpenHermes-2.5']}
+        expect = gguf.Metadata(datasets=[{'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}])
+        got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
+        self.assertEqual(got, expect)
+
+        # Dataset spec is given directly
+        model_card = {'datasets': [{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}]}
+        expect = gguf.Metadata(datasets=[{'name': 'OpenHermes 2.5', 'organization': 'Teknium', 'version': '2.5', 'repo_url': 'https://huggingface.co/teknium/OpenHermes-2.5'}])
+        got = gguf.Metadata.apply_metadata_heuristic(gguf.Metadata(), model_card, None, None)
         self.assertEqual(got, expect)
 
     def test_apply_metadata_heuristic_from_hf_parameters(self):