Skip to content

Commit 01cf0f7

Browse files
committed
Add support for tag, license customization through push_to_hub
1 parent b12ecbd commit 01cf0f7

File tree

1 file changed

+10
-2
lines changed

1 file changed

+10
-2
lines changed

timm/models/_hub.py

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -395,10 +395,18 @@ def push_to_hf_hub(
395395

396396

397397
def generate_readme(model_card: dict, model_name: str):
398+
tags = model_card.get('tags', None) or ['image-classification', 'timm']
398399
readme_text = "---\n"
399-
readme_text += "tags:\n- image-classification\n- timm\n"
400-
readme_text += "library_name: timm\n"
400+
if tags:
401+
readme_text += "tags:\n"
402+
for t in tags:
403+
readme_text += f"- {t}\n"
404+
readme_text += f"library_name: {model_card.get('library_name', 'timm')}\n"
401405
readme_text += f"license: {model_card.get('license', 'apache-2.0')}\n"
406+
if 'license_name' in model_card:
407+
readme_text += f"license_name: {model_card.get('license_name')}\n"
408+
if 'license_link' in model_card:
409+
readme_text += f"license_link: {model_card.get('license_link')}\n"
402410
if 'details' in model_card and 'Dataset' in model_card['details']:
403411
readme_text += 'datasets:\n'
404412
if isinstance(model_card['details']['Dataset'], (tuple, list)):

0 commit comments

Comments
 (0)