Add handler

handler.py

@@ -0,0 +1,30 @@
+from typing import  Dict, List, Any
+import torch
+from transformers import AutoProcessor, Pix2StructVisionModel
+from PIL import Image
+import pdb
+import requests
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        #self.processor = AutoProcessor.from_pretrained("jasper-lu/pix2struct_embedding")
+        #self.model = MarkupLMModel.from_pretrained("jasper-lu/pix2struct_embedding")
+        self.processor = AutoProcessor.from_pretrained("google/pix2struct-textcaps-base")
+        self.model = Pix2StructVisionModel.from_pretrained("google/pix2struct-textcaps-base")
+
+    def __call__(self, data: Any) -> List[List[Dict[str, float]]]:
+        image = Image.open(requests.get(data["url"], stream=True).raw)
+        inputs = self.processor(images=image, return_tensors="pt")
+
+        with torch.no_grad():
+            outputs = self.model(**inputs)
+        
+        last_hidden_state = outputs['last_hidden_state']
+        embedding = torch.mean(last_hidden_state, dim=1).flatten().tolist()
+        return {"embedding": embedding}
+
+"""
+handler = EndpointHandler()
+output = handler({"url": "https://figma-staging-api.s3.us-west-2.amazonaws.com/images/a8c6a0cc-c022-4f3a-9fc5-ac8582c964dd"})
+print(output)
+"""

requirements.txt

@@ -0,0 +1,24 @@
+certifi==2023.7.22
+charset-normalizer==3.3.2
+filelock==3.13.1
+fsspec==2023.10.0
+huggingface-hub==0.17.3
+idna==3.4
+Jinja2==3.1.2
+MarkupSafe==2.1.3
+mpmath==1.3.0
+networkx==3.2.1
+numpy==1.26.2
+packaging==23.2
+Pillow==10.1.0
+PyYAML==6.0.1
+regex==2023.10.3
+requests==2.31.0
+safetensors==0.4.0
+sympy==1.12
+tokenizers==0.14.1
+torch==2.1.0
+tqdm==4.66.1
+transformers==4.35.0
+typing_extensions==4.8.0
+urllib3==2.1.0