Factor out readme generation in hub helper, add more readme fields

timm/models/_hub.py

@@ -236,10 +236,29 @@ def push_to_hf_hub(
             model_card = model_card or {}
             model_name = repo_id.split('/')[-1]
             readme_path = Path(tmpdir) / "README.md"
+            readme_text = generate_readme(model_card, model_name)
+            readme_path.write_text(readme_text)
+
+        # Upload model and return
+        return upload_folder(
+            repo_id=repo_id,
+            folder_path=tmpdir,
+            revision=revision,
+            create_pr=create_pr,
+            commit_message=commit_message,
+        )
+
+
+def generate_readme(model_card, model_name):
     readme_text = "---\n"
     readme_text += "tags:\n- image-classification\n- timm\n"
     readme_text += "library_tag: timm\n"
     readme_text += f"license: {model_card.get('license', 'apache-2.0')}\n"
+    if 'details' in model_card and 'Dataset' in model_card['details']:
+        readme_text += 'datasets:\n'
+        readme_text += f"- {model_card['details']['Dataset'].lower()}\n"
+        if 'Pretrain Dataset' in model_card['details']:
+            readme_text += f"- {model_card['details']['Pretrain Dataset'].lower()}\n"
     readme_text += "---\n"
     readme_text += f"# Model card for {model_name}\n"
     if 'description' in model_card:
@@ -247,16 +266,32 @@ def push_to_hf_hub(
     if 'details' in model_card:
         readme_text += f"\n## Model Details\n"
         for k, v in model_card['details'].items():
+            if isinstance(v, (list, tuple)):
+                readme_text += f"- **{k}:**\n"
+                for vi in v:
+                    readme_text += f"  - {vi}\n"
+            elif isinstance(v, dict):
+                readme_text += f"- **{k}:**\n"
+                for ki, vi in v.items():
+                    readme_text += f"  - {ki}: {vi}\n"
+            else:
                 readme_text += f"- **{k}:** {v}\n"
-            if 'citation' in model_card:
+    if 'usage' in model_card:
-                readme_text += f"\n## Citation\n```\n{model_card['citation']}```\n"
+        readme_text += f"\n## Model Usage\n"
-            readme_path.write_text(readme_text)
+        readme_text += model_card['usage']
+        readme_text += '\n'
 
-        # Upload model and return
+    if 'comparison' in model_card:
-        return upload_folder(
+        readme_text += f"\n## Model Comparison\n"
-            repo_id=repo_id,
+        readme_text += model_card['comparison']
-            folder_path=tmpdir,
+        readme_text += '\n'
-            revision=revision,
+
-            create_pr=create_pr,
+    if 'citation' in model_card:
-            commit_message=commit_message,
+        readme_text += f"\n## Citation\n"
-        )
+        if not isinstance(model_card['citation'], (list, tuple)):
+            citations = [model_card['citation']]
+        else:
+            citations = model_card['citation']
+        for c in citations:
+            readme_text += f"```bibtex\n{c}\n```\n"
+    return readme_text