Add handler

handler.py

@@ -0,0 +1,23 @@
+from pylate import models
+from transformers import AutoTokenizer
+import torch
+
+class EndpointHandler:
+    def __init__(self, path=""):
+        self.tokenizer = AutoTokenizer.from_pretrained(path, use_fast=True)
+        self.model = models.ColBERT(model_name_or_path=path)
+        self.model.eval()
+
+    def __call__(self, data):
+        texts = data.get("inputs") or data.get("text") or data
+        if isinstance(texts, str):
+            texts = [texts]
+
+        with torch.no_grad():
+            emb = self.model.encode(
+                texts,
+                is_query=True,          # query-style encoding
+                batch_size=32,
+            )
+        # TEI expects JSON-serialisable output
+        return emb.cpu().tolist()

requirements.txt

@@ -0,0 +1,2 @@
+pylate>=0.4.0
+torch>=2.2