MNT Model card refactor to move plotting & table (#163)

Co-authored-by: Adrin Jalali <adrin.jalali@gmail.com>
skops-dev · Oct 5, 2022 · 8fe2577 · 8fe2577
1 parent f41f891
commit 8fe2577
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 10 deletions.
diff --git 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
@@ -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)