From 8fe25774f9e750808609c92821b2d3c09057ecc4 Mon Sep 17 00:00:00 2001
From: Merve Noyan <merveenoyan@gmail.com>
Date: Wed, 5 Oct 2022 11:45:03 +0200
Subject: [PATCH] MNT Model card refactor to move plotting & table (#163)

Co-authored-by: Adrin Jalali <adrin.jalali@gmail.com>
---
 skops/card/_model_card.py     | 21 +++++++++++----------
 skops/card/tests/test_card.py |  6 ++++++
 2 files changed, 17 insertions(+), 10 deletions(-)

diff --git a/skops/card/_model_card.py b/skops/card/_model_card.py
index a21353f8..7b5dcf18 100644
--- a/skops/card/_model_card.py
+++ b/skops/card/_model_card.py
@@ -254,15 +254,8 @@ def __init__(
         metadata: Optional[CardData] = None,
     ) -> None:
         self.model = model
-        self._hyperparameter_table = self._extract_estimator_config()
-        # the spaces in the pipeline breaks markdown, so we replace them
+        self.model_diagram = model_diagram
         self._eval_results = {}  # type: ignore
-        if model_diagram is True:
-            self._model_plot: str | None = re.sub(
-                r"\n\s+", "", str(estimator_html_repr(model))
-            )
-        else:
-            self._model_plot = None
         self._template_sections: dict[str, str] = {}
         self._extra_sections: list[tuple[str, Any]] = []
         self.metadata = metadata or CardData()
@@ -393,6 +386,7 @@ def _generate_card(self) -> ModelCard:
         # add evaluation results
 
         template_sections = copy.deepcopy(self._template_sections)
+
         if self.metadata:
             if self.metadata.to_dict().get("model_file"):
                 model_file = self.metadata.to_dict().get("model_file")
@@ -415,11 +409,18 @@ def _generate_card(self) -> ModelCard:
                         " json.load(f)\n"
                         'clf.predict(pd.DataFrame.from_dict(config["sklearn"]["example_input"]))'
                     )
+        if self.model_diagram is True:
+            model_plot: str | None = re.sub(
+                r"\n\s+", "", str(estimator_html_repr(self.model))
+            )
+        else:
+            model_plot = None
         template_sections["eval_results"] = tabulate(
             list(self._eval_results.items()),
             headers=["Metric", "Value"],
             tablefmt="github",
         )
+
         # if template path is not given, use default
         if template_sections.get("template_path") is None:
             template_sections["template_path"] = str(
@@ -446,8 +447,8 @@ def _generate_card(self) -> ModelCard:
 
             card = ModelCard.from_template(
                 card_data=self.metadata,
-                hyperparameter_table=self._hyperparameter_table,
-                model_plot=self._model_plot,
+                hyperparameter_table=self._extract_estimator_config(),
+                model_plot=model_plot,
                 **template_sections,
             )
         return card
diff --git a/skops/card/tests/test_card.py b/skops/card/tests/test_card.py
index fc87ecaa..9cb1103b 100644
--- a/skops/card/tests/test_card.py
+++ b/skops/card/tests/test_card.py
@@ -146,6 +146,12 @@ def test_metadata_keys(destination_path, model_card):
     assert "tags: dummy" in model_card
 
 
+def test_default_sections_save(model_card):
+    # test if the plot and hyperparameters are only added during save
+    assert "<style>" not in str(model_card)
+    assert "fit_intercept" not in str(model_card)
+
+
 def test_add_metrics(destination_path, model_card):
     model_card.add_metrics(**{"acc": 0.1})
     model_card.add_metrics(f1=0.1)