Create handler.py

handler.py

@@ -0,0 +1,30 @@
+from typing import Dict, List, Any
+from colpali_engine.models import ColQwen2, ColQwen2Processor
+import torch
+
+class EndpointHandler():
+    def __init__(self, path=""):
+        self.model = ColQwen2.from_pretrained(
+            path,
+            torch_dtype=torch.bfloat16,
+            device_map="cuda:0",  # or "mps" if on Apple Silicon
+            # attn_implementation="flash_attention_2", # should work on A100
+            ).eval()
+        self.processor = ColQwen2Processor.from_pretrained(path)
+
+    def __call__(self, data: Dict[str, Any]) -> List[Dict[str, Any]]:
+        """
+       data args:
+            inputs (:obj: `str`)
+      Return:
+            A :obj:`list` | `dict`: will be serialized and returned
+        """
+        # process input
+        images = data.pop("inputs", data)
+        batch_images = self.processor.process_images(images).to(model.device)
+        # Forward pass
+        with torch.no_grad():
+            image_embeddings = model(**batch_images)
+
+        return {"embeddings": list(torch.unbind(image_embeddings_doc.to("cpu")))}
+