changes from bobs repo to line up text embeddings

handler.py

@@ -5,7 +5,7 @@ import logging
 
 from PIL import Image
 import torch
-from transformers import CLIPProcessor, CLIPTextModel, CLIPVisionModelWithProjection
+from transformers import CLIPProcessor, CLIPModel
 
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 logger = logging.getLogger(__name__)
@@ -13,8 +13,7 @@ logger.setLevel('INFO')
 
 class EndpointHandler():
     def __init__(self, path=""):
-        self.text_model = CLIPTextModel.from_pretrained("rbanfield/clip-vit-large-patch14").to(device)
-        self.image_model = CLIPVisionModelWithProjection.from_pretrained("rbanfield/clip-vit-large-patch14").to(device)
+        self.model = CLIPModel.from_pretrained("rbanfield/clip-vit-large-patch14").to("cpu")
         self.processor = CLIPProcessor.from_pretrained("rbanfield/clip-vit-large-patch14")
 
     def __call__(self, data):
@@ -23,7 +22,7 @@ class EndpointHandler():
             inputs = data.pop("inputs", None)
             text_input = None
             image_data = None
-
+            logger.info('data contents: %s', data)
             if isinstance(inputs, Image.Image):
                 logger.info('image sent directly')
                 image = inputs
@@ -38,12 +37,12 @@ class EndpointHandler():
             if text_input:
                 processor = self.processor(text=text_input, return_tensors="pt", padding=True).to(device)
                 with torch.no_grad():
-                    return {'embeddings':self.text_model(**processor).pooler_output.tolist()[0]}
+                    return {"embeddings": self.model.get_text_features(**processor).tolist()}
             elif image:
                 # image = Image.open(image_data)
                 processor = self.processor(images=image, return_tensors="pt").to(device)
                 with torch.no_grad():
-                    return {'embeddings':self.image_model(**processor).image_embeds.tolist()[0]}
+                    return {"embeddings": self.model.get_image_features(**processor).tolist()}
             else:
                 return {'embeddings':None}
         except Exception as ex: