type variable df

knative · May 21, 2024 · 7b7b721 · 7b7b721
1 parent 818685b
commit 7b7b721
Showing 1 changed file with 10 additions and 1 deletion.
diff --git a/code-samples/community/serving/machinelearning-python-bentoml/iris_classifier.py b/code-samples/community/serving/machinelearning-python-bentoml/iris_classifier.py
@@ -1,4 +1,8 @@
+import numpy as np
 import bentoml
+from pydantic import Field
+from bentoml.validators import Shape
+from typing_extensions import Annotated
 import joblib
 
 
@@ -10,5 +14,10 @@ def __init__(self):
         self.model = joblib.load(self.iris_model.path_of("model.pkl"))
 
     @bentoml.api
-    def predict(self, df):
+    def predict(
+        self,
+        df: Annotated[np.ndarray, Shape((-1, 4))] = Field(
+            default=[[5.2, 2.3, 5.0, 0.7]]
+        ),
+    ) -> np.ndarray:
         return self.artifacts.model.predict(df)