fix(text): predicted ok with unordered labels (#12)

argilla-io · May 4, 2021 · e6e17f5 · e6e17f5
1 parent 0d6aaf9
commit e6e17f5
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 1 deletion.
diff --git a/src/rubrix/server/tasks/text_classification/api/model.py b/src/rubrix/server/tasks/text_classification/api/model.py
@@ -138,7 +138,7 @@ def predicted(self) -> Optional[PredictionStatus]:
         if self.predicted_as and self.annotated_as:
             return (
                 PredictionStatus.OK
-                if self.predicted_as == self.annotated_as
+                if set(self.predicted_as) == set(self.annotated_as)
                 else PredictionStatus.KO
             )
         return None

diff --git a/tests/server/text_classification/test_model.py b/tests/server/text_classification/test_model.py
@@ -123,3 +123,24 @@ def test_created_record_with_default_status():
 
     record = TextClassificationRecord.parse_obj(data)
     assert record.status == TaskStatus.default
+
+
+def test_predicted_ok_for_multilabel_unordered():
+    record = TextClassificationRecord(
+        inputs={"text": "The text"},
+        prediction=TextClassificationAnnotation(
+            agent="test",
+            labels=[
+                ClassPrediction(class_label="B"),
+                ClassPrediction(class_label="C", score=0.3),
+                ClassPrediction(class_label="A"),
+            ],
+        ),
+        annotation=TextClassificationAnnotation(
+            agent="test",
+            labels=[ClassPrediction(class_label="A"), ClassPrediction(class_label="B")],
+        ),
+        multi_label=True,
+    )
+
+    assert record.predicted == PredictionStatus.OK