diff --git a/skops/card/_model_card.py b/skops/card/_model_card.py index 7bf357f8..15a8fe04 100644 --- a/skops/card/_model_card.py +++ b/skops/card/_model_card.py @@ -1236,7 +1236,8 @@ def add_fairlearn_metric_frame( self, metric_frame, table_name: str = "Fairlearn MetricFrame Table", - transpose=True, + transpose: bool = True, + description: str | None = None, ) -> Self: """ Add a :class:`fairlearn.metrics.MetricFrame` table to the model card. The table contains @@ -1247,11 +1248,15 @@ def add_fairlearn_metric_frame( metric_frame: MetricFrame The Fairlearn MetricFrame to add to the model card. + table_name: str + The desired name of the table section in the model card. + transpose: bool, default=True Whether to transpose the table or not. - table_name: str - The desired name of the table section in the model card. + + description : str | None (default=None) + An optional description to be added before the table. Returns ------- @@ -1277,7 +1282,9 @@ def add_fairlearn_metric_frame( frame_dict = pd.DataFrame(frame_dict).T - return self.add_table(folded=True, **{table_name: frame_dict}) + return self.add_table( + folded=True, description=description, **{table_name: frame_dict} + ) def _add_metrics( self, diff --git a/skops/card/tests/test_card.py b/skops/card/tests/test_card.py index 15d73eb3..0a97a303 100644 --- a/skops/card/tests/test_card.py +++ b/skops/card/tests/test_card.py @@ -1721,8 +1721,8 @@ def card(self): card = Card(model=model) return card - @pytest.mark.parametrize("transpose", [True, False]) - def test_metric_table(self, card: Card, transpose): + @pytest.fixture + def metric_frame(self): metrics = import_or_raise( "fairlearn.metrics", "model card fairlearn metricframe" ) @@ -1734,13 +1734,17 @@ def test_metric_table(self, card: Card, transpose): metric_frame = metrics.MetricFrame( y_true=y_true, y_pred=y_pred, sensitive_features=sex, metrics=metric_dict ) + return metric_frame + + @pytest.mark.parametrize("transpose", [True, False]) + def test_metric_table(self, card: Card, transpose, metric_frame): card.add_fairlearn_metric_frame( metric_frame=metric_frame, transpose=transpose, table_name="Metric Frame Table", ) - actual_table = card.select("Metric Frame Table").content.format() + actual_table = card.select("Metric Frame Table").format() if transpose is True: expected_table = ( @@ -1757,3 +1761,19 @@ def test_metric_table(self, card: Card, transpose): ) assert expected_table == actual_table + + def test_metric_table_with_description(self, card: Card, metric_frame): + card.add_fairlearn_metric_frame( + description="An awesome table", + metric_frame=metric_frame, + table_name="Metric Frame Table", + ) + + actual_table = card.select("Metric Frame Table").format() + expected_table = ( + "An awesome table\n\n" + "
\n Click to expand \n\n| selection_rate" + " |\n|------------------|\n| 0.4 |\n| 0.8" + " |\n| 0.4 |\n| 0.5 |\n\n
" + ) + assert expected_table == actual_table