Update data contract

handler.py

@@ -118,19 +118,18 @@ class EndpointHandler():
         return self.response_pipeline.predict_proba(df1)[:,1]
     
     
-    def __call__(self, data: Dict[str, Union[Dict, List, pd.DataFrame]]):
+    def __call__(self, data: Dict[str, Union[Dict, List]]):
 
         inputs = data['inputs']
 
         is_dict =  isinstance(inputs, dict)
-        is_list = isinstance(inputs, list)
 
         if is_dict:
             df = pd.DataFrame([inputs])
-        elif is_list:
-            df = pd.DataFrame(inputs)
         else:
-            df = inputs
+            df = pd.DataFrame(inputs)
+
+        df = df.fillna('')
 
         if 'dataset' not in df.columns:
             df['dataset'] = ''
@@ -146,7 +145,5 @@ class EndpointHandler():
 
         if is_dict:
             return predictions[0]
-        elif is_list:
-            return predictions
         else:
-            return pd.DataFrame(predictions, index=df.index)
+            return predictions

tests.py

@@ -14,16 +14,9 @@ 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)}
+                       "response": "Drink lots of coffee!"},
+                      {"instruction": "What are some ways to stay energized throughout the day?",
+                       "response": "Buy lots of sailboats!"}]}
 
 # test the handler
 pred=response_model_handler(payload)