From 4f4f40baa69767db2fa272cf7ac037485514bd13 Mon Sep 17 00:00:00 2001 From: Ross Wightman Date: Fri, 27 Dec 2024 12:04:04 -0800 Subject: [PATCH] Add support for tag, license customization through push_to_hub --- timm/models/_hub.py | 12 ++++++++++-- 1 file changed, 10 insertions(+), 2 deletions(-) diff --git a/timm/models/_hub.py b/timm/models/_hub.py index 4922dfc0..809705ed 100644 --- a/timm/models/_hub.py +++ b/timm/models/_hub.py @@ -395,10 +395,18 @@ def push_to_hf_hub( def generate_readme(model_card: dict, model_name: str): + tags = model_card.get('tags', None) or ['image-classification', 'timm'] readme_text = "---\n" - readme_text += "tags:\n- image-classification\n- timm\n" - readme_text += "library_name: timm\n" + if tags: + readme_text += "tags:\n" + for t in tags: + readme_text += f"- {t}\n" + readme_text += f"library_name: {model_card.get('library_name', 'timm')}\n" readme_text += f"license: {model_card.get('license', 'apache-2.0')}\n" + if 'license_name' in model_card: + readme_text += f"license_name: {model_card.get('license_name')}\n" + if 'license_link' in model_card: + readme_text += f"license_link: {model_card.get('license_link')}\n" if 'details' in model_card and 'Dataset' in model_card['details']: readme_text += 'datasets:\n' if isinstance(model_card['details']['Dataset'], (tuple, list)):