Create handler.py

handler.py

@@ -0,0 +1,48 @@
+from typing import Dict, List, Any
+from transformers import AutoTokenizer
+from translate import EncoderCT2fromHfHub
+import os
+
+class EndpointHandler():
+    def __init__(self, path="", local_test=False):
+        path = "tmp/ct2fast-e5-small-v2-hfie"
+        snapshot_id = "1"
+        if local_test:
+            repo_dir = os.getcwd()
+        else:
+            repo_dir = "/repository/"
+        cache_dir = os.path.join(os.path.expanduser("~/.cache/huggingface/hub/models--" + path.replace("/", "--")))
+        snapshot_dir = cache_dir + "/snapshots/" + snapshot_id
+        os.makedirs(cache_dir + "/refs", exist_ok=True)
+        os.makedirs(snapshot_dir, exist_ok=True)
+        with open(cache_dir + "/refs/main", 'w') as filee:
+            filee.write(snapshot_id)
+        for filee in "config.json", "model.bin", "tokenizer_config.json", "tokenizer.json", "vocabulary.txt":
+            # Make symbolic links
+            link = os.path.join(snapshot_dir, filee)
+            if not(os.path.exists(link)): os.symlink(os.path.join(repo_dir,filee), link)
+        self.model = EncoderCT2fromHfHub(
+                model_name_or_path=path,
+                device="cuda",
+                compute_type="int8_float16"
+        )
+
+    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
+        """
+        inputs = data.pop("inputs",data)
+        outputs = self.model.generate(text=[inputs])
+        return outputs["pooler_output"].tolist()
+
+# Test code
+if __name__ == '__main__':
+    from handler import EndpointHandler
+    my_handler = EndpointHandler(path=".", local_test=True)
+    inputs = ['The quick brown fox jumps over the lazy dog']
+    for input in inputs:
+        response = my_handler({"inputs": input})
+    print(response)