@@ -209,6 +209,7 @@ def push_to_hf_hub(
209209 private : bool = False ,
210210 create_pr : bool = False ,
211211 model_config : Optional [dict ] = None ,
212+ model_card : Optional [dict ] = None ,
212213):
213214 # Create repo if it doesn't exist yet
214215 repo_url = create_repo (repo_id , token = token , private = private , exist_ok = True )
@@ -232,9 +233,23 @@ def push_to_hf_hub(
232233
233234 # Add readme if it does not exist
234235 if not has_readme :
236+ model_card = model_card or {}
235237 model_name = repo_id .split ('/' )[- 1 ]
236238 readme_path = Path (tmpdir ) / "README.md"
237- readme_text = f'---\n tags:\n - image-classification\n - timm\n library_tag: timm\n ---\n # Model card for { model_name } '
239+ readme_text = "---\n "
240+ readme_text += "tags:\n - image-classification\n - timm\n "
241+ readme_text += "library_tag: timm\n "
242+ readme_text += f"license: { model_card .get ('license' , 'apache-2.0' )} \n "
243+ readme_text += "---\n "
244+ readme_text += f"# Model card for { model_name } \n "
245+ if 'description' in model_card :
246+ readme_text += f"\n { model_card ['description' ]} \n "
247+ if 'details' in model_card :
248+ readme_text += f"\n ## Model Details\n "
249+ for k , v in model_card ['details' ].items ():
250+ readme_text += f"- **{ k } :** { v } \n "
251+ if 'citation' in model_card :
252+ readme_text += f"\n ## Citation\n ```\n { model_card ['citation' ]} ```\n "
238253 readme_path .write_text (readme_text )
239254
240255 # Upload model and return
0 commit comments