add output type annotation

vespa-engine · Jun 2, 2022 · 015feaf · 015feaf
1 parent 7361065
commit 015feaf
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/vespa/ml.py b/vespa/ml.py
@@ -90,7 +90,7 @@ def export_to_onnx(self, output_path: str) -> None:
             pipeline, opset=11, output=Path(output_path), use_external_format=False
         )
 
-    def predict(self, text: str):
+    def predict(self, text: str) -> List:
         """
         Predict using a local instance of the model
 
@@ -102,7 +102,7 @@ def predict(self, text: str):
         return [x["score"] for x in predictions]
 
     @staticmethod
-    def parse_vespa_prediction(prediction):
+    def parse_vespa_prediction(prediction) -> List:
         return [cell["value"] for cell in prediction["cells"]]
 
     def create_url_encoded_tokens(self, x):