Merge pull request #365 from vespa-engine/tgm/fix-ml-test

fix broken test due to transformer library changes
vespa-engine · Aug 1, 2022 · ad53d7a · ad53d7a
2 parents 04a0417 + d63abe1
commit ad53d7a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vespa/ml.py b/vespa/ml.py
@@ -100,8 +100,8 @@ def predict(self, text: str) -> List:
         pipeline = self._create_pipeline()
         predictions = pipeline(text)
         return [
-            [x["score"] for x in predictions if x["label"] == "LABEL_0"][0],
-            [x["score"] for x in predictions if x["label"] == "LABEL_1"][0],
+            [x["score"] for x in predictions[0] if x["label"] == "LABEL_0"][0],
+            [x["score"] for x in predictions[0] if x["label"] == "LABEL_1"][0],
         ]
 
     @staticmethod