custom pipeline

pipeline.py

@@ -0,0 +1,40 @@
+import torch
+from torch import Tensor
+from transformers import AutoTokenizer, AutoModel
+from typing import List
+import os
+
+device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+
+class PreTrainedPipeline():
+    def __init__(self, path=""):
+        # load the optimized model
+        self.model_path = os.path.join("", '.')
+        self.tokenizer = AutoTokenizer.from_pretrained(self.model_path)
+        self.model = AutoModel.from_pretrained(self.model_path)
+        self.model.eval()
+        self.model = self.model.to(device)
+
+    def __call__(self, inputs: str) -> List[float]:
+        """
+        Args:
+            data (:obj:):
+                includes the input data and the parameters for the inference.
+        Return:
+            A :obj:`dict`:. The object returned should be a dict like {"feature_vector": [0.6331314444541931,0.8802216053009033,...,-0.7866355180740356,]} containing :
+                - "feature_vector": A list of floats corresponding to the image embedding.
+        """
+
+        batch_dict = self.tokenizer(inputs, max_length=512,
+                                    padding=True, truncation=True, return_tensors='pt')
+        with torch.no_grad():
+            outputs = self.model(**batch_dict)
+            embeddings = self.average_pool(outputs.last_hidden_state,
+                                           batch_dict['attention_mask'])
+        return embeddings.cpu().numpy().tolist()
+
+    def average_pool(self,  last_hidden_states: Tensor, attention_mask: Tensor) -> Tensor:
+        last_hidden = last_hidden_states.masked_fill(
+            ~attention_mask[..., None].bool(), 0.0)
+        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]

requirements.txt

@@ -0,0 +1,2 @@
+transformers
+torch