add all options to handler

handler.py

@@ -71,11 +71,7 @@ class EndpointHandler:
 
         return encoding
 
-    def predict(self, data: Dict[str, Any], preprocessed: Dict[str, Any]) -> Dict[str, Any]:
-        task = data.get('task', None)
-
-        if task is None or task not in self.task_config.keys():
-            raise ValueError(f"Invalid task: {task}")
+    def predict(self, task:str, preprocessed: Dict[str, Any]) -> Dict[str, Any]:
 
         logits = self.task_heads[task](self.model(**preprocessed).last_hidden_state[:, 0, :])
         config = self.task_config[task]
@@ -89,19 +85,34 @@ class EndpointHandler:
             pred_idx = int(np.argmax(probs))
             return {'label': config['label_map'][pred_idx], 'confidence': float(probs[pred_idx])}
 
-    def postprocess(self, outputs: Dict[str, Any]) -> List[Dict[str, Any]]:
-        return [outputs]
+    # def postprocess(self, outputs: Dict[str, Any]) -> List[Dict[str, Any]]:
+    #     return [outputs]
+
+
 
     def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
         task = data.get('task', None)
         print(f"Task: {task}")
         if task is None:
             raise ValueError("'task' key is required in the input dictionary")
-        if task not in self.task_config.keys():
+        task = task.lower()
+
+        if task == "all":
+            results = {}
+
+            for _t in self.task_config.keys():
+                data['task'] = _t
+                preprocessed = self.preprocess(data)
+                outputs = self.predict(_t, preprocessed)
+                results[_t] = outputs
+            return results
+        elif task not in self.task_config.keys():
             raise ValueError(f"Invalid task: {task}")
         preprocessed = self.preprocess(data)
-        outputs = self.predict(data, preprocessed)
-        return self.postprocess(outputs)
+
+        outputs = self.predict(task, preprocessed)
+        # return self.postprocess(outputs)
+        return outputs
 
 class TaskClassificationHead(torch.nn.Module):
     def __init__(self, hidden_size: int, num_labels: int, dropout: float):