Upload pipeline.py

pipeline.py

@@ -0,0 +1,29 @@
+import torch
+from transformers import BertTokenizer
+from foody_bert import FoodyBertForSequenceClassification
+
+class PreTrainedPipeline():
+
+    def __init__(self, path=""):
+
+        """
+        Initialize model
+        """
+        self.bert_tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
+        self.model = FoodyBertForSequenceClassification.from_pretrained(".")
+    def __call__(self, inputs: str) -> List[float]:
+
+        """
+        Args:
+            inputs (:obj:`str`):
+                a string to get the features of.
+        Return:
+            A :obj:`list` of floats: The features computed by the model.
+        """
+        input_ids = self.bert_tokenizer.encode(inputs, add_special_tokens=True)
+        X = torch.tensor([input_ids])
+        with torch.no_grad():
+            predicted_class_id = self.model(X).logits.argmax().item()
+            labels = ['positive','neutral','negative']
+            reps   = labels[predicted_class_id]
+        return self.model.get_sentence_vector(inputs).tolist()