Ensure that textcat tasks scores are always a float (#61)

explosion · May 11, 2023 · 1217ca6 · 1217ca6
1 parent cde7eeb
commit 1217ca6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/spacy_llm/tasks/textcat.py b/spacy_llm/tasks/textcat.py
@@ -137,7 +137,7 @@ def _format_response(self, response: str) -> Dict[str, float]:
             categories = {label: score}
         else:
             # Multilabel classification
-            categories = {label: 0 for label in self._label_dict.values()}
+            categories = {label: 0.0 for label in self._label_dict.values()}
 
             pred_labels = response.split(",")
             if self._exclusive_classes and len(pred_labels) > 1: