handler output reformating

handler.py

@@ -47,6 +47,7 @@ class EndpointHandler:
         output = self.model(**tokenized_inputs)
         predictions = torch.argmax(output.logits, dim=-1).numpy(force=True)
         batch['label'] = [self.model.config.id2label[p] for p in predictions]
-        batch.pop('text')
         
-        return batch
+        n = len(test_dict['id'])
+        result = [{key : value[i] for key, value in test_dict.items() if key not in ('text')} for i in range(n)]
+        return result