Improve inference output

EveripediaNetwork · Jun 17, 2024 · ba1672d · ba1672d
1 parent 5313ef9
commit ba1672d
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/fastc/classifiers/centroids.py b/fastc/classifiers/centroids.py
@@ -79,7 +79,12 @@ def predict(
                 for label, score in zip(dot_products.keys(), softmax_scores)
             }
 
-            yield scores
+            result = {
+                'label': max(scores, key=scores.get),
+                'scores': scores,
+            }
+
+            yield result
 
     def predict_one(self, text: str) -> Dict[int, float]:
         return list(self.predict([text]))[0]