Add special fields for skops template

- It's now possible to enter metrics in text area
- It's now possible to toggle model diagram

edit.py

@@ -37,6 +37,7 @@ from skops.card._model_card import PlotSection, split_subsection_names
 
 from utils import iterate_key_section_content, process_card_for_rendering
 from tasks import (
+    AddMetricsTask,
     AddSectionTask,
     AddFigureTask,
     DeleteSectionTask,
@@ -279,33 +280,107 @@ def add_download_model_card_button():
     )
 
 
+def display_edit_buttons():
+    # first row: undo + redo + reset
+    col_0, col_1, col_2, *_ = st.columns([2, 2, 2, 2])
+    undo_disabled = not bool(st.session_state.task_state.done_list)
+    redo_disabled = not bool(st.session_state.task_state.undone_list)
+    with col_0:
+        name = f"UNDO ({len(st.session_state.task_state.done_list)})"
+        tip = "Undo the last edit"
+        st.button(name, on_click=undo_last, disabled=undo_disabled, help=tip)
+    with col_1:
+        name = f"REDO ({len(st.session_state.task_state.undone_list)})"
+        tip = "Redo the last undone edit"
+        st.button(name, on_click=redo_last, disabled=redo_disabled, help=tip)
+    with col_2:
+        tip = "Undo all edits"
+        st.button("Reset", on_click=reset_model_card, help=tip)
+
+    # second row: download + delete
+    col_0, col_1, *_ = st.columns([2, 2, 2, 2])
+    with col_0:
+        add_download_model_card_button()
+    with col_1:
+        tip = "Start over from scratch (lose all progress)"
+        st.button("Delete", on_click=delete_model_card, help=tip)
+
+
+def _update_model_diagram():
+    val = st.session_state.get("special_model_diagram", True)
+    model_card = st.session_state.model_card
+    model_card.model_diagram = val
+
+    # TODO: this may no longer be necesssary once this issue is solved:
+    # https://github.com/skops-dev/skops/issues/292
+    if val:
+        model_card.add_model_plot()
+    else:
+        model_card.delete("Model description/Training Procedure/Model Plot")
+
+
+def _parse_metrics(metrics: str) -> dict[str, str | float]:
+    # parse metrics from text area, one per line, into a dict
+    metrics_table = {}
+    for line in metrics.splitlines():
+        line = line.strip()
+        val: str | float
+        name, _, val = line.partition("=")
+        try:
+            # try to coerce to float but don't error if it fails
+            val = float(val.strip())
+        except ValueError:
+            pass
+        metrics_table[name.strip()] = val
+    return metrics_table
+
+
+def _update_metrics():
+    metrics = st.session_state.get("special_metrics_text", {})
+    model_card = st.session_state.model_card
+    metrics_table = _parse_metrics(metrics)
+
+    # check if any change
+    if metrics_table == model_card._metrics:
+        return
+
+    task = AddMetricsTask(model_card, metrics_table)
+    st.session_state.task_state.add(task)
+
+
+def display_skops_special_fields():
+    st.checkbox(
+        "Show model diagram",
+        value=True,
+        on_change=_update_model_diagram,
+        key="special_model_diagram",
+    )
+
+    with st.expander("Add metrics"):
+        with st.form("special_metrics", clear_on_submit=False):
+            st.text_area(
+                "Add one metric per line, e.g. 'accuracy = 0.9'",
+                key="special_metrics_text",
+            )
+            st.form_submit_button(
+                "Update",
+                on_click=_update_metrics,
+            )
+
+
 def edit_input_form():
     if "task_state" not in st.session_state:
         st.session_state.task_state = TaskState()
 
     with st.sidebar:
-        col_0, col_1, col_2, *_ = st.columns([2, 2, 2, 2])
-        undo_disabled = not bool(st.session_state.task_state.done_list)
-        redo_disabled = not bool(st.session_state.task_state.undone_list)
-        with col_0:
-            name = f"UNDO ({len(st.session_state.task_state.done_list)})"
-            tip = "Undo the last edit"
-            st.button(name, on_click=undo_last, disabled=undo_disabled, help=tip)
-        with col_1:
-            name = f"REDO ({len(st.session_state.task_state.undone_list)})"
-            tip = "Redo the last undone edit"
-            st.button(name, on_click=redo_last, disabled=redo_disabled, help=tip)
-        with col_2:
-            tip = "Undo all edits"
-            st.button("Reset", on_click=reset_model_card, help=tip)
-
-        col_0, col_1, *_ = st.columns([2, 2, 2, 2])
-        with col_0:
-            add_download_model_card_button()
-        with col_1:
-            tip = "Start over from scratch (lose all progress)"
-            st.button("Delete", on_click=delete_model_card, help=tip)
+        # TOP ROW BUTTONS
+        display_edit_buttons()
+
+        # SHOW SPECIAL FIELDS IF SKOPS TEMPLATE WAS USED
+        if st.session_state.get("model_card_type", "") == "skops":
+            display_skops_special_fields()
 
+        # SHOW EDITABLE SECTIONS
         if "model_card" in st.session_state:
             display_sections(st.session_state.model_card)
 

start.py

@@ -107,6 +107,7 @@ def create_skops_model_card() -> None:
     metadata = card.metadata_from_config(hf_path)
     model_card = card.Card(model=st.session_state.model, metadata=metadata)
     st.session_state.model_card = model_card
+    st.session_state.model_card_type = "skops"
 
 
 def create_empty_model_card() -> None:
@@ -117,6 +118,7 @@ def create_empty_model_card() -> None:
     )
     model_card.add(**{"Untitled": "[More Information Needed]"})
     st.session_state.model_card = model_card
+    st.session_state.model_card_type = "empty"
 
 
 def create_hf_model_card() -> None:
@@ -128,6 +130,7 @@ def create_hf_model_card() -> None:
     path = hf_hub_download(repo_id, "README.md")
     model_card = card.parse_modelcard(path)
     st.session_state.model_card = model_card
+    st.session_state.model_card_type = "loaded"
 
 
 def start_input_form():

tasks.py

@@ -208,3 +208,24 @@ class UpdateFigureTask(Task):
 
         self.path.unlink(missing_ok=True)
         section.content = self.old_data
+
+
+class AddMetricsTask(Task):
+    """Add new metrics"""
+
+    def __init__(
+        self,
+        model_card: card.Card,
+        metrics: dict[str, str | int | float],
+    ) -> None:
+        self.model_card = model_card
+        self.old_metrics = model_card._metrics.copy()
+        self.new_metrics = metrics
+
+    def do(self) -> None:
+        self.model_card._metrics.clear()
+        self.model_card.add_metrics(**self.new_metrics)
+
+    def undo(self) -> None:
+        self.model_card._metrics.clear()
+        self.model_card.add_metrics(**self.old_metrics)