Added binary file handing to inference endpoint and made return value a dictionary

handler.py

@@ -14,18 +14,25 @@ class EndpointHandler():
         self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
 
     def __call__(self, data):
-        inputs = data.pop("inputs", None)
-        text_input = inputs["text"] if "text" in inputs else None
-        image_input = inputs["image"] if "image" in inputs else None
+
+        text_input = None
+        if isinstance(data, dict):
+            inputs = data.pop("inputs", None)
+            text_input = inputs.get('text',None)
+            image_data = BytesIO(base64.b64decode(inputs['image'])) if 'image' in inputs else None
+        else:
+            # assuming its an image sent via binary
+            image_data = BytesIO(data)
+
 
         if text_input:
             processor = self.processor(text=text_input, return_tensors="pt", padding=True).to(device)
             with torch.no_grad():
-                return self.text_model(**processor).pooler_output.tolist()
-        elif image_input:
-            image = Image.open(BytesIO(base64.b64decode(image_input)))
+                return {'embeddings':self.text_model(**processor).pooler_output.tolist()[0]}
+        elif image_data:
+            image = Image.open(image_data)
             processor = self.processor(images=image, return_tensors="pt").to(device)
             with torch.no_grad():
-                return self.image_model(**processor).image_embeds.tolist()
+                return {'embeddings':self.image_model(**processor).image_embeds.tolist()[0]}
         else:
-            return None
+            return {'embeddings':None}