fix: use actual median for model-card word count

src/setfit/model_card.py

@@ -4,6 +4,7 @@
 from dataclasses import dataclass, field, fields
 from pathlib import Path
 from platform import python_version
+from statistics import median
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Tuple, Union
 
 import datasets
@@ -312,7 +313,7 @@ def add_naive_word_count(sample: Dict[str, Any]) -> Dict[str, Any]:
             {
                 "Training set": "Word count",
                 "Min": min(dataset["word_count"]),
-                "Median": sum(dataset["word_count"]) / len(dataset),
+                "Median": median(dataset["word_count"]),
                 "Max": max(dataset["word_count"]),
             },
         ]

tests/model_card_pattern.py

@@ -99,7 +99,7 @@
 ### Training Set Metrics
 \| Training set \| Min \| Median \| Max \|
 \|:-------------\|:----\|:-------\|:----\|
-\| Word count   \| 3   \| 7.875  \| 18  \|
+\| Word count   \| 3   \| [\d.]+ +\| 18  \|
 
 \| Label    \| Training Sample Count \|
 \|:---------\|:----------------------\|

tests/span/aspect_model_card_pattern.py

@@ -111,7 +111,7 @@
 ### Training Set Metrics
 \| Training set \| Min \| Median \| Max \|
 \|:-------------\|:----\|:-------\|:----\|
-\| Word count   \| 5   \| 14.5   \| 23  \|
+\| Word count   \| 5   \| [\d.]+ +\| 23  \|
 
 \| Label     \| Training Sample Count \|
 \|:----------\|:----------------------\|

tests/span/polarity_model_card_pattern.py

@@ -111,7 +111,7 @@
 ### Training Set Metrics
 \| Training set \| Min \| Median \| Max \|
 \|:-------------\|:----\|:-------\|:----\|
-\| Word count   \| 8   \| 16.8   \| 28  \|
+\| Word count   \| 8   \| [\d.]+ +\| 28  \|
 
 \| Label    \| Training Sample Count \|
 \|:---------\|:----------------------\|

tests/test_model_card.py

@@ -62,6 +62,26 @@ def test_model_card_languages() -> None:
     assert "**Languages:** en, nl, de" in model_card
 
 
+def test_train_set_metrics_word_count() -> None:
+    # The word counts are 2, 3, 4 and 11, i.e. a median of 3.5 and a mean of 5.0
+    train_dataset = Dataset.from_dict(
+        {
+            "text": [
+                "one two",
+                "one two three",
+                "one two three four",
+                "one two three four five six seven eight nine ten eleven",
+            ],
+            "label": [0, 1, 0, 1],
+        }
+    )
+    model_card_data = SetFitModelCardData()
+    model_card_data.set_train_set_metrics(train_dataset)
+    assert model_card_data.train_set_metrics_list == [
+        {"Training set": "Word count", "Min": 2, "Median": 3.5, "Max": 11}
+    ]
+
+
 def test_is_on_huggingface_edge_case() -> None:
     assert not is_on_huggingface("test_value")
     assert not is_on_huggingface("a/test/value")