Add table of contents

edit.py

@@ -223,10 +223,23 @@ def create_form_from_section(
 
 
 def display_sections(model_card: card.Card) -> None:
-    for key, section_name, content, is_fig in iterate_key_section_content(
-        model_card._data
-    ):
-        create_form_from_section(model_card, key, section_name, content, is_fig)
+    for section_info in iterate_key_section_content(model_card._data):
+        create_form_from_section(
+            model_card,
+            key=section_info.return_key,
+            section_name=section_info.title,
+            content=section_info.content,
+            is_fig=section_info.is_fig,
+        )
+
+
+def display_toc(model_card: card.Card) -> None:
+    elements = []
+    for section_info in iterate_key_section_content(model_card._data):
+        title, level = section_info.title, section_info.level
+        section_name = split_subsection_names(title)[-1]
+        elements.append("  " * level + "- " + section_name)
+    st.markdown("\n".join(elements))
 
 
 def display_model_card(model_card: card.Card) -> None:
@@ -236,6 +249,10 @@ def display_model_card(model_card: card.Card) -> None:
     # strip metadata
     with st.expander("show metadata"):
         st.text(metadata)
+
+    with st.expander("Table of Contents"):
+        display_toc(model_card)
+
     st.markdown(rendered, unsafe_allow_html=True)
 
 

utils.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import base64
 import re
+from dataclasses import dataclass
 from pathlib import Path
 from typing import Iterator
 
@@ -61,11 +62,21 @@ def process_card_for_rendering(rendered: str) -> tuple[str, str]:
     return metadata, rendered_with_img
 
 
+@dataclass(frozen=True)
+class SectionInfo:
+    return_key: str
+    title: str
+    content: str
+    is_fig: bool
+    level: int
+
+
 def iterate_key_section_content(
     data: dict[str, Section],
     parent_section: str = "",
     parent_keys: list[str] | None = None,
-) -> Iterator[tuple[str, str, str, bool]]:
+    level=0,
+) -> SectionInfo:
     parent_keys = parent_keys or []
 
     for key, val in data.items():
@@ -79,11 +90,18 @@ def iterate_key_section_content(
 
         return_key = key if not parent_keys else "/".join(parent_keys + [key])
         is_fig = getattr(val, "is_fig", False)
-        yield return_key, title, val.content, is_fig
+        yield SectionInfo(
+            return_key=return_key,
+            title=title,
+            content=val.content,
+            is_fig=is_fig,
+            level=level,
+        )
 
         if val.subsections:
             yield from iterate_key_section_content(
                 val.subsections,
                 parent_section=title,
                 parent_keys=parent_keys + [key],
+                level=level + 1,
             )