test: update test data

argilla-io · Jun 1, 2022 · fb592eb · fb592eb
1 parent 9964390
commit fb592eb
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/src/rubrix/server/apis/v0/models/metrics/text_classification.py b/src/rubrix/server/apis/v0/models/metrics/text_classification.py
@@ -37,7 +37,6 @@ def apply(self, records: Iterable[TextClassificationRecord]) -> Any:
         if not len(ds_labels):
             return {}
 
-        labels_mapping = {label: i for i, label in enumerate(ds_labels)}
         y_true, y_pred = ([], [])
         for record in filtered_records:
             annotations = record.annotated_as
@@ -65,7 +64,7 @@ def apply(self, records: Iterable[TextClassificationRecord]) -> Any:
 
         per_label = {}
         for label, p, r, f, s in zip(
-            labels_mapping.keys(),
+            ds_labels,
             *precision_recall_fscore_support(
                 y_true=y_true,
                 y_pred=y_pred,

diff --git a/tests/metrics/test_text_classification.py b/tests/metrics/test_text_classification.py
@@ -45,13 +45,15 @@ def test_metrics_for_text_classification(mocked_client):
         "ham_f1": 1.0,
         "ham_precision": 1.0,
         "ham_recall": 1.0,
+        "ham_support": 1,
         "precision_macro": 1.0,
         "precision_micro": 1.0,
         "recall_macro": 1.0,
         "recall_micro": 1.0,
         "spam_f1": 1.0,
         "spam_precision": 1.0,
         "spam_recall": 1.0,
+        "spam_support": 1,
     }
     results.visualize()
 
@@ -63,13 +65,15 @@ def test_metrics_for_text_classification(mocked_client):
         "ham_f1": 1.0,
         "ham_precision": 1.0,
         "ham_recall": 1.0,
+        "ham_support": 1,
         "precision_macro": 1.0,
         "precision_micro": 1.0,
         "recall_macro": 1.0,
         "recall_micro": 1.0,
         "spam_f1": 1.0,
         "spam_precision": 1.0,
         "spam_recall": 1.0,
+        "spam_support": 1,
     }
     results.visualize()