Update data contract

handler.py

@@ -1,5 +1,5 @@
 
-from typing import Dict, Union, Optional
+from typing import Dict, List, Union, Optional
 from pathlib import Path
 import json
 import joblib
@@ -114,12 +114,19 @@ class EndpointHandler():
         return self.response_pipeline.predict_proba(df1)[:,1]
     
     
-    def __call__(self, df: Union[pd.DataFrame, Dict]):
+    def __call__(self, data: Dict[str, Union[Dict, List, pd.DataFrame]]):
 
-        is_dict =  isinstance(df, dict)
+        inputs = data['inputs']
+
+        is_dict =  isinstance(inputs, dict)
+        is_list = isinstance(inputs, list)
 
         if is_dict:
-            df = pd.DataFrame([df])
+            df = pd.DataFrame([inputs])
+        elif is_list:
+            df = pd.DataFrame(inputs)
+        else:
+            df = inputs
 
         if 'dataset' not in df.columns:
             df['dataset'] = ''
@@ -135,5 +142,7 @@ class EndpointHandler():
 
         if is_dict:
             return predictions[0]
+        elif is_list:
+            return predictions
         else:
             return pd.DataFrame(predictions, index=df.index)

tests.py

@@ -1,11 +1,29 @@
 from handler import EndpointHandler
+import pandas as pd
 
 # init handler
 response_model_handler = EndpointHandler()
 
 # prepare sample payload
-payload = {"instruction": "What are some ways to stay energized throughout the day?",
-           "response": "Drink lots of coffee!"}
+payload = {'inputs': {"instruction": "What are some ways to stay energized throughout the day?",
+           "response": "Drink lots of coffee!"}}
+
+# test the handler
+pred=response_model_handler(payload)
+
+print(pred)
+
+payload = {'inputs': [{"instruction": "What are some ways to stay energized throughout the day?",
+           "response": "Drink lots of coffee!"}]*2}
+
+# test the handler
+pred=response_model_handler(payload)
+
+print(pred)
+
+
+payload = {'inputs': pd.DataFrame([{"instruction": "What are some ways to stay energized throughout the day?",
+           "response": "Drink lots of coffee!"}]*2)}
 
 # test the handler
 pred=response_model_handler(payload)